[BUG] Flaky: TestZero2CPUOffloadGradAccumNorm::test fails intermittently (accumulated_grads_in_cpu should be non-zero after first backward)

Summary

tests/unit/v1/zero/test_zero_cpu_offload_grad_accum.py::TestZero2CPUOffloadGradAccumNorm::test fails intermittently on the modal-torch-latest workflow and blocks the merge queue. It has been observed at least twice in the last ~100 runs of that workflow, both times with the same assertion and nothing else failing in the suite.

The failure looks like a timing race between the asynchronous GPU→CPU offload copy and a host-side read of optimizer.accumulated_grads_in_cpu, so a green run proves nothing and the fix needs a repeated-run protocol (see "Required verification" at the end).

Failed test

tests/unit/v1/zero/test_zero_cpu_offload_grad_accum.py::TestZero2CPUOffloadGradAccumNorm::test
AssertionError: accumulated_grads_in_cpu should be non-zero after first backward

____________________ TestZero2CPUOffloadGradAccumNorm.test _____________________
[gw2] linux -- Python 3.10.13 /usr/local/bin/python
multiprocessing.pool.RemoteTraceback:
"""
Traceback (most recent call last):
  ...
  File "/workspace/deepspeed/tests/unit/v1/zero/test_zero_cpu_offload_grad_accum.py", line 85, in test
    assert norms[0] > 0, "accumulated_grads_in_cpu should be non-zero after first backward"
AssertionError: accumulated_grads_in_cpu should be non-zero after first backward
"""
=========================== short test summary info ============================
FAILED tests/unit/v1/zero/test_zero_cpu_offload_grad_accum.py::TestZero2CPUOffloadGradAccumNorm::test
= 1 failed, 1693 passed, 193 skipped, 2 xfailed, 52 warnings in 4460.96s (1:14:20) =
DS_CI_FAILURE_CLASS=test: candidate failed

Evidence

Run 2 was a direct master push, so this is not caused by any single in-flight PR. It is classified as test rather than infra/timeout by the failure triage added in #8606, i.e. a test really ran and really failed.

For reference, the other failures of that workflow in the same window were unrelated infra problems (SandboxStartTimeout, runner shutdown), so this is the only test-level flake seen recently.

Analysis (suspected root cause, not yet verified on GPU)

Line numbers are from origin/master.

The test reads pinned host memory that a GPU-side asynchronous copy fills, with no synchronization anywhere in between:

# tests/unit/v1/zero/test_zero_cpu_offload_grad_accum.py:22-26, 77-89
def _cpu_grad_norm(engine):
    total_norm_sq = 0.0
    for grad in engine.optimizer.accumulated_grads_in_cpu.values():
        total_norm_sq += grad.float().norm(2).item()**2   # host read, no synchronize
    return total_norm_sq**0.5

for batch in data_loader:
    loss = model(batch[0], batch[1])
    model.backward(loss)
    norms.append(_cpu_grad_norm(model))        # read immediately after backward returns
...
assert norms[0] > 0, "accumulated_grads_in_cpu should be non-zero after first backward"

backward() populates that buffer through this chain:

engine._backward_epilogue()                                 # engine.py:3124-3136
  └─ allreduce_gradients()                                  # engine.py:3048, called at :3136
      └─ optimizer.reduce_gradients()                       # engine.py:3067
          └─ overlapping_partition_gradients_reduce_epilogue()   # stage_1_and_2.py:1165
              └─ independent_gradient_partition_epilogue()       # stage_1_and_2.py:984
                  └─ self.reduce_ipg_grads()                     # stage_1_and_2.py:986
                      └─ copy_grads_in_partition(param)          # stage_1_and_2.py:2018
                          └─ async_accumulate_grad_in_cpu_via_gpu(param)   # stage_1_and_2.py:1924

The copy itself is deliberately asynchronous, into a pinned buffer whose initial content is zeros:

# deepspeed/runtime/zero/stage_1_and_2.py:1622-1626
def buffer_to_accumulate_to_in_cpu():
    if not self.low_precision_master_weights_and_grads:
        buffer = torch.zeros(param.numel(), dtype=param.dtype, device=self.device)
        return get_accelerator().pin_memory(buffer) if self.cpu_offload_pin_memory else buffer

# deepspeed/runtime/zero/stage_1_and_2.py:1643-1651
def copy_gradients_to_cpu():
    grad_accum = self.get_param_gradient_attribute(param)   # param.grad_accum / param.grad, on the GPU
    ...
    self.accumulated_grads_in_cpu[param_id].data.copy_(grad_accum.data.view(-1), non_blocking=True)

• offload_optimizer.pin_memory defaults to True (deepspeed/runtime/zero/offload_config.py:45,75), so non_blocking=True is a genuinely asynchronous D2H copy and copy_ returns before the data has landed.
• Because the initial content is zeros, a host read that wins the race observes exactly 0.0, which is the reported symptom.
• contiguous_gradients is forced on by offload (stage_1_and_2.py:272), and copy_grads_in_partition takes the accumulate branch for micro_step_id > 0 or not is_gradient_accumulation_boundary() (stage_1_and_2.py:1923), so this path is exercised on every backward, including the first micro-step of the gas=4 window.

The only whole-device synchronization nearby is skipped for this configuration:

# deepspeed/runtime/zero/stage_1_and_2.py:994-1000
if self.overlap_comm:
    if not get_accelerator().resolves_data_dependency():
        get_accelerator().synchronize()      # only under overlap_comm
    self._clear_previous_reduced_grads()

if self.cpu_offload is False:                # offload skips this branch entirely
    ...

The test config does not set overlap_comm, and for ZeRO-2 the dynamic default is overlap_comm = (stage == weights), i.e. False (deepspeed/runtime/zero/config.py:384-385). So on the Modal CUDA runner nothing guarantees that accumulated_grads_in_cpu is complete when backward() returns, and the assertion is evaluated on memory that may still be in flight. Under pytest -n 4 on a loaded sandbox the host can win that race. Note that the race needs an accelerator that does not resolve data dependencies by itself (resolves_data_dependency() == False — CUDA, MLU, MPS, NPU, SDAA); CPU and HPU return True and cannot reproduce it, which is why a CPU-only validation of the fix would be meaningless.

The reverse direction has the same hazard and is the other half of this bug: at micro_step_id > 0, accumulate_gradients() copies the same pinned buffer back to the GPU with non_blocking=True (stage_1_and_2.py:1631-1641), which reads stale data if the previous D2H has not landed, silently corrupting gradient accumulation.

Required verification for a fix

This is a timing race, so a single passing CI run is not evidence. Any PR claiming to fix this must include the following in its description:

1. Measured pre-fix reproduction. Run the test repeatedly before the fix and report the command, iteration count, pass/fail counts, and failure message. Any CUDA-capable device works; keep the torch/CUDA stack comparable to the MODAL_TORCH_PRESET. If it does not reproduce, report the iteration count you ran.

2. Report the environment: GPU model/count, torch and CUDA versions, and the exact commands, per the project's integration-test rules.

Candidate fix directions (for discussion, not prescriptive)

• Test-side: synchronize before reading host accumulators, e.g. get_accelerator().synchronize() at the start of _cpu_grad_norm. This makes the test deterministic without changing the engine contract: the assertion is about the content of accumulated_grads_in_cpu, and the synchronization only makes the observation point well defined.

I do not have GPU access to reproduce this locally, so the analysis above is inferred from the CI logs and the code; the pre-fix reproduction in the verification protocol is where it should be confirmed or refuted.

How the test runs in CI

# ci/torch_latest.py (master), the Modal sandbox command
("pytest", "-n", "4", "--verbose",
 "--ignore=tests/unit/v1/nvme/test_gds.py",
 f"--torch_ver={preset['torch_test_version']}",
 f"--cuda_ver={preset['cuda_test_version']}",
 "--", *targets)

pytest -n 4 with pytest-randomly in the dev requirements, so worker load and test order vary between runs.