[faresobeid]

Summary

Two Muon optimizer-step speedups in trainer/optim. They help SFT, RL, and the planned pretraining data mode. There are no config changes and no loss changes, and the param groups and checkpoint layout are unchanged.

1. Symmetric-GEMM Newton–Schulz (newton_schulz_symmetric). Dion's dense NS is kept, except that X·Xᵀ and b·A + c·A·A run on QuACK's gemm_symmetric. That kernel computes one triangle and mirrors it. The a·X + B·X step runs on addmm/baddbmm with the scalar folded in. QuACK's symmetric GEMM has existed since quack-kernels 0.4.1, which is the current lower bound, so the pin is unchanged. The fork pin S1ro1/quack@f62009f is 0.6.5. It supports SM80/90/100/110/120. The code falls back to Dion's dense zeropower_via_newtonschulz5 in these cases: CPU, no QuACK, pre-SM80, or a matrix dimension not divisible by 8 (QuACK needs 16-byte aligned rows).
2. Expert stacks are orthogonalized where they are sharded. A Muon(dion.Muon) subclass sends a parameter to a local update (muon_update_batch_async with world_size=1) when it is a 3D+ DTensor whose shard placements all split dim 0. In that case each rank already holds whole [fan_out, fan_in] matrices, and NS treats the leading axis as a batch. Everything else still uses Dion's all-to-all path.
   • Dion's path is worse than a plain gather under EP. With EP=8 on one node, expert placements are (_StridedShard(0) on dp_shard_mod_ep=1, Shard(0) on ep=8). Dion's padding then zero-pads each rank's 16-expert shard up to all 128 experts and runs NS on that. Every rank does 8× the needed expert NS work. Dion's logged NS input is (128, 1536, 2048); the local path's is (16, 1536, 2048).

Results (Qwen3-30B-A3B, random init, 1×8 H200, FSDP + ep as listed)

Optimizer-step wall time is measured with the real setup_model + setup_optimizer, random grads, and torch.cuda.synchronize. Each number is the max over ranks, averaged over steady steps 4–8.

ep	main (Dion dense)	+ symmetric NS	+ local experts	both (this PR)
1	0.268 s	0.215 s	0.249 s	0.222 s (1.21×)
2	0.464 s	0.363 s	0.250 s	0.222 s (2.09×)
8	1.634 s	1.226 s	0.250 s	0.221 s (7.4×)

NS microbenchmark on one H200 (ms per call; relative Frobenius error vs an fp32 NS with the same coefficients):

shape	dense	dion triton	QuACK	speedup	err dense / QuACK
4096×4096	3.21	2.35	2.71	1.18×	9.3e-3 / 6.0e-3
8192×8192	25.5	18.8	17.0	1.50×	8.7e-3 / 5.7e-3
4096×14336	8.74	6.71	7.07	1.24×	8.9e-3 / 5.4e-3
128×1536×2048	35.9	25.6	25.0	1.44×	7.5e-3 / 6.2e-3
128×2048×768	12.9	8.8	8.6	1.51×	7.8e-3 / 5.6e-3
32×7168×2048	49.5	35.8	33.0	1.50×	7.7e-3 / 5.4e-3

QuACK is closer to the fp32 reference than the dense bf16 path, because the alpha/beta epilogue runs in fp32. QuACK and the dense path differ from each other by about 1e-2 relative (max abs ≤ 3e-3). That is the size of bf16 NS noise.

First-step update equivalence vs main, as ||Δupdate|| / ||update||, max over ranks, layers 1 and 40, for ep ∈ {1, 2, 8}:

• local-experts path alone: experts 4.0–4.9e-3 (different bmm batch sizes round differently in bf16); every other parameter bit-identical (0).
• symmetric NS: 1.0–1.2e-2 (router 2.3e-2), consistent with the microbenchmark.

Tests

• tests/unit/train/test_muon.py (gpu): symmetric NS is at least as accurate as Dion's dense NS against an fp32 reference, on 2D (wide and tall) and 3D shapes. Passes on H200.
• pytest -m "not gpu" tests/unit/orchestrator tests/unit/test_configs.py tests/unit/train/test_muon.py: 291 passed, 3 deselected. ruff check / ruff format --check are clean.

Notes

• Dion's own Triton NS (use_triton=True, which is also symmetric) is within a few percent of QuACK at most shapes. QuACK is faster on large or batched matrices. I used QuACK because it is already a dependency and it is more accurate.
• Overlap: fix(optim): stream Muon gradients and state for 8-node GLM-5.3 SFT #3719 (draft) also adds optim/muon.py with a Muon(DionMuon) subclass that overrides step. The two combine directly: its streamed path calls _create_muon_tasks, which this PR overrides. Whichever lands second rebases trivially. refactor(trainer): share RL/SFT loop setup, resume step, trace and step logging #3851 touches setup_optimizer's signature but not the lines changed here. There are no conflicts with refactor(configs): one trainer-common config base for TrainerConfig and SFTConfig #3849 / Chore/remove transformers modeling #3833 / refactor(lora): single-adapter LoRA layers, drop modules_to_save #3855 / refactor(trainer): add iter_moe_blocks for the MoE layer lookup #3843.