Skip to content

nightly-last-green

nightly-last-green: Keep Muon's update at full size under fp16 loss scaling (#8655)
@alanhuangyoo alanhuangyoo tagged this 30 Sep 05:59 
Under fp16, Muon parameters barely train with ZeRO-1/2/3 unless the
optimizer is offloaded. Newton-Schulz returns the same update at any
loss scale, but the step then treats it like a scaled gradient:
`unscale_and_clip_grads` divides it by the loss scale, and the global
norm takes it as scaled too. So a Muon parameter moves by 1/loss_scale
of its update, 1/65536 with the default dynamic loss scaler.

On 2×H20, the norm of one Muon step on a small model, static loss scale
1 vs 1024, no clipping:

| | scale 1 | scale 1024 |
|---|---|---|
| ZeRO-1/2/3, master | 4.47e-2 | 4.37e-5 |
| ZeRO-1/2/3, this PR | 4.47e-2 | 4.47e-2 |
| ZeRO-1/2/3 with CPU offload, master | 4.47e-2 | 4.47e-2 |

And 60 steps of a small regression with the default fp16 config (dynamic
loss scale, clipping 1.0), ZeRO-1 and ZeRO-3 alike:

| | loss at step 0 → 59 |
|---|---|
| bf16 | 0.403 → 0.295 |
| fp16, master | 0.403 → 0.385 |
| fp16, this PR | 0.403 → 0.295 |

It's the same with gradient accumulation (gas=2: bf16 and this PR 0.362
→ 0.320, master 0.362 → 0.418), and with an initial scale of 2^30 that
overflows for the first steps this PR trains normally afterwards with no
NaN.

The CPU-offload path from #8464 already handles this: it unscales before
Newton-Schulz and scales the update back before the norm. Without
offload the update sits in the fp16 gradient buffers, where scaling it
back up could overflow, so this does it after the cast to fp32, for the
norm and before `unscale_and_clip_grads`. In the 2^30 run the dynamic
scaler settles at 2^22, so scaling the update up in fp16 would overflow
it. The momentum keeps the scaled gradient as before, which in fp16
keeps small gradients out of the subnormal range. An invalid group norm
stays at its -1 sentinel. bf16 and fp32 have a loss scale of 1 and don't
reach the new code.

`TestMuonLossScaling`, next to the offload test, checks ZeRO-1/2/3 with
clipping on: on master the update under scale 1024 is 3e-3 of the one
under 1, here they match. The Muon tests, the fp16 dynamic loss scale
tests and the ZeRO cases of `test_fp16.py` pass on 2×H20 (294 passed, 8
skipped).

Not covered: ZenFlow and SuperOffload, which have their own step. Also,
with fp16 the momentum is in scaled units here but unscaled under
offload, so a universal checkpoint moved between the two would carry it
off by the loss scale; that was already the case before this.

---------

Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>
Assets 2
Loading