From ca1594753dfd4d3f58607f7c275251b4861c6420 Mon Sep 17 00:00:00 2001 From: Guokai Ma Date: Fri, 4 Sep 2026 08:57:37 +0800 Subject: [PATCH] Gate the offload-state memory deltas on allocator-backed stats The dynamic offload tests assert strict allocated-memory deltas around offload_states()/reload_states(). That contract assumes memory_allocated() is allocator bookkeeping, which holds on cuda but not on cpu: CPU_Accelerator reports process RSS there, and RSS does not shrink when tensors are freed, so every delta assert fails even when the offload itself is correct. Gate the deltas on whether the accelerator's torch device module exposes memory_allocated (cuda does, torch.cpu does not), mirroring the fork_rng capability probe in train_cifar. Device-placement validations keep running on every backend; only the unobservable memory deltas are skipped. Exposed by the LOCAL_SIZE=4 multi-rank CPU run in #8381 (92 failures). Signed-off-by: Guokai Ma --- tests/unit/v1/zero/test_offload_states.py | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/tests/unit/v1/zero/test_offload_states.py b/tests/unit/v1/zero/test_offload_states.py index bf51d6df2e77..86e70c981e1b 100644 --- a/tests/unit/v1/zero/test_offload_states.py +++ b/tests/unit/v1/zero/test_offload_states.py @@ -15,6 +15,12 @@ from deepspeed.utils import safe_get_local_fp32_param, safe_get_local_optimizer_state from deepspeed.runtime.zero.offload_states import get_state_devices +# The strict allocated-memory deltas asserted in this file assume memory_allocated() +# is allocator bookkeeping (cuda); on cpu it reports process RSS, which does not +# shrink when tensors are freed. +accelerator_device_mod = torch.get_device_module(get_accelerator().device_name()) +allocator_backed_memory_stats = hasattr(accelerator_device_mod, 'memory_allocated') + # ============================================================================== # ZeRO-1 and ZeRO-2 TESTS # ============================================================================== @@ -80,7 +86,8 @@ def run_model_zero12(model, param_groups, config_dict, hidden_dim, dtype, offloa optimizer_device = offload_torch_device if is_offload_optimizer_enabled(config_dict) else accelerator_device offload_only_optimizer_states = is_only_offload_optimizer_states( offloaded_states, [OffloadStateTypeEnum.optim_states, OffloadStateTypeEnum.hp_params]) - expect_memory_change = not (is_offload_optimizer_enabled(config_dict) and offload_only_optimizer_states) + expect_memory_change = allocator_backed_memory_stats and not (is_offload_optimizer_enabled(config_dict) + and offload_only_optimizer_states) model, _, _, _ = deepspeed.initialize(model=model, model_parameters=param_groups, config=config_dict) @@ -116,14 +123,16 @@ def run_model_zero12(model, param_groups, config_dict, hidden_dim, dtype, offloa alloc_after_offload = get_accelerator().memory_allocated() if grad_numel > 0: - assert alloc_after_offload < alloc_before_offload, f"FAIL: Allocated memory for grads should decrease after offload {alloc_after_offload=} < {alloc_before_offload=}" + if allocator_backed_memory_stats: + assert alloc_after_offload < alloc_before_offload, f"FAIL: Allocated memory for grads should decrease after offload {alloc_after_offload=} < {alloc_before_offload=}" validate_grad_device(model, offload_torch_device) model.reload_states() alloc_after_reload = get_accelerator().memory_allocated() if grad_numel > 0: - assert alloc_after_reload > alloc_after_offload, f"FAIL: Allocated memory for grads should increase after reload {alloc_after_reload=} > {alloc_after_offload=}" + if allocator_backed_memory_stats: + assert alloc_after_reload > alloc_after_offload, f"FAIL: Allocated memory for grads should increase after reload {alloc_after_reload=} > {alloc_after_offload=}" validate_grad_device(model, accelerator_device) reloaded_grads = [ @@ -279,7 +288,8 @@ def run_model_zero3(model, param_groups, config_dict, hidden_dim, dtype, offload offload_only_optimizer_states = is_only_offload_optimizer_states( offloaded_states, [OffloadStateTypeEnum.optim_states, OffloadStateTypeEnum.hp_params, OffloadStateTypeEnum.lp_grads]) - expect_memory_change = not (is_offload_optimizer_enabled(config_dict) and offload_only_optimizer_states) + expect_memory_change = allocator_backed_memory_stats and not (is_offload_optimizer_enabled(config_dict) + and offload_only_optimizer_states) offload_state_device: dict[OffloadStateTypeEnum, torch.device] = { OffloadStateTypeEnum.hp_params: offload_torch_device,