[BUG] Muon + fp16 does not train under ZeRO 1/2: the first loss-scale overflow permanently poisons the momentum buffer

[alanhuangyoo]

Summary

Muon under ZeRO 1/2 with fp16 does not train. The first loss-scale overflow folds an overflowed gradient into Muon's momentum buffer, the buffer stays NaN for the rest of the run, and every subsequent step therefore overflows too. The loss scaler backs off until it gives up:

Exception: Current loss scale already at minimum - cannot decrease scale anymore. Exiting run.

No parameter is ever updated. This is with the configuration tests/unit/ops/muon/test_muon.py itself uses, only run for longer.

Isolation

Same configuration, only the initial loss scale differs. 30 steps, 2 GPUs, ZeRO-2, SimpleModel(hidden_dim=128, nlayers=5), lr=0.05:

initial_scale_power	parameters changed	momentum finite	outcome
16 (the default, 65536)	0 / 10	False	dies at minimum loss scale
0 (scale 1, no overflow)	10 / 10	True	trains normally

So the failure is entirely the overflow interaction, not Muon itself.

Mechanism

An overflowed step is supposed to be discarded. Muon consumes the gradient before that happens.

get_flat_partition (stage_1_and_2.py) calls muon_update on the accumulated gradient while filling the partition, which runs in independent_gradient_partition_epilogue, i.e. before the overflow check that decides whether to skip the step. muon_update then does:

momentum.lerp_(grad, 1 - beta)                      # inf/nan enters the momentum
update = grad.lerp_(momentum, beta) if nesterov     # ...and is written back into grad

Both halves matter. The first poisons the momentum permanently, since nothing ever resets it. The second writes NaN back into the gradient in place, so the overflow check fires again on the next step regardless of what the loss scale has been reduced to — the failure sustains itself.

Step trace, default scale:

step   1: loss_scale=65536    overflow=True  changed=0/10  momentum_finite=False
step   2: loss_scale=32768    overflow=True  changed=0/10  momentum_finite=False
step  10: loss_scale=128      overflow=True  changed=0/10  momentum_finite=False
...        loss scale reaches the minimum and the run raises

Why the suite is green

Two things hide it.

The run is too short. TestMuonConfigs takes 5 steps. The scaler needs about sixteen halvings to reach the minimum, so the exception never fires inside the test.

The parameter-change assertion cannot fail. initial_params is captured before deepspeed.initialize, which casts the model to fp16, so the comparison is fp32 against fp16 and torch.equal is False whatever happened in between:

initial_params = [p.clone().cpu() for p in model.parameters()]   # fp32, pre-initialize
...
assert not torch.equal(initial.cpu(), final.cpu())               # final is fp16

Measured:

initial dtype (pre-init)            torch.float32
after-training dtype                torch.float16
repo assertion (pre-init vs after)  10/10 "changed"
same-dtype comparison               0/10 actually changed
cast alone, no training at all      10/10 "changed" by the same assertion

The last line is the point: casting a fresh model to fp16 and training it zero steps satisfies the assertion.

Expected behaviour

A discarded step should leave the optimizer state as it was. Muon's momentum should not absorb a gradient whose step is about to be skipped, and a skipped step should not make the next one fail.

Environment

• DeepSpeed at 32110a9 (current master)
• torch 2.9.1+cu128, 2 × NVIDIA H20-3e, NCCL
• ZeRO stages 1 and 2 both reproduce; stage 3 uses a different momentum path and is not affected

Related

The same code path has a second, independent defect I am fixing separately: a restored checkpoint returns the momentum buffer in fp32 while the gradients are in the gradient accumulation dtype, so resuming a bf16 Muon run raises on momentum.lerp_(grad). Both are invisible to the current suite for the same reason — it never checkpoints and only runs fp16.

I am happy to take this one.

[alanhuangyoo]

Correcting one line in the issue body, and adding the fix result.

On ZeRO-3. I wrote that stage 3 "uses a different momentum path and is not affected". That was not measured — it came from having verified stage 3 in a bf16 run, where there is no loss scaling to overflow in the first place. When I went back to test it under fp16, my harness hit an unrelated failure of its own (RuntimeError: Function LinearFunctionForZeroStage3Backward returned an invalid gradient) on both the affected and unaffected configurations, so I have no evidence about stage 3 either way and the sentence should not have been there.

What is measured is stages 1 and 2, both of which reproduce.

The fix. Keeping the momentum out of a step the scaler discards, and letting the non-finite gradient through so the overflow is still seen, is enough. Same configuration, 30 steps, 2 × H20:

	master	with the fix
stage 1, default scale	0/10 parameters moved, momentum non-finite, dies at minimum scale	10/10 moved, momentum finite, no error
stage 2, default scale	0/10, non-finite, dies	10/10, finite, no error
stage 1/2, scale 1 (no overflow)	10/10	10/10 — unchanged

Both halves of the guard turned out to be needed. Protecting the momentum alone is not sufficient: on the non-nesterov path update is the momentum, so a protected momentum would produce a finite update for an overflowed step and the step would be applied rather than skipped. The non-finiteness has to keep propagating.

PR shortly, once the regression tests finish on both trees.