[BUG]_prepare_fp32_grad_for_sub_group clears every parameter's gradient once per sub-group

ZeRO-3: _prepare_fp32_grad_for_sub_group clears every parameter's gradient once per sub-group

Describe the bug

In ZeRO-3, the optimizer step clears gradients for every parameter in every sub-group, once per
sub-group. The number of parameter visits grows as sub-groups x parameters instead of
parameters. Restricting the clear to the current sub-group removes the extra visits.

Measured on two models, the time saved depends strongly on the model:

model, ZeRO-3 + AutoEP, 4 x 8 H200	optimizer step, fix off -> on	saved per step
GLM-4.5-Air	399.78 -> 202.44 ms	197.34 ms, about 4% of the 5.0 s training step
Qwen3.5-397B-A17B, 26 layers	188.88 -> 186.15 ms	2.74 ms, about 0.07% of the 4.2 s training step

The GLM-4.5-Air saving is large enough to matter. The Qwen3.5 saving is real but negligible.

Where

deepspeed/runtime/zero/stage3.py, upstream master at f5af15c2b (2026-09-16):

def _prepare_fp32_grad_for_sub_group(self, sub_group_id):          # line 2405
    ...
    self.fp32_partitioned_groups_flat[sub_group_id].grad = single_grad_partition

    # release all the gradient since we have already created a necessary copy in dp_grad_partition
    self.zero_grad(set_to_none=True)                                 # line 2418
    ...
    self.averaged_gradients[sub_group_id] = None                     # line 2424

zero_grad (line 2202) walks all parameters in all groups:

for group in self.fp16_groups:
    for p in group:
        ...
        p.grad = None

step() calls _prepare_fp32_grad_for_sub_group once for each sub-group:

for sub_group_id, group in enumerate(self.fp16_groups):
    self._prepare_sub_group(sub_group_id, timer_names)   # -> _prepare_fp32_grad_for_sub_group

So each step makes (number of sub-groups) x (number of parameters) parameter visits. Only
(number of parameters) are needed, because each sub-group's own iteration can clear its own
parameters.

When it matters

The cost is small with one or a few sub-groups. It becomes large when a model has many
sub-groups. AutoEP produces many: split_params_into_different_moe_groups_for_optimizer
(deepspeed/moe/utils.py) packs expert tensors under max_group_size = 178,956,971 elements, and
when two expert tensors do not fit together, each one becomes its own group.

model, 4 nodes	sub-groups	parameters	visits per step, current	visits needed
GLM-4.5-Air, full, expert parallel 8	137	735	100,695	735
Qwen3.5-397B-A17B, 26 layers, expert parallel 32	41	479	19,639	479

Both rows were read at optimizer setup, from the number of sub-groups and the number of
parameters in fp16_groups, on the same build used for the timings below. For GLM-4.5-Air the
count follows 3 x (MoE layers) + 2 dense groups: 45 MoE layers give 137, and a 4-layer
configuration gives 11, matching 3 x 3 + 2.

The saving is not proportional to the number of visits. Per wasted visit, GLM-4.5-Air costs
14 times more than Qwen3.5:

model	wasted visits per step	time saved	cost per wasted visit
GLM-4.5-Air	99,960	197.34 ms	1.97 µs
Qwen3.5-397B-A17B	19,160	2.74 ms	0.14 µs

The reason is not established. 0.14 µs is about the cost of a Python loop iteration that finds
p.grad is None and does nothing else. 1.97 µs suggests that on GLM-4.5-Air many visits find a
live gradient and do real work: a record_stream call and a tensor release. That is a hypothesis,
not a measurement. Counting how many visits find p.grad is not None on each model would settle
it, and would show which models are affected. On GLM-4.5-Air the phase is CPU-bound: a PyTorch
profiler trace of _prepare_sub_group, taken on an earlier build, showed 215.9 ms of CPU against
5.2 ms of GPU work per step.

Measured

Same code, same settings. The only difference is whether the clear is restricted to the current
sub-group. Timer: DeepSpeed's step_microstep (wall_clock_breakdown), 4 warmup and 6 measured
steps. The first step is excluded because it allocates optimizer state. Mean of steps 2 to 10.

model, 4 x 8 H200	fix on	fix off (upstream behaviour)	saved per step
GLM-4.5-Air	202.44 ms	399.78 ms	197.34 ms
Qwen3.5-397B-A17B, 26 layers	186.15 ms	188.88 ms	2.74 ms

model	fix-on range	fix-off range
GLM-4.5-Air	191.47 - 210.40 ms	384.76 - 439.35 ms
Qwen3.5-397B-A17B	184.55 - 189.92 ms	187.58 - 190.87 ms

On GLM-4.5-Air the ranges do not overlap: the slowest step with the fix is 174 ms faster than the
fastest step without it. On Qwen3.5 the ranges overlap, but the 2.74 ms difference in the means is
still clear (t about 4.2 over 9 steps each).

Each pair ran on the same four nodes, so node-to-node variation does not affect either comparison.
Whole-step times are not used for Qwen3.5: there the whole step differed by 17 ms in the opposite
direction
(4.1821 s fix on, 4.1651 s fix off), which is run-to-run variation larger than the effect.

Settings: ZeRO-3, AutoEP, micro-batch 1, sequence 8192, full activation recomputation, bf16,
torch 2.8.0+cu126. GLM-4.5-Air used cuDNN attention and deepseek_v3 preset. Qwen3.5 used
flash-linear-attention and the tiled loss.

Numerics are unchanged. The change only affects which param.grad references are set to None
and when, not any value the optimizer reads. These runs do not use deterministic algorithms, so
values differ slightly between any two runs.

model	step 1	largest loss difference, steps 2-10	largest grad-norm difference
GLM-4.5-Air	bit-identical (loss 12.748915, grad norm 0.7232)	9.9e-5	3e-4
Qwen3.5-397B-A17B	loss identical (13.250000), grad norm 9.1958 vs 9.1960	6.9e-3	1.7e-2

All of these are run-to-run noise:

• GLM-4.5-Air: two identical runs of a 4-layer GLM-4.5-Air configuration, on an earlier build,
  differed by 3.7e-4 in step-2 loss (12.780490 and 12.780859). That is more than the 9.9e-5 here.
• Qwen3.5, step 1: it comes before any optimizer update, so the change cannot affect it. The 2e-4
  grad-norm difference comes from non-deterministic kernels (flash-linear-attention and TileLang),
  which GLM-4.5-Air does not use.
• Qwen3.5, loss: the values behind the 6.9e-3 difference, 13.243056 and 13.243750, also appear in
  a separate run with the fix on, and other fix-on runs show similar values.
• Qwen3.5, grad norm: within a single fix-on run of this configuration it ranges over 0.032 to
  0.063 across 10 steps, over three separate runs. That is more than the 1.7e-2 difference
  here.

Proposed fix

Clear only the current sub-group's parameters, and keep the two resets zero_grad does:

self.fp32_partitioned_groups_flat[sub_group_id].grad = single_grad_partition

# Only this sub-group's parameters need clearing: step() visits every sub-group in turn.
self.micro_step_id = 0
self._epilogue_ran_this_backward = False
for p in self.fp16_groups[sub_group_id]:
    if p.grad is not None:
        if get_accelerator().on_accelerator(p.grad):
            p.grad.record_stream(get_accelerator().current_stream())
        p.grad = None

Why it is safe

path through step()	effect of the change
normal GPU path	same result: the loop visits every sub-group, so every sub-group is cleared
overflow	unaffected: _overflow_clean_up calls its own full zero_grad(set_to_none=True)
CPU optimizer offload	unaffected: _prepare_sub_group does not call _prepare_fp32_grad_for_sub_group on this path, before or after
NVMe swap with some sub-groups swappable and some not	the old full sweep inside a non-swappable sub-group also cleared the swappable ones; the new code does not

The last row is covered anyway. DeepSpeedEngine._take_model_step (deepspeed/runtime/engine.py)
calls a full zero_grad after optimizer.step() on every path: bf16, ZeRO, fp16, AMP and the
default. So every step already ends with one full sweep, and the per-sub-group sweep was redundant.
The NVMe mixed case clears at the end of the step instead of during it. It would be worth one test
on that configuration before merging, because it is the only path where the timing of the clear
changes.

To reproduce

1. Run a ZeRO-3 + AutoEP model with wall_clock_breakdown: true on upstream master, and record
   step_microstep.
2. Apply the change above and run again with identical settings.

Notes

• This is a ZeRO-3 issue, not an AutoEP one. AutoEP only makes it visible by creating many
  sub-groups. Any ZeRO-3 run with many sub-groups makes the extra visits, for example with a small
  sub_group_size. How much time they cost depends on the model, as the two measurements show.
• A separate, smaller point: under ZeRO-3 the MoE max_group_size cap is redundant, because ZeRO-3
  already bounds group size with sub_group_size. Removing the cap on Qwen3.5 saved 4.4 ms per
  optimizer step, which is too small to act on. It is listed here only
  because it is the reason the sub-group count is high.

System info

• DeepSpeed 0.19.6, built from upstream master f5af15c2b (2026-09-16) with local patches. The
  two arms of each comparison differ only in _prepare_fp32_grad_for_sub_group. The affected code
  is unchanged in f5af15c2b.
• torch 2.8.0+cu126, transformers 5.15.1
• 4 nodes x 8 NVIDIA H200, EFA