[muon] Keep the momentum out of steps the loss scaler discards

[alanhuangyoo]

Fixes #8432.

Problem

Muon under ZeRO 1/2 with fp16 does not train. Running the configuration tests/unit/ops/muon/test_muon.py itself uses, for longer than its five steps:

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

No parameter is ever updated. The first loss-scale overflow is folded into Muon's momentum before the overflow check decides to discard the step, and the buffer stays non-finite for the rest of the run.

The failure sustains itself because both halves of muon_update touch the gradient:

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

so the next step's gradient is already non-finite whatever the loss scale has been reduced to. Backing off cannot help.

Fix

The momentum stays out of a step whose gradient is not finite, and the non-finite gradient is still returned so the overflow is seen and the step skipped.

Both halves are needed. Protecting the momentum alone is not enough: on the non-nesterov path update is the momentum, so a protected momentum would hand an overflowed step a finite update and the step would be applied rather than skipped. The non-finiteness has to keep propagating. Evaluated on device, so this costs no synchronization.

Verification

30 steps, 2 × H20, ZeRO 1/2, SimpleModel(hidden_dim=128, nlayers=5), lr=0.05:

stage	initial scale	master	this branch
1	65536 (default)	0/10 moved, momentum non-finite, dies	10/10 moved, momentum finite
2	65536 (default)	0/10 moved, momentum non-finite, dies	10/10 moved, momentum finite
1	1 (no overflow)	10/10 moved, finite	10/10 moved, finite
2	1 (no overflow)	10/10 moved, finite	10/10 moved, finite

The scale-1 rows are the control: with no overflow there was never a problem, so the failure is entirely the overflow interaction rather than anything about Muon.

Tests

tests/unit/ops/muon/test_muon_overflow.py. Three of the four fail on the parent commit:

master   3 failed, 1 passed
           FAILED test_an_overflowed_gradient_does_not_enter_the_momentum
                  - AssertionError: an overflowed step must not move the momentum
           FAILED test_training_recovers_from_the_initial_overflow[1]
                  - Exception: Current loss scale already at minimum
           FAILED test_training_recovers_from_the_initial_overflow[2]
                  - Exception: Current loss scale already at minimum
this PR  4 passed

The unit-level case pins both halves directly: the momentum must not move, and the returned update must stay non-finite. test_a_finite_gradient_still_moves_the_momentum is the guard against the guard — it would catch a fix that simply disabled the optimizer.

Existing suite on this branch, non-offload configurations:

2 failed, 80 passed

Both failures are op_builder.builder.CUDAMismatchException in test_muon_reduce_scatter_with_optimizer_offload_raises, from this box's system CUDA not matching the one torch was built against, so CPUAdam will not build. They are unrelated to this change and reproduce on master.

Note on the suite

Worth recording, since it is why this survived: TestMuonConfigs captures initial_params before deepspeed.initialize, which casts the model to fp16. The assertion is therefore fp32 against fp16 and torch.equal is False whatever happened in between:

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

Casting a fresh model to fp16 and training it zero steps satisfies it. That is out of scope here — the new file captures after initialize and says why in a comment — but the assertion is worth tightening separately.

[ebarkhordar]

grad_is_finite is per tensor, but the decision it protects the momentum from is global. has_overflow (stage_1_and_2.py:2482) sums _has_inf_or_nan over every partitioned gradient in every group and all-reduces MAX across the DP and model-parallel groups, and step() discards the whole step on that one flag at :2307. The ordering is not in doubt: the flag is computed from averaged_gradients, which is what get_flat_partition returns, and that is where muon_update is applied per tensor at :2167.

So on a step discarded because some other matrix overflowed, this matrix's gradient was finite, its momentum has already moved, and the parameter update it moved for is thrown away. The momentum then carries a step that never happened. That is narrower than the invariant your test module states ("A step the loss scaler discards must leave Muon's momentum as it was"), and it is the case the tests do not reach: both tensor-level tests use a single tensor, and the training test only asserts the momentum is finite, which the mixed case satisfies.

Is the per-tensor scope deliberate? One absorbed gradient is cheap next to the poisoned-forever behaviour you are fixing, so this is not a blocker either way. I read this rather than ran it, and I have no multi-GPU box here, so I have not seen the mixed case happen.

[alanhuangyoo]

Yes, deliberate — but you are right that the module claimed more than it delivers, and that is fixed in 4ce8e62. Apologies for not answering here; I pushed it and never replied in the thread.

Two changes, both from your reading:

The docstring no longer states the invariant you quoted. It is now "A tensor whose own gradient overflowed must not absorb it into its momentum", and it says the scope out loud — that the guard is per tensor, that has_overflow reduces over every partitioned gradient and step() discards on that one flag, and that a finite tensor therefore still advances its momentum on a step discarded for another.

And TestMuonMixedOverflow::test_a_finite_tensor_still_absorbs_a_step_discarded_for_another pins that case rather than leaving it to reading. You said you have no multi-GPU box, so here is the measurement — one finite matrix (calm) and one whose gradient overflows (boom), momentum read where muon_update writes it (inside get_flat_partition, during backward) rather than around engine.step():

this branch:  step 1  overflow=True  calm 55.159618 -> 104.801094   boom 42.780441 -> 42.780441  discarded
master:       step 1  overflow=True  calm 55.159618 -> 104.801094   boom 42.780441 -> inf
              step 2  overflow=True  calm 104.801094 -> 149.487137  boom inf -> nan
              step 3  overflow=True  calm 149.487137 -> 162.125839  boom nan -> nan

calm moves on a discarded step on both — that is the residue you describe, and it is the same on master, so this PR does not make it worse. What changes is the boom column: inf -> nan -> nan forever on master versus held at its pre-overflow value here.

On whether to close the residue too: it needs the momentum write deferred until the step is known to survive, which costs a second buffer the size of the momentum for every Muon parameter. That did not seem worth trading for one absorbed gradient per overflow event, so the docstring records the choice instead of hiding it. Happy to be argued out of that if you or a maintainer would rather pay the memory.

[alanhuangyoo]

You are right, and it was not deliberate — I put the guard where muon_update already was and did not think about the flag being global.

Tracing it the same way you did, the ordering is forced, not incidental:

• self.averaged_gradients[i] = self.get_flat_partition(...) (:956, :989) — muon_update runs inside get_flat_partition (:2167), so the momentum moves at gradient-reduction time.
• has_overflow_partitioned_grads_serial (:2474) reads self.averaged_gradients[i], i.e. after that.
• step() discards on the result at :2307.

So the global flag is computed from post-muon_update gradients by construction. There is no ordering of the current code that lets a per-tensor call see it.

What a global guard would actually cost, since that is the part your comment leaves open:

1. Check overflow before applying Muon. Means an extra all-reduce inside the reduction path, on every step, to buy a rollback that only matters on discarded steps. Wrong trade at fp16 loss-scale frequencies.
2. Invert the update. momentum.lerp_(grad, 1 - beta) inverts exactly in real arithmetic ((m_new - (1-beta)*g) / beta), and in the mixed case this tensor's grad is finite so it is well-defined — but it is not bit-exact in fp32, so "leave the momentum as it was" would become "leave it approximately as it was".
3. Defer the write. Compute the update from a copy and commit the momentum only once the step is known to survive. Exact, no extra collective, costs one buffer the size of the momentum — which is already a full copy of the Muon parameters, so roughly +1× on those groups.

(3) is the only one that delivers what my test module claims. Whether that memory is worth it for a case that costs one absorbed gradient is a call I would rather you and @delock make than assume.

Two things I am doing either way:

• The invariant in the test module is wrong as written. "A step the loss scaler discards must leave Muon's momentum as it was" is stronger than a per-tensor guard delivers. I will narrow it to what is actually true — a tensor whose own gradient overflowed does not absorb it — and add the mixed case explicitly, so the gap is recorded rather than implied.
• Measure it. You said you read this rather than ran it and have no multi-GPU box; I do. Two 2-D parameters, only one fed an overflowing input, fp16 + ZeRO-1, and check whether the finite one's momentum moves on the discarded step. My box is unreachable right now, so this is a promise rather than a result — I will post the numbers, including if they contradict the reading.

Thanks for reading it this closely. This is the second time on this stack I have claimed something wider than I measured, and both times it was someone else who noticed.

[alanhuangyoo]

Numbers, as promised. Your case reproduces exactly.

Two 2-D parameters in one group, calm and boom. Only boom is fed an input that overflows in fp16, at step 1. Muon, fp16, ZeRO-1, single GPU. The momentum is read from the group's flat partition buffer before and after each backward, because that is where muon_update writes — not inside step().

On this branch:

step  overflow   calm momentum (before bwd -> after bwd)   boom momentum (before -> after)   params
   0     False          None -> 55.159618                       None -> 42.780441            changed
   1      True     55.159618 -> 104.801094                 42.780441 -> 42.780441            NO (discarded)
   2     False    104.801094 -> 149.487137                 42.780441 -> 81.290314            changed
   3     False    149.487137 -> 189.690079                 81.290314 -> 115.945297           changed

Step 1 is discarded and the parameters do not move, boom is protected — and calm moves anyway, 55.159618 -> 104.801094, for an update that is thrown away. That is your case, and it is not hypothetical.

On master, same script:

step  overflow   calm momentum                            boom momentum                     params
   0     False          None -> 55.159618                       None -> 42.780441            changed
   1      True     55.159618 -> 104.801094                 42.780441 -> inf                  NO (discarded)
   2      True    104.801094 -> 149.487137                        inf -> nan                 NO (discarded)
   3      True    149.487137 -> 162.125839                        nan -> nan                 NO (discarded)

So the same run shows both things at once: the poisoning this PR is for (inf -> nan, every later step discarded, loss scale halving 16 -> 8 -> 4, no recovery), and the gap you found (calm advancing on a discarded step, identically on both sides — the fix does nothing for it).

What I am changing here: the invariant in the test module, which as written promises more than the code delivers, and a test that pins the mixed case at the behaviour above so it is recorded rather than discovered again.

What I am not changing without a word from you and @delock: the scope. Of the three ways to make it global, only deferring the momentum write is exact, and it costs one buffer the size of the momentum. One absorbed gradient per discarded step against +1x optimizer memory on the Muon groups is a trade I would rather not make unilaterally inside a PR that is meant to stop a run from dying.

Thanks — you found this by reading, without a box to run it on, and you were right on every detail including which tests could not reach it.

[ebarkhordar]

Thanks for running it. Your step 1 row is the case exactly: calm moving 55.159618 to 104.801094 on a step whose parameter update is discarded, and identical on both sides, so the guard does not touch it.

On scope, since you asked. Deferring the momentum write is the only one of the three that delivers what the test module claims. I would avoid inverting the lerp for the reason you gave, that it turns an exact invariant into an approximate one, and a rollback bought with an extra collective on every step is paying on the common path for the rare one. Whether the buffer is worth one absorbed gradient per discarded step is a memory call I have no numbers for, so that part is yours and @delock's.

Narrowing the invariant and pinning the mixed case is worth doing either way.

[alanhuangyoo]

Done, in 4ce8e62 — the uncontested half.

The module docstring now says what the guard covers (a tensor whose own gradient overflowed does not absorb it) instead of what it does not (a step the scaler discards leaving the momentum as it was), and spells out why the two differ.

test_a_finite_tensor_still_absorbs_a_step_discarded_for_another pins the case: two 2-D parameters in one group, only one fed an overflowing input, fp16 + ZeRO-1, momentum read on either side of backward. It asserts that the overflowing tensor's momentum holds, that no parameter moves, and that the finite tensor's momentum does advance — with a message saying that if this starts failing, the guard has become global and the docstring should follow.

So the gap is a failing test away from being noticed rather than a paragraph someone has to remember.

tests/unit/ops/muon/test_muon_overflow.py: 3 passed, 2 skipped (fp16 world_size=2 cases need the second GPU).

Leaving the scope decision where you put it — @delock, the question is whether one absorbed gradient per discarded step is worth a second buffer the size of the momentum on the Muon groups. I have no preference strong enough to spend your memory budget on it.