Muon runs no Newton-Schulz at ZeRO stage 0, which is the default: the plainest Muon config trains with SGD

[alanhuangyoo]

Summary

{"optimizer": {"type": "Muon"}} with no zero_optimization block trains with SGD. Newton-Schulz never runs. There is no error, no warning, and the loss goes down.

zero_optimization.stage defaults to 0 (config.py: stage: ZeroStageEnum = 0), so this is the plainest Muon configuration there is.

Measurement

Two Linear layers, one step, lr=0.02, gradient_clipping: 0 so nothing else rescales the step. The Newton-Schulz kernels in original_muon are wrapped with a counter, and the resulting weights are compared against plain SGD, w - lr * grad, computed from the gradient captured by a tensor hook.

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

1.49e-08 on fp32 is the reduction-order difference between my reference and DeepSpeed's, i.e. the weights are the SGD weights. It is seven orders of magnitude below the stage-1 row, which is what a real Muon step looks like against the same gradient. bf16's 4.88e-04 is that dtype's rounding. The fp16 row differs for a different reason — the loss-scale unscaling, which my reference does not model — but its Newton-Schulz count is still zero, so it is not Muon either.

The decisive column is the middle one. Every stage-0 row is zero.

Why

deepspeed/runtime/zero/muon/muon_optimizer.py:

    @torch.no_grad()
    def step(self, closure=None, step_id=None):
        ...
        for group in self.param_groups:
            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"])

The comment is correct about ZeRO: get_flat_partition in stage_1_and_2.py and the sub-group loop in stage3.py do call muon_update, so by the time this runs the gradient already holds the orthogonalized update and applying it is the right thing.

It is only correct when one of those ran. At stage 0 there is no ZeRO optimizer:

• fp32 leaves MuonWithAuxAdam unwrapped, and its own step is the one above
• bf16 and fp16 wrap it in FP16_UnfusedOptimizer, which knows nothing about use_muon and hands it flat fp32 partitions

In all three cases nothing applies Newton-Schulz, and p.add_(p.grad, alpha=-lr) with a raw gradient is SGD. The base class original_muon.MuonWithAuxAdam.step does implement the real update, and the subclass overrides it.

use_muon is set correctly here — the group exists, holds the four 2-D weights, and is marked use_muon=True. The tagging is not the problem; nothing consumes it.

What I think should happen

Immediately: refuse the configuration. Silently training a different optimizer than the one asked for is the worst available outcome, and an error at deepspeed.initialize costs nothing. Muon is only implemented for the wrappers that apply the update — DeepSpeedZeroOptimizer and DeepSpeedZeroOptimizer_Stage3 — so anything else should say so. This is a small, self-contained change and I am happy to send it.

Then, if wanted: implement the update for those paths. The parameters MuonWithAuxAdam.step sees are flat under every wrapper, including FP16_UnfusedOptimizer, so a shape test is not enough to tell "already updated" from "nobody updated it" — it needs the wrappers that do the work to say so, or the flat partitions to carry the shapes to unflatten by. That is a feature rather than a bug fix, so I would keep it separate.

Note that the tests do not catch this: tests/unit/ops/muon/ and tests/unit/v1/ops/muon/ parametrize over ZeRO stages 1, 2 and 3, and stage 0 is not among them. A test that counts Newton-Schulz calls, rather than checking that training progresses, would have caught it — progress is exactly what SGD also produces.

Reproduction

import torch, deepspeed
import deepspeed.runtime.zero.muon.original_muon as om

NS = {"n": 0}
for name in ("zeropower_via_gram_newtonschulz", "zeropower_via_newtonschulz5"):
    f = getattr(om, name)
    setattr(om, name, (lambda f: (lambda *A, **K: (NS.__setitem__("n", NS["n"] + 1), f(*A, **K))[1]))(f))

model = torch.nn.Sequential(torch.nn.Linear(128, 128, bias=False), torch.nn.Linear(128, 128, bias=False))
engine, _, _, _ = deepspeed.initialize(
    model=model, model_parameters=model.parameters(),
    config={"train_micro_batch_size_per_gpu": 4, "gradient_accumulation_steps": 1,
            "gradient_clipping": 0.0,
            "optimizer": {"type": "Muon", "params": {"lr": 0.02}}})   # no zero_optimization

x = torch.randn(4, 128, device=engine.device)
engine.backward(engine(x).square().sum())
engine.step()
print("newton_schulz calls:", NS["n"])     # 0

[alanhuangyoo]

Correcting something I wrote in the issue body, because it was the reason I gave for not fixing this properly and it is wrong.

I said the parameters step() sees are flat under every wrapper, FP16_UnfusedOptimizer included, so a shape test could not tell "already updated" from "nobody updated it". FP16_UnfusedOptimizer does not flatten — it keeps per-parameter fp32 clones:

# deepspeed/runtime/fp16/unfused_optimizer.py
fp32_group = [p.clone().float().detach() for p in self.fp16_groups[i]]
...
param_group['params'] = self.fp32_groups[i]

Measured, printing the ranks of the parameters in the Muon group:

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]

A flat ZeRO partition is 1-D and everything else hands over the matrix, so the shape does separate them cleanly. That makes this fixable rather than only refusable, which is the better outcome — #8442 now runs Newton-Schulz in the non-flat case instead of rejecting the configuration. Stage 0 fp32 then produces the same update as stage 1 (max|w - SGD| = 9.772e-02 on both).

The measurement in the issue body stands; only my claim about why it was hard to fix was wrong.