Muon is silently disabled under ZeRO-3 when the model is built with zero.Init

[alanhuangyoo]

Muon is silently inactive under ZeRO-3 whenever the model is built with deepspeed.zero.Init. Training runs, the loss falls, and every parameter is on the AdamW branch.

What happens

set_optimizer_flags decides use_muon per parameter:

if p.ndim >= 2 and not any(keyword in name.lower() for keyword in ("embed", "lm_head")):

The test is right — Muon is defined on matrices — but p.shape is not the layer's shape under zero.Init. ZeRO-3 replaces a partitioned parameter's data with a flat placeholder and records the shape it has as a layer in ds_shape:

q_proj: shape=(0,) ndim=1 ds_shape=torch.Size([256, 256]) ds_numel=65536 status=ZeroParamStatus.NOT_AVAILABLE

So every parameter in the model reports as 1-D, none are tagged, and stage3.py's

if getattr(param_group['params'][0], 'use_muon', False):

finds no sub-group using Muon. muon_update is never reached.

Measurement

2-layer Llama, ZeRO-3, {"optimizer": {"type": "Muon"}}, muon_update wrapped with a counter, 6 steps. Only the model construction differs between the two rows:

	parameters with use_muon	muon_update calls
model built normally	14	84
model built under zero.Init	0	0

With this change the second row reads 14 and 84, and the loss falls over the 6 steps (4.8924 -> 4.1229).

zero.Init is how a model that does not fit on one device is built, which is the case ZeRO-3 exists for, so this is not an edge configuration.

The change

Read the shape the parameter has as a layer — ds_shape when ZeRO-3 has recorded one, shape otherwise — the same thing _shape_before_zero3_partition in module_inject/layers.py already does for AutoTP. Nothing else changes: the rank test and the embed/lm_head exclusions are unchanged, and a genuine 1-D parameter is still excluded because ds_shape is read for its rank rather than assumed to be a matrix.

Tests

tests/unit/runtime/zero/test_muon_use_muon_under_zero_init.py, 5 cases. Three are CPU-only and build the ZeRO-3 shape directly (a flat placeholder plus ds_shape): matrices are tagged unpartitioned, matrices are still tagged once partitioned, and a partitioned vector is not promoted. Two run a real zero.Init and deepspeed.initialize under ZeRO-3 and assert muon_update is reached.

On master, test_muon_update_is_called[True] fails with

AssertionError: Muon never ran; every parameter was left on the AdamW branch

and test_matrices_are_still_tagged_once_zero3_has_partitioned_them fails too; the other three pass on both sides, which is what pins that nothing widened. 5 passed with the change. yapf and flake8 clean.

[alanhuangyoo]

Note on the red check: the modal-torch-latest / DeepSpeedAI CI job was cancelled at the 1h30m job timeout, not failed — the API reports "conclusion": "cancelled" for it, and there is no FAILED line anywhere in the log. collect tests and DCO pass on the same commit.

Three other branches of mine hit the same timeout today while four went green, so it is the runner rather than this tree. I cannot re-run it myself; any maintainer re-running the job should be enough.

[alanhuangyoo]

Ran this together with the other two Muon fixes I have open (#8440, #8442) on a real Llama, since all three touch optimizer setup and I wanted to know they compose. Merged onto master they apply cleanly, and:

zero.Init + stage 3   use_muon=14  NS calls=56  group weight_decays=[0.0, 0.1]  loss 5.18 -> 4.24  finite
plain     + stage 3   use_muon=14  NS calls=56  group weight_decays=[0.0, 0.1]  loss 4.92 -> 3.34  finite
plain     + stage 2   use_muon=14  NS calls=56  group weight_decays=[0.0, 0.1]  loss 4.82 -> 3.26  finite

14 Muon matrices over 4 steps is 56 orthogonalizations, which is the right count, and the two weight decays are the ones the caller asked for rather than the config's. The zero.Init row is the one this PR is about: on master it is use_muon=0.

zero.Init with stage 1 fails at RuntimeError: 'weight' must be 2-D in the forward pass, but that is pre-existing and has nothing to do with the optimizer — the identical run on master fails the same way, and so does AdamW. zero.Init is a stage-3 mechanism.

The full Muon suite on the merged tree is running; I will post the number when it finishes.

[alanhuangyoo]

The merged-tree number I promised: 190 passed in 28:51.

That is tests/unit/ops/muon/ and tests/unit/v1/ops/muon/ plus the three PRs' own test files, run against master with #8438, #8440 and #8442 all applied. No failures, no skips beyond the usual world-size ones.

Re-ran it after #8442 changed shape (it now implements the stage-0 Newton-Schulz rather than refusing the configuration), so the number covers the current state of all three.

[alanhuangyoo]

Rebased onto current master (4 commits since this branch was cut) — no conflicts, and it re-triggers the CI job that was cancelled at the 1h30m timeout on the previous run, so the red check should resolve on its own.

5 passed on the rebased branch; yapf and flake8 clean.