Support Muon in BF16_Optimizer with FP32 gradient accumulation

[0z5a]

BF16 weights with FP32 gradient accumulation and ZeRO stage 1 select BF16_Optimizer. Muon currently refuses that configuration because this wrapper hands the optimizer flat partitions without first applying Newton–Schulz. This change computes Muon updates on the original matrices after gradient reduction and global clipping, then copies only each rank's owned intersections into the optimizer update.

Momentum remains FP32 and partitioned like the master weights. The implementation gathers one group's momentum as temporary workspace, stages local momentum until the optimizer step, and excludes alignment padding. Auxiliary Adam groups retain their existing update path. The initial support boundary is dense data parallelism with two-dimensional Muon matrices; model-parallel MPU, expert groups and graph harvesting remain rejected. This adds communication and temporary memory proportional to the largest Muon group; it is a correctness change, with no performance claim.

Validation on two NVIDIA L20 GPUs (PyTorch 2.13.0+cu130):

Test	Configuration / coverage	Result
BF16 wrapper regressions	Three pytest cases, each exercised at world sizes 1 and 2; standard/Gram NS, accumulation boundaries, partition-sized FP32 momentum, and the neighboring BF16-gradient configuration	3 passed
Eager engine E2E	DP=2, accumulation over 3 microbatches, clipping=0.05, mixed Muon/Adam groups, and 17×13 / 9×17 matrices crossing partition boundaries; standard and Gram NS	All 3 optimizer steps match the independent full-matrix reference exactly
Checkpoint continuation	Save, continue for 2 steps, reload, and replay; standard and Gram NS	Losses, weights, and optimizer state reproduce exactly
Final source validation	Repeat the eager E2E using the final source files in an isolated directory	Passed, including exact reference comparison and checkpoint replay
Compiled NS	Same-input, per-step FP32-master comparison with upstream compilation enabled	Passed at atol=1e-4, rtol=1e-3; checkpoint replay remains exact
Repository checks	All configured pre-commit hooks on the 4 changed files	Passed

The compiled comparison is a same-input per-step check. Independently evolving BF16 trajectories can diverge after a rounding difference; these results do not establish bitwise compiled equivalence.

Repository regression command:

TORCHDYNAMO_DISABLE=1 PYTHONPATH=.:tests python -m pytest -q \
  tests/unit/runtime/zero/test_muon_without_zero_optimizer.py \
  -k TestMuonBF16Optimizer

This addresses the missing orthogonalization in BF16_Optimizer. It is separate from #8483's ZeRO-1/2 gradient/momentum dtype reconciliation and #7748's checkpoint dtype conversion.

Fixes #8461.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 319f913637

ℹ️ About Codex in GitHub

Codex has been enabled to automatically 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 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Add the required sign-off trailer

This non-merge commit has no Signed-off-by: trailer, so it does not satisfy the repository's mandatory commit-signing requirement and may be rejected by DCO enforcement. Recreate the commit with --signoff using the configured author identity.

AGENTS.md reference: AGENTS.md:L8-L8

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Add the mandatory license header to the new tutorial

This newly added file begins directly with front matter and omits the license header required for all new files in this repository. Add the prescribed SPDX and DeepSpeed Team header without disrupting the documentation front matter.

AGENTS.md reference: AGENTS.md:L12-L16

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Avoid retaining every staged momentum partition

When an optimizer has multiple Muon parameter groups, each group's full local FP32 momentum clone is appended here and retained until after optimizer.step(). Consequently, the staging peak is proportional to the sum of all local Muon partitions—not the largest group as documented—and can add an entire FP32 model shard in common per-layer or per-rate groupings, causing otherwise viable large-model runs to OOM.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Validate updates for matrices split across partitions

At world size two, the two 32×32 matrices deliberately end exactly at partition boundaries, so this test never exercises the newly added left/right slicing for a matrix owned by multiple ranks. Counting private kernel calls and checking that some weight changed would still pass if the wrong cross-rank slice were copied; use irregular matrix sizes and compare complete weights and momentum against an independently derived Muon step.

AGENTS.md reference: AGENTS.md:L30-L33

Useful? React with 👍 / 👎.

[delock]

Hi @0z5a , all_gather_into_tensor fetches the full group momentum, but each rank only needs the slices of matrices it owns part of.

In _prepare_muon_updates(), every rank all-gathers the entire group's momentum, including matrices it owns none of. What each rank actually needs is only the cross-partition slices of matrices its partition intersects. A single dist.all_to_all_single with precomputed send/recv split tables (layout is static, computed once in init) delivers exactly that, still as one fused collective:

comm volume drops from ~full group momentum per rank to just the cross-boundary matrix slices;
the gathered workspace shrinks from group-sized to a few boundary matrices.
This also aligns with the granularity of the ZeRO-1/2 fix in #8141, which only pays extra communication for split matrices. If the all-gather is a deliberate trade (few boundary matrices, simpler code), could you note that in the docstring? As written, "one group's gathered momentum is temporary workspace" undersells the staging cost.

[0z5a]

Thanks for the suggestion. I updated the implementation in 6006dc2 to avoid gathering the full momentum buffer.

The BF16 Muon path now precomputes static send/receive split tables and uses a single dist.all_to_all_single per Muon group, exchanging only the matrix slices that cross DP partition boundaries. Each rank reconstructs only the split matrices it owns before applying the full-matrix Muon update.

I also extended the distributed test with irregular matrix shapes that cross the partition boundary and compare the complete weights and gathered momentum against an independent full-matrix reference. In addition, I ran a real 2-GPU BF16 ZeRO-1 training job for 5 optimizer steps with mixed Muon/Adam parameter groups and three microbatches per step. The weights matched the reference exactly at every step (maximum error 0.0), and the gathered momentum matched exactly as well. The full test file passes with 10 tests.

[delock]

We don't need a seperate tutorial for muon+bf16. Also should not explain internals in tutorial. Can we remove this file?

[0z5a]

removed.

[delock]

@0z5a can you take a look at test failure? Thanks!

[0z5a]

@0z5a can you take a look at test failure? Thanks!

ok, i will see tomorrow.

[0z5a]

@0z5a can you take a look at test failure? Thanks!

Fixed CI error, PTAL @delock .

[0z5a]

@delock , tried to fix again, PTAL.

[delock]

@0z5a let me investigate sandbox in another context, can you revert sandbox timeout change? Thanks!

[0z5a]

Reverted — thanks @delock.