Gate the offload-state memory deltas on allocator-backed stats - #8409
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>
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.
Avoid requiring torch.get_device_module during collection
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.
Preserve the ZeRO-3 offload placement check
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.
Add the required Signed-off-by trailer
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 👍 / 👎.
e6d2d40
…pspeedai#8559) ## Description `bf16_required_version_check()` (`tests/unit/util.py`) requires torch >= 1.10, **CUDA >= 11.0 and NCCL >= 2.10.3**. On the cpu accelerator, bf16 collectives run over gloo/ccl and none of those transport dependencies exist, so the check always returns False and **every bf16 test is skipped — about 40 call sites across 15 files**. `test_zero_autocast.py` is worse off: it **raises** instead of skipping, so each of its cases counts as a failure (24 on the multi-rank CPU run in deepspeedai#8381). This PR scopes the version floors inside the check itself: ```python if (cpu_accelerator and accelerator_pass) or (torch_version_available and cuda_version_available and nccl_version_available and accelerator_pass): return True ``` - **cpu**: only the accelerator's own bf16 support (`is_bf16_supported()`) is required — the torch/CUDA/NCCL floors are transport dependencies that gloo/ccl does not have - **every other accelerator** (cuda, npu, hpu, xpu, mlu, ...): evaluates the exact original floors; `cpu_accelerator` is False so the expression is bit-identical to before, and the npu/hpu/xpu exemption branches are untouched - `test_zero_autocast.py`: the bf16 gate goes back to the bare call every other caller uses, and skips instead of raising The hardcoded `init_distributed(dist_backend='nccl')` in the same test is deliberately left alone: it is a no-op (the harness already initialized the process group, `comm.py:838-839`), and deriving the backend per accelerator would change behavior on non-cuda accelerators (npu/hpu resolve to hccl etc.). The baseline `DDP(device_ids=[i])` pinning is left for a follow-up (deepspeedai#8399 fixed the same pattern elsewhere). ## Validation (executed on real hardware) - 20-core x86_64 CPU, torch 2.13.0+cpu, gloo backend - Direct call: `bf16_required_version_check()` on cpu returns `False` before, `True` after. Note: `CPU_Accelerator.is_bf16_supported()` is currently a stub that always returns True, so the cpu path does not gate on the hardware's bf16 instructions — giving it a real capability probe is left as a follow-up - Non-cpu equivalence: with `cpu_accelerator == False` the new expression reduces exactly to the original `A and C and N and P` - pre-commit (yapf / flake8 / check-torchdist / codespell) passes on both changed files Sibling PRs from the same series: deepspeedai#8397, deepspeedai#8398, deepspeedai#8399, deepspeedai#8407, deepspeedai#8409. Exposed by the `LOCAL_SIZE=4` multi-rank CPU run in deepspeedai#8381. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Every test in this file drives offload_states()/reload_states(), whose contract presumes two memory tiers: offload frees accelerator-side state and reload restores it. On the cpu accelerator the offload target is the accelerator itself, so the contract is not observable there: memory deltas have no allocator-backed metric (RSS does not shrink on free, deepspeedai#8409) and the freed-vs-restored lifecycle cannot be keyed on device placement. Running the file on cpu only produced failures that say more about the degenerate setup than about the engine. Skip the module on cpu, mirroring the hpu module skip in test_onebit.py. Other accelerators keep the full suite. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
…dai#8684) ## Description Every fp16-config test that reaches `deepspeed.initialize` crashes its sanity check (`Type fp16 is not supported on your device.`) on accelerators whose `is_fp16_supported()` is false. On CPU that maps to the AVX512-FP16 capability of the host, and GitHub's `ubuntu-24.04` runners are hardware-heterogeneous, so these tests flip between failure and skip depending on which runner they land on. deepspeedai#8398 added the first skipifs; the multi-rank CPU run in deepspeedai#8381 flushed out six more files: | File | Shape of the gap | |---|---| | `checkpoint/test_universal_checkpoint.py` | fp16 parametrizations fail **inside the baseline DistributedFixture's distributed run** — pytest reports a setup **ERROR** for every dependent test (48 on the multi-rank run) instead of a skip | | `v1/zero/test_zero_coalesce_grad_reduction.py` | `TestCoalesceFP16` forces an fp16 config | | `runtime/test_no_sync_ctxt.py` | dtype=float16 parametrizations of three methods; stages 2/3 additionally never reach their expected no_sync AssertionError on such hosts | | `checkpoint/test_moe_checkpoint.py` | whole class hardcodes fp16 | | `runtime/zero/test_zero_offloadpp.py` | `TestZeroPartialOffloadConfigSweep` hardcodes fp16 | | `checkpoint/test_pipeline.py` | fp16 enabled for zero_stage > 0; only that parametrization skips, zero_stage=0 keeps running | With this PR, **every fp16-config test in the suite guards on `is_fp16_supported()`** — the capability gap is a skip, not a failure, on any accelerator. ## Validation (executed on real hardware) - 20-core x86_64 CPU without AVX512-FP16, gloo, 2-4 ranks (`LOCAL_SIZE=2/4`) - Before: 18 failures + 48 setup ERRORs across these files on the multi-rank CPU run - After: every affected parametrization skips; adjacent non-fp16 parametrizations keep passing (e.g. pipeline zero_stage=0 runs to completion) - Full-suite evidence in deepspeedai#8381: the multi-rank CPU run went 8 failures -> 0 with these guards Sibling PRs from the same series: deepspeedai#8397, deepspeedai#8398, deepspeedai#8399, deepspeedai#8407, deepspeedai#8409, deepspeedai#8559, deepspeedai#8648. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
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.