Gate the offload-state memory deltas on allocator-backed stats by delock · Pull Request #8409 · deepspeedai/DeepSpeed

delock · 2026-09-04T01:10:16Z

Description

The dynamic offload-state tests assert strict allocated-memory deltas around offload_states() / reload_states():

alloc_after_offload < alloc_before_offload
alloc_after_reload > alloc_after_offload

That 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.cpu does not), mirroring the capability probe used for fork_rng in train_cifar (#8407):

cuda and other allocator-backed backends: behavior unchanged
cpu: the unobservable deltas are skipped; all device-placement validations still run

Exposed by the LOCAL_SIZE=4 multi-rank CPU run in #8381 (92 of the 131 v1-half failures there).

Validation (executed on real hardware)

20-core x86_64 CPU, torch 2.13.0+cpu, gloo backend, 2 ranks (LOCAL_SIZE=2)
Before: TestDynamicOffloadStatesZero12[False-1-False-False-optim_states] fails on the persistent-state delta assert
After: 5 representative cases pass (persistent and grad paths, ZeRO stage 1/2/3, static_offload_optimizer=True branch) — 5 passed in 55.6s
pre-commit (yapf / flake8 / check-torchdist / codespell) passes on the changed file

Sibling PRs from the same series: #8397, #8398, #8399, #8407.

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

💡 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".

chatgpt-codex-connector · 2026-09-04T01:13:04Z

+accelerator_device_mod = torch.get_device_module(get_accelerator().device_name())
+allocator_backed_memory_stats = hasattr(accelerator_device_mod, 'memory_allocated')


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 👍 / 👎.

chatgpt-codex-connector · 2026-09-04T01:13:04Z

+    expect_memory_change = allocator_backed_memory_stats and not (is_offload_optimizer_enabled(config_dict)
+                                                                  and offload_only_optimizer_states)


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 👍 / 👎.

chatgpt-codex-connector · 2026-09-04T01:13:04Z

+# 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())


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 👍 / 👎.

…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>

…es (deepspeedai#8702) ## Description Two latent bugs in Ulysses SP, both backend-independent (neither is CPU-specific — the CPU/gloo multi-rank run in deepspeedai#8381 is simply what surfaced them), plus test adaptations so the file can also run on non-CUDA accelerators. ### 1. `register_with_transformers` rejected every string model path `deepspeed/runtime/sequence_parallel/ulysses_sp.py` compared `core_attn_implementation` against the config's `_attn_implementation` even when the caller passed a **model path string**: ```python hf_model_config = AutoConfig.from_pretrained(model_name_or_path) # string branch model_attn_implementation = getattr(hf_model_config, "_attn_implementation", None) if model_attn_implementation is not None and model_attn_implementation != core_attn_implementation: raise ValueError(...) ``` In transformers, `_attn_implementation` is a **property resolved at model load time** — `from_pretrained` resolves which implementation the runtime can actually use and writes it back to the config. A bare `AutoConfig` never goes through that resolution, so the property returns its fallback value for "unresolved": ``` AutoConfig.from_pretrained(path)._attn_implementation -> eager (fallback, not a choice) AutoModelForCausalLM.from_pretrained(path).config._attn... -> sdpa (resolved) ``` So every string-path caller asking for `sdpa`/`flex_attention`/FA2 was rejected with a spurious mismatch against `'eager'`. The check is only meaningful for a config that was resolved, i.e. when the caller handed us a constructed model (or a PEFT wrapper exposing `.config`); scope it to that case. The comparison expression itself is unchanged, and `TestUlyssesSPHFAttnImplMismatch` (which passes a model object) still raises as intended. ### 2. `UlyssesSPDataLoaderAdapter` used a 0-dim tensor in `all_gather` ```python seqlen = torch.tensor(n, ...) # 0-dim seqlens = [torch.zeros(1, ...) for _ in range(sp_world_size)] # shape (1,) dist.all_gather(seqlens, seqlen, group=self.sp_group) ``` `c10d.all_gather` requires each receive tensor to have the same size as the send tensor. NCCL moves flattened bytes and tolerates the mismatch; gloo validates shapes and rejects it (`ProcessGroupGloo::allgather: invalid tensor size at index 0 (expected (), got (1))`). Send a 1-element tensor to match. ### 3. Tests - `TestUlyssesSPHFFlexAttention` needs the CUDA flex kernel -> skip on non-CUDA accelerators (same guard idiom as `tests/unit/util.py`, the nvme/triton CUDA-only tests). - `TestUlyssesSPHFPEFT`: the mock attached a **bare** `AutoConfig`, whose `_attn_implementation` is the unresolved `'eager'` — so this test failed on every backend, not just CPU. Give the mock the load-resolved value a real PEFT wrapper carries (the sibling `TestUlyssesSPHFAttnImplMismatch` in the same file already sets it the same way). - `TestUlyssesSPHFDisableInEval` pinned its tensors/models to `cuda:<rank>` although it runs sdpa and nothing in it is CUDA-specific. Use `get_accelerator().current_device_name()` like the rest of the file (on CUDA that is exactly `cuda:<rank>`, since the test harness sets the device per rank in `tests/unit/common.py`), so the test runs on any backend instead of erroring. ## Validation Hardware: 20-core x86_64 CPU without AVX512-FP16, gloo backend, 2 ranks (`LOCAL_SIZE=2`), transformers 4.51.3, torch 2.13. - `tests/unit/ulysses_alst/test_ulysses_sp_hf.py`: **5 passed, 2 skipped** (flex skipped by the new guard). The sdpa classes pass end to end including their numerical-parity assertions; before the fix they died at registration. - `TestUlyssesSPHFDisableInEval` now actually **executes** on CPU rather than being skipped. - `TestUlyssesSPHFAttnImplMismatch` still raises the expected `ValueError` (the mismatch check is preserved where a resolved config exists). - Full multi-rank CPU CI in deepspeedai#8381: this file went from 5 failures to 0, with the overall job green. (`TestUlyssesSPHFHubKernel` is not exercised by the version pin above — it needs the `lazy_import_flash_attention` API from a newer transformers; it passes in CI, which installs transformers from git main.) Sibling PRs from the same series: deepspeedai#8397, deepspeedai#8398, deepspeedai#8399, deepspeedai#8407, deepspeedai#8409, deepspeedai#8559, deepspeedai#8648, deepspeedai#8684. --------- Signed-off-by: Guokai Ma <guokai.ma@intel.com> Co-authored-by: pengdurice <pengduhit@gmail.com>

The dynamic offload-state tests already relax their memory-delta assertions for non-allocator backends (deepspeedai#8409), but validate_device's contiguous_grad_buffer check still assumed the state was moved out of the accelerator: when it is reported on cpu, the buffer is required to be gone. On the cpu accelerator the offload target is the accelerator itself, so the buffer is never moved and that assertion fails. Gate the check on the offload actually crossing memory tiers. On CUDA the new flag is always true, so the assertion behaves exactly as before; on cpu the state falls through to the device comparison, which holds. Measured on cpu/gloo with 2 ranks (master's version, no skip): 112 of the 128 parametrizations already pass, the 16 failures are all this one assertion. With the guard the whole file passes, so the cpu run keeps the coverage a module-level skip would have dropped. Signed-off-by: Guokai Ma <guokai.ma@intel.com>

…deepspeedai#8738) ## Description `validate_device` in `tests/unit/v1/zero/test_offload_states.py` requires the contiguous grad buffer to be **gone** once it is reported on cpu, because a cpu device means `offload_states()` moved it to the host tier: ```python if state == OffloadStateTypeEnum.contiguous_grad_buffer and state_device[state] == torch.device("cpu"): assert len(get_state_devices(model, state)) == 0, f"State {state} must be removed after offload_states()" ``` On the cpu accelerator the offload target **is** the accelerator, so `state_device[state] == torch.device("cpu")` holds for everything, yet the buffer was never moved anywhere — the assertion fails for every parametrization that includes it. This is the device-state counterpart of deepspeedai#8409, which gated the **memory-delta** assertions on `allocator_backed_memory_stats` for the same reason. Add a sibling `offload_leaves_accelerator` flag (next to that one) and require it for the removal check. ## Behavior | Accelerator | `offload_leaves_accelerator` | Effect | |---|---|---| | CUDA / XPU | `True` | assertion unchanged, bit for bit | | cpu | `False` | state falls through to the device comparison instead, which holds | ## Validation (executed on real hardware) 20-core x86_64 CPU without AVX512-FP16, gloo backend, 2 ranks (`LOCAL_SIZE=2`), `-n 4`: - **Without** this change on master: `16 failed, 112 passed` — and all 16 failures are this single assertion (`State OffloadStateTypeEnum.contiguous_grad_buffer must be removed after offload_states()`). - **With** this change: `128 passed, 0 failed`. So the cpu accelerator keeps the coverage of the other 112 parametrizations instead of losing them to a blanket skip. The enclosing cpu CI run is deepspeedai#8381, which exercises this file at multi-rank and reports the same result (`128 PASSED`). Signed-off-by: Guokai Ma <guokai.ma@intel.com>

delock requested review from loadams, tjruwase and tohtana as code owners September 4, 2026 01:10

chatgpt-codex-connector Bot reviewed Sep 4, 2026

View reviewed changes

pengdurice approved these changes Sep 14, 2026

View reviewed changes

pengdurice added this pull request to the merge queue Sep 14, 2026

github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Sep 14, 2026

delock added this pull request to the merge queue Sep 15, 2026

Merged via the queue into deepspeedai:master with commit e6d2d40 Sep 15, 2026
13 of 15 checks passed

delock deleted the pr-e-offload-memory-deltas branch September 15, 2026 01:32

This was referenced Sep 16, 2026

[DON'T MERGE] Run multi-rank CPU unit tests in CI via LOCAL_SIZE #8381

Open

Derive the autocast test backend and scope the bf16 gate to NCCL #8559

Merged

delock mentioned this pull request Sep 27, 2026

Skip fp16-config tests on accelerators without fp16 support #8684

Merged

delock mentioned this pull request Sep 29, 2026

Fix Ulysses SP registration for string model paths and gloo collectives #8702

Merged

delock mentioned this pull request Oct 4, 2026

Assert offload-state removal only when offload leaves the accelerator #8738

Merged

Navigation Menu

Repository navigation

Gate the offload-state memory deltas on allocator-backed stats - #8409

Gate the offload-state memory deltas on allocator-backed stats#8409
delock merged 1 commit into
deepspeedai:masterfrom
delock:pr-e-offload-memory-deltas

delock commented Sep 4, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

chatgpt-codex-connector Bot Sep 4, 2026

Uh oh!

chatgpt-codex-connector Bot Sep 4, 2026

Uh oh!

chatgpt-codex-connector Bot Sep 4, 2026

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		accelerator_device_mod = torch.get_device_module(get_accelerator().device_name())
		allocator_backed_memory_stats = hasattr(accelerator_device_mod, 'memory_allocated')

		expect_memory_change = allocator_backed_memory_stats and not (is_offload_optimizer_enabled(config_dict)
		and offload_only_optimizer_states)

Conversation

delock commented Sep 4, 2026

Description

Validation (executed on real hardware)

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Sep 4, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Sep 4, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Sep 4, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants