fix: Ensure full gradient reduction for Muon with reduce_scatter by nathon-lee · Pull Request #7808 · deepspeedai/DeepSpeed

nathon-lee · 2026-01-23T07:59:53Z

fix(zero): Ensure full gradient reduction for Muon optimizer with reduce_scatter

This commit addresses the issue where cross-partition parameters received incorrect updates when using ZeRO-1/ZeRO-2 with reduce_scatter=true and Muon optimizer. The Newton-Schulz orthogonalization in Muon requires complete gradient information, which wasn't available when reduce_scatter was enabled.

The fix introduces a check for Muon parameters and forces full all-reduce gradient reduction for these cases, ensuring consistent parameter updates across all ranks.

Closes #7807

nathon-lee · 2026-01-23T08:18:50Z

Hi, @tohtana @tjruwase Thanks for reviewing this PR which fixes the Muon gradient reduction issue with ZeRO-1/2 and reduce_scatter (#7807). to resolve cross-partition parameter inconsistencies. Let me know if any changes are needed. Thanks!

nathon-lee · 2026-01-24T01:28:02Z

@sfc-gh-truwase Thanks for the review and suggestion! I've updated the implementation to detect Muon usage during initialization and added an assertion to prevent incompatible configurations with reduce_scatter. Also simplified the average_tensor method using the pre-detected flag. Let me know if any further changes are needed!

nathon-lee · 2026-01-24T10:08:27Z

I've made simple formatting adjustments to comply with the project's YAPF style requirements, including:

Fixed indentation to use 4 spaces consistently
Adjusted line breaks for function parameters and long expressions
Removed trailing whitespace from all lines
Ensured consistent spacing around operators and parentheses
These changes only affect code formatting and do not alter any functionality.

sfc-gh-truwase · 2026-01-24T20:26:35Z

@nathon-lee please see this formatting: https://github.com/deepspeedai/DeepSpeed/blob/master/CONTRIBUTING.md#prerequisites

sfc-gh-truwase · 2026-01-25T17:10:52Z


-        self.low_precision_master_weights_and_grads = self.master_weights_and_grads_dtype != torch.float32
+        # Check for Muon optimizer usage
+        self.uses_muon = any(


I think it would be better to maintain this state on a per param group granularity.

Please see #7776 for the context.

Thank you for your valuable feedback! I appreciate you pointing out the need for per-parameter group tracking. I'll implement the Muon state management at the parameter group level as suggested and reference PR #7776 to ensure alignment with the project's architecture. Let me know if you need any further adjustments!

I think that is all the is needed. Please ping me when ready for review again. Thanks!

Signed-off-by: leejianwoo-collab <leejianwoo@gmail.com>

@sfc-gh-truwase

Use ZeRO stage 1 to use BF16 optimizer. (We should have switched to ZeRO1 in deepspeedai#7788, but I missed the change. @sfc-gh-truwase) - deepspeedai#7790 removed the fallback that allowed bf16 model + fp32 grad accumulation without ZeRO, so that combo now raises NotImplementedError. - deepspeedai#7788 changed test_bf16_optimizer_fragments to force BF16_Optimizer by setting grad_accum_dtype=fp32, but it kept ZeRO stage 0, which is now invalid after deepspeedai#7790. Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: leejianwoo-collab <leejianwoo@gmail.com>

Signed-off-by: leejianwoo-collab <leejianwoo@gmail.com>

Evoformer tests fail when we run them in parallel with other tests. ``` RuntimeError: Cannot re-initialize CUDA in forked subprocess. ``` This PR adds `@pytest.mark.sequential` to the tests. See the full test log for details: https://github.com/deepspeedai/DeepSpeed/actions/runs/21303530770/job/61326548592 Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: leejianwoo-collab <leejianwoo@gmail.com>

Fix deepspeedai#7812: This PR makes DeepSpeedEngine cleanup safe for partial initialization. This prevents destructor-time tracebacks by guarding access to unitialized attributes of DeepSpeed engine. Signed-off-by: Masahiro Tanaka <mtanaka@anyscale.com> Signed-off-by: leejianwoo-collab <leejianwoo@gmail.com>

Signed-off-by: leejianwoo-collab <leejianwoo@gmail.com>

nathon-lee · 2026-02-27T06:45:42Z

This branch has accumulated quite a few conflicts.
I’m going to create a new branch from the latest base and resolve the issue there.

@PKUWZP

## 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 deepspeedai#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 deepspeedai#7919): ```python 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 deepspeedai#7878 and deepspeedai#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 (deepspeedai#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 deepspeedai#7807 cc @PKUWZP @pengdurice (ZeRO-3 Muon guard, deepspeedai#7919) @tohtana Signed-off-by: whycoming <alwaysxd666@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>

@PKUWZP

## 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 deepspeedai#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 deepspeedai#7919): ```python 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 deepspeedai#7878 and deepspeedai#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 (deepspeedai#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 deepspeedai#7807 cc @PKUWZP @pengdurice (ZeRO-3 Muon guard, deepspeedai#7919) @tohtana Signed-off-by: whycoming <alwaysxd666@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com> Signed-off-by: nathon-lee <leejianwoo@gmail.com>

@PKUWZP

## 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 deepspeedai#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 deepspeedai#7919): ```python 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 deepspeedai#7878 and deepspeedai#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 (deepspeedai#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 deepspeedai#7807 cc @PKUWZP @pengdurice (ZeRO-3 Muon guard, deepspeedai#7919) @tohtana Signed-off-by: whycoming <alwaysxd666@gmail.com> Co-authored-by: Ma, Guokai <guokai.ma@gmail.com>

nathon-lee requested review from tjruwase and tohtana as code owners January 23, 2026 07:59

nathon-lee force-pushed the fix_issue_7807 branch from bc2d301 to ceb84ba Compare January 23, 2026 08:00

sfc-gh-truwase reviewed Jan 23, 2026

View reviewed changes

Comment thread deepspeed/runtime/zero/stage_1_and_2.py Outdated

nathon-lee force-pushed the fix_issue_7807 branch from 08f4845 to 44fc221 Compare January 24, 2026 01:29

sfc-gh-truwase reviewed Jan 24, 2026

View reviewed changes

Comment thread deepspeed/runtime/zero/stage_1_and_2.py Outdated

nathon-lee force-pushed the fix_issue_7807 branch from 9ee47b4 to e0248ce Compare January 24, 2026 04:47

nathon-lee requested a review from loadams as a code owner January 24, 2026 04:47

sfc-gh-truwase approved these changes Jan 24, 2026

View reviewed changes

nathon-lee force-pushed the fix_issue_7807 branch from 2a0f659 to 1979f00 Compare January 24, 2026 10:01

sfc-gh-truwase reviewed Jan 25, 2026

View reviewed changes

nathon-lee and others added 8 commits January 25, 2026 21:32

fix: Ensure full gradient reduction for Muon with reduce_scatter

e5c6d1e

Signed-off-by: leejianwoo-collab <leejianwoo@gmail.com>

Update stage_1_and_2.py

fcdc7eb

Signed-off-by: leejianwoo-collab <leejianwoo@gmail.com>

Update stage_1_and_2.py

d87433f

Signed-off-by: leejianwoo-collab <leejianwoo@gmail.com>

fix: update 1 file reformatted.

ff88670

Signed-off-by: leejianwoo-collab <leejianwoo@gmail.com>

Refactor gradient reduction logic for Muon parameters

6c06319

Signed-off-by: leejianwoo-collab <leejianwoo@gmail.com>

nathon-lee force-pushed the fix_issue_7807 branch from aba6a8d to 6c06319 Compare January 26, 2026 02:33

tohtana and others added 2 commits January 27, 2026 15:55

Merge branch 'master' into fix_issue_7807

2e73be8

Merge branch 'deepspeedai:master' into fix_issue_7807

8ad2f70

Copilot AI mentioned this pull request Jan 28, 2026

Fix Muon unit tests for reduce_scatter incompatibility assertion nathon-lee/DeepSpeed_woo#3

Closed

Merge branch 'master' into fix_issue_7807

19f6149

fy817 mentioned this pull request Jan 29, 2026

Fix Muon optimizer conflict with gradient clipping in ZeRO 1/2 #7776

Open

nathon-lee closed this Feb 27, 2026

whycoming mentioned this pull request Jun 25, 2026

Reject Muon optimizer with reduce_scatter in ZeRO-1/2 #8090

Merged

Navigation Menu

fix: Ensure full gradient reduction for Muon with reduce_scatter - #7808

fix: Ensure full gradient reduction for Muon with reduce_scatter#7808
nathon-lee wants to merge 11 commits into
deepspeedai:masterfrom
nathon-lee:fix_issue_7807

nathon-lee commented Jan 23, 2026

Uh oh!

nathon-lee commented Jan 23, 2026 •

edited

Loading

Uh oh!

Uh oh!

nathon-lee commented Jan 24, 2026

Uh oh!

Uh oh!

nathon-lee commented Jan 24, 2026

Uh oh!

sfc-gh-truwase commented Jan 24, 2026

Uh oh!

sfc-gh-truwase Jan 25, 2026

Uh oh!

sfc-gh-truwase Jan 25, 2026

Uh oh!

nathon-lee Jan 26, 2026

Uh oh!

sfc-gh-truwase Jan 28, 2026

Uh oh!

nathon-lee commented Feb 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

nathon-lee commented Jan 23, 2026

Uh oh!

nathon-lee commented Jan 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

nathon-lee commented Jan 24, 2026

Uh oh!

Uh oh!

nathon-lee commented Jan 24, 2026

Uh oh!

sfc-gh-truwase commented Jan 24, 2026

Uh oh!

sfc-gh-truwase Jan 25, 2026

Choose a reason for hiding this comment

Uh oh!

sfc-gh-truwase Jan 25, 2026

Choose a reason for hiding this comment

Uh oh!

nathon-lee Jan 26, 2026

Choose a reason for hiding this comment

Uh oh!

sfc-gh-truwase Jan 28, 2026

Choose a reason for hiding this comment

Uh oh!

nathon-lee commented Feb 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

nathon-lee commented Jan 23, 2026 •

edited

Loading