ZeRO-2 + bf16 silently computes incorrect gradients at multi-rank when a submodule is used more than once per step, a regression from #7665, still present in 0.19.3 #8224

Summary

Under ZeRO stage 2 + bf16, multi-rank, no offload, training that was correct on DeepSpeed <= 0.18.2 diverges from 0.18.3 onward and is still wrong on the latest release (0.19.3). The failure is a silent gradient-reduction correctness bug (no error, no warning): the cross-rank reduced gradient is missing one subgraph's contribution.

It reproduces for any model whose backward runs a parameter's gradient-accumulation hook in more than one subgraph of a single backward(): for us, an encoder submodule is applied twice per step (once in the main path, once to encode a second input for an auxiliary loss), so the encoder's parameters receive gradients from two subgraphs. We believe #7665 is the cause (details below).

Environment

• DeepSpeed: 0.18.2 correct, 0.18.3 through 0.19.3 incorrect
• PyTorch 2.11.0, CUDA 13.2.1, Python 3.12
• ZeRO stage 2, bf16.enabled=true, overlap_comm=true, gradient clipping 1.0, no offload
• Multi-rank: 16 ranks (2 nodes x 8 GPUs), gradient accumulation 8, client AdamW + LR scheduler
• Loss terms use mean reduction

How it presents

Training is resumed from a pretrained checkpoint with a warmup that holds LR at 0 for the first phase, then ramps. On the affected versions:

• While LR is held at 0 (no weight updates), every loss ted baseline, so data loading and the forward pass areunaffected.
• The divergence begins exactly when the LR leaves 0 and tng updates, and it plateaus. This is the signature of theoptimizer being fed a wrong gradient, not a data or forward-pass problem.
• The primary reconstruction loss roughly **doubles and never nd auxiliary-latent terms collapse (the model is pushedtoward a degenerate solution because the reduced gradient is biased). It is not uniform degradation; the two auxiliary terms move opposite to the
  reconstruction term.

Medians over a matched window (16 logged points, same config a

The transition is sharp and monotonic: everything <= 0.18.2 is18.3 is wrong, by a factor of ~2.1x on the primary loss with nooverlap between the two groups.

Evidence that it is a gradient-reduction bug (not model / data / LR / local math)

We ruled these out by controlled experiment before blaming reduction:

1. Single-rank is bitwise identical. At world_size == 1, losses, gradient norms, and the weight trajectory over 6 optimizer steps (grad-accum 8)
   are bitwise identical between a good version and 0.18.8. Slip / optimizer math is unchanged. The bug only appears oncecross-rank reduction does real work, i.e. it is in the reduction path.
2. overlap_comm=false does not help (still incorrect on overlapped/streamed-reduction path.
3. Disabling activation checkpointing does not help (still incorrect on 0.18.8), so it is not reentrant-checkpointing recomputation.
4. Forward path is equivalent across versions (identical lt 0).

Summary of the controls (all on 0.18.8 unless noted):

That isolates the fault to cross-rank gradient reduction correctness for this backward-graph shape, at multi-rank only.

Suspected root cause

The only change to the gradient-reduction path between 0.18.2 (correct) and 0.18.3 (first incorrect) is #7665 "PyTorch-compatible backward API"
(commit 53e91a09, first released in 0.18.3). To let loss.baexplicit engine.backward(), DeepSpeed now has to **infer**when a backward is complete: it replaced unconditional per-hook gradient reduction with an **end-of-backward epilogue gated by a participating-parameter count**: count_used_parameters_in_backward()inruntime/utitate machinery in base_optimizer.py(should_refresh_expected_hook_count / update_hook_state_and_maybe_run_epilogue) and create_gradient_handling_hooks() in
runtime/zero/stage_1_and_2.py.

Our hypothesis for the mechanism: when a parameter's grad hooksubgraph** of a single backward (a submodule applied twice, ormultiple loss terms over shared params), the number of hook firings does not match the predicted participating-parameter count, so the reduction
epilogue runs **before all of that parameter's gradient has act cross-rank reduction therefore averages an incompletegradient. This is consistent with every observation above (single-rank correct; only reduction affected; independent of overlap_comm and activation
checkpointing).

Why the existing follow-up fixes do not cover this

There is a cluster of fixes hardening the #7665 multi-backwards the no-offload ZeRO-2 automatic path this hits:

• Fix ZeRO-1/2 CPU-offloaded gradient loss with multiple backward() per step #7981: ZeRO-1/2 CPU-offload gradient loss with multiple -specific).
• Fix BF16_Optimizer last-microbatch grad leak under ZeRO-1 #7985: ZeRO-1 BF16_Optimizer last-microbatch grad leak (stage-1-specific).
• Add engine.coalesce_grad_reduction() for ZeRO 1/2/3 multi-backward #7992: coalesce_grad_reduction(), an opt-in context maath.
• [BUG]Incorrect gradient computation in ZeRO-2 with DeepSpeed ≥ 0.17.6 #7718: the inverse case that PyTorch-compatible backward API #7665 fixed for a different graph shape.

Confirmed still incorrect on 0.19.3 with all of the above merged.

Impact

Silent: no exception, just wrong gradients and degraded training. Easy to miss without a single-rank or DDP reference, or (as in our case) a
checkpoint-resume warmup that makes the divergence visually obegin. Any ZeRO-2 + bf16 multi-rank job whose backward graphreuses a submodule or shares parameters across loss terms is exposed.

Workaround

Pin deepspeed==0.18.2 (last release before #7665).

Repro

We can share a minimal synthetic reproducer (a tiny model that applies one nn.Linear twice, ZeRO-2 + bf16, world_size >= 2, comparing safe_get_full_grad against a manually all-reduced DDP referedd it on request.