Reject Muon optimizer with reduce_scatter in ZeRO-1/2

[whycoming]

Summary

ZeRO-1/2 silently produces incorrect, rank-divergent parameter updates when the Muon optimizer is used together with reduce_scatter (the default). This adds an explicit error at initialization, mirroring the existing ZeRO-3 guard, and includes a regression test. Closes #7807.

Root cause

Muon's Newton-Schulz orthogonalization is a whole-matrix operation: the rank that updates a parameter must hold that parameter's complete, fully-reduced gradient matrix, then take its partition slice of the orthogonalized result.

• get_flat_partition() (deepspeed/runtime/zero/stage_1_and_2.py) applies muon_update() to each parameter's gradient reshaped to its full 2-D shape, and only then narrows to this rank's partition.
• With reduce_scatter=True, average_tensor() reduce-scatters the gradients: each rank receives the averaged values only for its own partition slice. For the rest of a parameter whose flattened gradient crosses a partition boundary, the rank still holds its local, un-all-reduced gradient.
• So for any cross-partition parameter, no rank holds the full reduced matrix. muon_update orthogonalizes a partly-reduced, rank-divergent matrix, and each rank silently applies a different, incorrect update. Parameters that lie wholly inside one partition are unaffected — exactly matching the report.

ZeRO-3 already guards this exact conflict in deepspeed/runtime/zero/stage3.py (added in #7919):

if self.use_muon and self.reduce_scatter:
    raise ValueError("Muon and reduce scatter cannot be used together")

ZeRO-1/2 had no equivalent. The existing Muon unit tests pin "reduce_scatter": false everywhere, which implicitly acknowledges the path is unsupported but never enforces it for users — and since reduce_scatter defaults to true, a default Muon + ZeRO-1/2 run is silently wrong.

Fix

Mirror the ZeRO-3 guard in ZeRO-1/2: raise the same ValueError at initialization when the optimizer is MuonWithAuxAdam and reduce_scatter is enabled. To run Muon under ZeRO-1/2, set "reduce_scatter": false (as the Muon tests already do). The change is the import plus the guard, with no other behavioral change.

Verification (2x RTX 4090, torch 2.9.1+cu128, ZeRO stage 1 and 2)

• Before: deepspeed.initialize with Muon + reduce_scatter=true succeeds silently. With world_size=2 and a model sized so a 2-D weight straddles the gradient-partition boundary, that weight's post-step update diverges by ~0.67 in relative Frobenius norm from the correct full-gradient result, while wholly-owned weights are unaffected — confirming the silent cross-partition corruption.
• After: the same configuration raises ValueError: Muon and reduce scatter cannot be used together for both ZeRO stage 1 and 2. The existing Muon tests (which use reduce_scatter: false) remain green.

Notes

This supersedes #7878 and #7808, which aimed at the same issue by trying to force a full all-reduce for Muon but ended up with a self-contradictory guard. Aligning ZeRO-1/2 with the merged ZeRO-3 behavior (#7919) keeps the two code paths consistent and turns silent numerical corruption into a clear, actionable error.

A follow-up PR adds a numerical-correctness regression test for the supported reduce_scatter: false Muon path, since the current Muon tests only assert that parameters changed.

Closes #7807

cc @PKUWZP @pengdurice (ZeRO-3 Muon guard, #7919) @tohtana

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 3aa6f82598

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Only reject active Muon parameter groups

This check treats every MuonWithAuxAdam instance as unsafe, but the ZeRO-1/2 whole-matrix Muon path is only taken for params/groups with use_muon=True in get_flat_partition(). When users select the muon optimizer for a model whose parameters are all excluded from Muon (for example embeddings/lm_head/1-D params) or pass a MuonWithAuxAdam with all groups marked use_muon=False, training falls back to the auxiliary Adam path, which is elementwise and compatible with reduce_scatter; because reduce_scatter defaults to true, those valid runs now fail during initialization. Gate the error on an active use_muon group instead of the optimizer class alone.

Useful? React with 👍 / 👎.

[whycoming]

get_flat_partition does gate the Newton-Schulz path on use_muon, so strictly the hazard only exists when an active Muon group is present. I kept the condition as isinstance(MuonWithAuxAdam) here to mirror the merged ZeRO-3 guard (stage3.py, #7919), which uses the same check — so the two stages stay consistent, and it matches the conservative style of the existing reduce_scatter guards (e.g. the MoE assertion in stage_1_and_2.py).

The only configuration this over-rejects is a MuonWithAuxAdam whose groups are all use_muon=False — i.e. selecting Muon for a model with no Muon-eligible 2-D params — which is degenerate, and the remedy (reduce_scatter: false) is identical to every other Muon run, so the practical cost is nil.

If you'd prefer the precise gating (reject only when an active use_muon group exists), I'm happy to apply it to both ZeRO-1/2 and ZeRO-3 so they remain consistent — just let me know.

[delock]

@whycoming I agree with you that there is no need for precise gating which would only complicate the mechanism. Besides. keeping zero1/2/3 with same behavior could avoid surprise when switch from zero 1/2 to zero3.

[delock]

@whycoming Thanks for looking into this bug and provide the fix! I have put your fix and test case into the merge queue. Thanks!