Muon runs no Newton-Schulz at ZeRO stage 0, the default: run it (#8442)
Closes #8441.
## The problem
`MuonWithAuxAdam.step` applies an update it assumes has been
orthogonalized already:
```python
# 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.
---------
Signed-off-by: alanhuangyoo <alanhuangyoo@gmail.com>
DeepSpeed
cbd303eb - Muon runs no Newton-Schulz at ZeRO stage 0, the default: run it (#8442)
