Skip to content

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

Merged
delock merged 1 commit into
deepspeedai:masterfrom
delock:pr-e-offload-memory-deltas
Sep 15, 2026
Merged

delock merged 1 commit into
deepspeedai:masterfrom
delock:pr-e-offload-memory-deltas

Conversation

@delock

@delock delock commented Sep 4, 2026

Copy link
Copy Markdown
Collaborator

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 chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment on lines +21 to +22
accelerator_device_mod = torch.get_device_module(get_accelerator().device_name())
allocator_backed_memory_stats = hasattr(accelerator_device_mod, 'memory_allocated')

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

Comment on lines +291 to +292
expect_memory_change = allocator_backed_memory_stats and not (is_offload_optimizer_enabled(config_dict)
and offload_only_optimizer_states)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge 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())

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

@pengdurice
pengdurice added this pull request to the merge queue Sep 14, 2026
@github-merge-queue
github-merge-queue Bot removed this pull request from the merge queue due to failed status checks Sep 14, 2026
@delock
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
delock deleted the pr-e-offload-memory-deltas branch September 15, 2026 01:32
pull Bot pushed a commit to QSLee-Net/DeepSpeed that referenced this pull request Sep 17, 2026
…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>
delock added a commit to delock/DeepSpeedSYCLSupport that referenced this pull request Sep 18, 2026
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>
yermakoffivan pushed a commit to yermakoffivan/deepspeed that referenced this pull request Sep 28, 2026
…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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants