Muon runs no Newton-Schulz at ZeRO stage 0, the default: run it

[alanhuangyoo]

Closes #8441.

The problem

MuonWithAuxAdam.step applies an update it assumes has been orthogonalized already:

# deepspeed/runtime/zero/muon/muon_optimizer.py
if group["use_muon"]:
    # we move the muon update part to the deepspeed's optimizer since the parameter here is a flat version
    # thus not suitable for muon update
    for p in group["params"]:
        p.mul_(1 - group["lr"] * group["weight_decay"])
        p.add_(p.grad.reshape(p.shape), alpha=-group["lr"])

That holds under ZeRO: get_flat_partition in stage_1_and_2.py and the sub-group loop in stage3.py call muon_update, so by then the gradient holds the orthogonalized update.

With no ZeRO optimizer nothing does, and p.add_(p.grad, alpha=-lr) on a raw gradient is SGD. zero_optimization.stage defaults to 0, so a config that just names Muon gets that. Counting the Newton-Schulz kernel calls on one step:

config	wrapper	Newton-Schulz calls	max|w - SGD|
no zero_optimization block, fp32	MuonWithAuxAdam	0	1.49e-08
stage: 0, bf16	FP16_UnfusedOptimizer	0	4.88e-04
stage: 0, fp16	FP16_UnfusedOptimizer	0	—
stage: 1, fp32	DeepSpeedZeroOptimizer	2	9.77e-02

1.49e-08 is reduction ordering: those are the SGD weights, seven orders below what a real Muon step does to the same gradient. Training runs and the loss falls either way.

The change

The two cases are distinguishable by shape, which I initially thought they were not. ZeRO hands step() a flat 1-D partition. An unwrapped optimizer hands it the model's weight, and FP16_UnfusedOptimizer hands it a per-parameter fp32 clone — p.clone().float().detach(), same shape — not a flat buffer. Measured:

stage 0 / fp32   MuonWithAuxAdam         ndims=[2]
stage 0 / bf16   FP16_UnfusedOptimizer   ndims=[2]
stage 0 / fp16   FP16_UnfusedOptimizer   ndims=[2]
stage 1 / fp32   DeepSpeedZeroOptimizer  ndims=[1]

So: orthogonalize when the parameter is a matrix, and keep applying the update as-is when it is a partition. After the change, stage 0 fp32 produces max|w - SGD| = 9.772e-02 — the same value stage 1 gives, i.e. the same update.

Newton-Schulz is scale-invariant and the momentum starts at zero, so initialize_optimizer_states' warm-up step on zero gradients stays a no-op.

num_heads is deliberately not threaded through here: it does not exist on muon_update on master. Once #8384 lands, this call site is where per-head would be added for the unwrapped path.

Tests

tests/unit/runtime/zero/test_muon_without_zero_optimizer.py, 7 cases: Newton-Schulz runs at stage 0 for fp32, bf16 and fp16 — all three wrappers; it runs for a config with no zero_optimization block, which is the plainest form; and it still runs on stages 1, 2 and 3.

On master, 4 fail and 3 pass. The four that fail are the stage-0 ones, with Newton-Schulz ran 0 times for two Muon matrices; the three that pass are the ZeRO stages, which is the control that says the test measures the right thing.

The counter is started after deepspeed.initialize, because FP16_UnfusedOptimizer steps once at construction to allocate state and that call would otherwise satisfy the assertion on its own. It is also patched inside the test body rather than in a fixture, since DistributedTest runs the body in a worker a parent-process fixture would not reach.

This is the assertion the existing Muon tests were missing: tests/unit/ops/muon/ parametrizes stages [1, 2, 3] and checks that the loss moves, which SGD also does — which is why stage 0 went unnoticed.

7 passed. yapf and flake8 clean.

[alanhuangyoo]

Checked the obvious follow-up question given #8443, which reports ZeRO-3 running Newton-Schulz once per micro-batch: this path does not have that problem. Two optimizer steps, two Muon matrices, so 4 orthogonalizations is correct:

no zero_optimization  gas=1   newton_schulz calls=4   expected 4
no zero_optimization  gas=4   newton_schulz calls=4   expected 4
stage 0               gas=1   newton_schulz calls=4   expected 4
stage 0               gas=4   newton_schulz calls=4   expected 4

MuonWithAuxAdam.step is reached once per optimizer step here, since engine.step() no-ops until the accumulation boundary, so the momentum advances once per step as configured. ZeRO-3's problem is that its update sits in the IPG reduce path, which is a different call site with no boundary guard.

[alanhuangyoo]

Ran the full Muon suite with this PR merged onto master alongside the other two open Muon fixes (#8438, #8440), since all three touch optimizer setup: 190 passed in 28:51 — tests/unit/ops/muon/, tests/unit/v1/ops/muon/, and the three PRs' own test files. The three merge cleanly onto master and onto each other.

[alanhuangyoo]

Before you spend time on this one — I found a hole in it myself and want to flag it rather than have you find it.

I claimed the two cases were distinguishable by shape: a matrix is the weight and gets orthogonalized here, a 1-D tensor is a ZeRO partition whose update the ZeRO optimizer already applied. That is true for every wrapper except one. BF16_Optimizer also replaces the param groups with flat partitions —

# deepspeed/runtime/bf16_optimizer.py
param_group['params'] = [self.fp32_groups_flat_partition[i]]

— and knows nothing about use_muon, so it never applies the update. My shape test reads its partitions as "already done" and the step stays SGD. Enumerating every wrapper a Muon config can select, on master:

s0-fp32        MuonWithAuxAdam              ndims=[2]  NS=0
s0-bf16        FP16_UnfusedOptimizer        ndims=[2]  NS=0
s0-fp16        FP16_UnfusedOptimizer        ndims=[2]  NS=0
s1-bf16        DeepSpeedZeroOptimizer       ndims=[1]  NS=2
s1-bf16-ga32   BF16_Optimizer               ndims=[1]  NS=0     <- this one
s1-fp16        DeepSpeedZeroOptimizer       ndims=[1]  NS=2
s2-bf16        DeepSpeedZeroOptimizer       ndims=[1]  NS=2
s3-bf16        DeepSpeedZeroOptimizer_S3    ndims=[]   NS=2

s1-bf16-ga32 is bf16 with grad_accum_dtype: fp32 at stage 1 — a normal enough recipe, and it was already broken before this PR, in the same silent way. Its weights come out at max|w - SGD| = 4.9e-04, which is bf16 rounding away from plain SGD, against 6.8e-02 for the same config one flag apart.

I have pushed a commit that refuses that combination at deepspeed.initialize rather than leaving it silently wrong. It cannot be fixed in step — the original shapes are not recoverable from a flat partition — so implementing Muon inside BF16_Optimizer would be its own change, and I would rather not fold it in here.

So the PR now reads: orthogonalize where the parameters are the real weights, refuse the one place where they are not and nobody else does the work. Two tests added for it, including the neighbouring config without grad_accum_dtype, so the refusal is shown to be narrow.

Sorry for the churn on the one I asked you to look at first.