Skip to content

Keep Muon's update at full size under fp16 loss scaling - #8655

Open
alanhuangyoo wants to merge 1 commit into
deepspeedai:masterfrom
alanhuangyoo:fix/muon-fp16-loss-scale
Open

alanhuangyoo wants to merge 1 commit into
deepspeedai:masterfrom
alanhuangyoo:fix/muon-fp16-loss-scale

Conversation

@alanhuangyoo

Copy link
Copy Markdown
Contributor

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.

Newton-Schulz returns the same update at any loss scale, but without
optimizer offload ZeRO-1/2/3 then treated it like a gradient: the global
norm took it as scaled and unscale_and_clip_grads divided it by the loss
scale. Under fp16 a Muon parameter moved by 1/loss_scale of its update
(1/65536 by default) and effectively stopped training. Scale it back in
fp32 for both, as the CPU-offload path already does.

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

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants