Gate the offload-state memory deltas on allocator-backed stats - #8409
Gate the offload-state memory deltas on allocator-backed stats#8409delock wants to merge 1 commit into
Conversation
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 deepspeedai#8381 (92 failures). Signed-off-by: Guokai Ma <guokai.ma@intel.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ca1594753d
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| accelerator_device_mod = torch.get_device_module(get_accelerator().device_name()) | ||
| allocator_backed_memory_stats = hasattr(accelerator_device_mod, 'memory_allocated') |
There was a problem hiding this comment.
On supported PyTorch versions that predate torch.get_device_module (including PyTorch 2.0, which remains allowed by requirements/requirements.txt:10), importing this test module now raises AttributeError during pytest collection, so none of its tests can run. Probe the method with getattr or determine the capability through the DeepSpeed accelerator abstraction instead.
Useful? React with 👍 / 👎.
| expect_memory_change = allocator_backed_memory_stats and not (is_offload_optimizer_enabled(config_dict) | ||
| and offload_only_optimizer_states) |
There was a problem hiding this comment.
Under the CPU accelerator—the environment this change is intended to unblock—this makes expect_memory_change false, but the existing validate_device(model, offload_state_device, offloaded_states) call is inside the same conditional at lines 336–338. Consequently, all ZeRO-3 CPU cases stop verifying that the requested states actually moved to CPU; the later equality checks after reload can pass even if offloading was a no-op. Gate only the memory-delta assertion and keep the observable device-placement validation unconditional.
AGENTS.md reference: AGENTS.md:L30-L32
Useful? React with 👍 / 👎.
| # 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()) |
There was a problem hiding this comment.
This is a one-parent, non-merge commit, but its commit message has no Signed-off-by trailer. Add the required signoff before merging so the commit satisfies the repository's mandatory commit policy.
AGENTS.md reference: AGENTS.md:L8-L8
Useful? React with 👍 / 👎.
1 participant
Description
The dynamic offload-state tests assert strict allocated-memory deltas around
offload_states()/reload_states():alloc_after_offload < alloc_before_offloadalloc_after_reload > alloc_after_offloadThat contract assumes
memory_allocated()is allocator bookkeeping, which holds on cuda (torch.cuda.memory_allocated()). On cpu,CPU_Accelerator.memory_allocated()reports process RSS (psutil), and RSS does not shrink when tensors are freed — so all 92 parameterized cases fail even when the offload itself is correct (the device-placement and data-integrity checks in the same tests pass).Gate only the memory-delta asserts on whether the accelerator's torch device module exposes
memory_allocated(cuda does;torch.cpudoes not), mirroring the capability probe used forfork_rngintrain_cifar(#8407):Exposed by the
LOCAL_SIZE=4multi-rank CPU run in #8381 (92 of the 131 v1-half failures there).Validation (executed on real hardware)
LOCAL_SIZE=2)TestDynamicOffloadStatesZero12[False-1-False-False-optim_states]fails on the persistent-state delta assertstatic_offload_optimizer=Truebranch) — 5 passed in 55.6sSibling PRs from the same series: #8397, #8398, #8399, #8407.