[faresobeid]

Replaces the dion dependency (samsja/dion@0c801b0) with a Muon optimizer in prime-rl: src/prime_rl/trainer/optim/muon.py.

Changes vs main

• Muon (prime-rl) replaces dion.Muon. Its defaults give the same update as main:
  • momentum M = mu*M + G, with Nesterov (mu*M + G) when optim.nesterov is set;
  • bf16 Newton-Schulz with dion's 5-step coefficients and eps 1e-8. Each iteration's scale-and-add is folded into the GEMMs (baddbmm);
  • update scaled by 0.2*sqrt(max(fan_out, fan_in)) per matrix (dion's rms_norm), with decoupled weight decay;
  • each matrix_partitions entry (fused qkv / gate_up) is orthogonalized and scaled on its own. Equal partitions run as one batched Newton-Schulz;
  • AdamW for the non-matrix params, with the same routing as main.
• Distribution comes from each parameter's DTensor placement:
  • A parameter sharded on a matrix dim (FSDP Shard(0), or Shard(1)) is batched with same-layout parameters, in batches of its shard-group size. An all-to-all over that group gives each rank one whole matrix. The rank orthogonalizes it, and a second all-to-all sends the result back. Shards may be uneven or empty.
    • Short batches send only to the ranks that own a matrix. There are no dummy matrices. Owners continue round-robin from one batch to the next, so short batches land on different ranks.
  • Parameters sharded only on batch dims, such as experts under EP, are orthogonalized locally. There is no communication and no padded expert buffers, as in dion@0c801b0.
  • Up to 3 batches are in flight, so all-to-alls overlap Newton-Schulz.
• optimizer.chunks() / optimizer.step_chunk(chunk): step() walks a fixed list of chunks. A chunk is one batch of Muon matrices or one AdamW group. step_chunk updates only that chunk's params, on the current stream.
• optimizer.orthogonalize_chunk(params, U, first_owner): Newton-Schulz plus the exchange, without the momentum update or the parameter update.
• Config: optim.newton_schulz_coefficients, a list of (a, b, c) triples, one per Newton-Schulz iteration. The default is dion's 5-step schedule, written out. optim.update_rms defaults to 0.2 (Kimi/Moonlight). DeepSeek-V4's setting is newton_schulz_coefficients = 8 x (3.4445, -4.7750, 2.0315) + 2 x (2, -1.5, 0.5) with update_rms = 0.18: about 2x the Newton-Schulz time, singular values 0.998-1.003.
• The NCCL warm-up now runs on every group that Muon's all-to-alls use. These groups are read from the parameters, not from dp_shard_cp / dp_shard_mod_ep.
• dion is removed from pyproject.toml and uv.lock.
• Checkpoints: state keys (momentum, variance) and param-group keys and layout are the same as dion's. A dion checkpoint therefore resumes directly (checked below).

Validation

Optimizer-step benchmark: the real setup_model + setup_optimizer, random init, random grads, Nesterov, wd 0.1. Times are the median over steps 2+, max over ranks. H200.

run	impl	step	peak GiB	step transient GiB	all-to-all GB/rank
Qwen3-30B-A3B, 1 node EP8, offload none	dion	0.267 s	43.8	0.87	0.42
	this PR, compiled NS	0.261 s	43.5	0.50	0.42
	this PR	0.215 s	43.3	0.35	0.42
Qwen3-30B-A3B, offload optimizer	dion	1.404 s	43.8	15.37	0.42
	this PR, compiled NS	1.369 s	43.5	15.01	0.42
	this PR	1.479 s (compiled NS, separate job: 1.692 s)	43.3	14.85	0.42
GLM-4.5-Air, 2 nodes cp4 EP8, offload none	dion	0.639 s	76.9	1.96	1.75
	this PR, compiled NS	0.613 s	76.2	1.28	1.56
	this PR	0.534 s	76.2	1.22	1.40
GLM-4.5-Air, offload optimizer	dion	2.569 s	76.9	27.12	1.75
	this PR, compiled NS	2.544 s	76.2	26.45	1.56
GLM-5.3, 6 layers, 2 nodes EP8, offload none	this PR, compiled NS	0.260 s	27.5	4.23	0.82
	this PR, compiled NS + owner rotation	0.240 s	27.5	4.20	0.48
	this PR	0.215 s	26.3	2.95	0.48

Newton-Schulz per call (H200), compiled -> baddbmm: (32,768,2048) 2.12 -> 2.10 ms, (16,2048,768) 1.67 -> 1.21 ms, (16,1408,4096) 6.25 -> 5.80 ms, (8,2048,6144) 9.54 -> 8.57 ms, (6144,16384) 24.1 -> 23.1 ms. Relative error against fp64 Newton-Schulz with the same coefficients is 0.54-0.58%, down from 0.77-0.98% with the compiled version, which rounded A @ A and B @ X to bf16 before the adds. Dropping the compile also shortens the first step (Qwen3-30B-A3B 8.0 s vs 22.7 s, GLM-4.5-Air 15.0 s vs 44.3 s).

Against the compiled-NS version of this PR (same weights and grads, 3 steps, Qwen3-30B-A3B EP8, 4 layers), ||Δ - Δ_compiled|| / ||Δ_compiled|| is 0.36 / 0.47 / 0.54%, below the 0.9% bf16 reference below. Owner rotation alone gives bitwise-identical updates for every matrix that goes through the all-to-all (Qwen3-30B-A3B, GLM-4.5-Air, GLM-5.3).

Update A/B: same weights and grads in one process, 4 layers, 5 steps. The metric is ||Δ_this - Δ_dion|| / ||Δ_dion|| over all Muon matrices. As a bf16-noise reference, dion is run against itself with grads scaled by 1.001. Muon's output does not depend on that scale except through rounding.

model	this PR vs dion (steps 0-4)	dion vs dion x1.001	AdamW params
Qwen3-30B-A3B EP8	0.92% (flat)	0.91%	1e-6
GLM-4.5-Air cp4 EP8, 2 nodes	0.89% (flat)	0.90%	1e-6

The max per parameter is 0.94% vs 1.0%. With offload = "optimizer" the results are the same.

SFT, Qwen3-0.6B Reverse-Text-SFT, Muon lr 5e-5, bs 32, 150 steps. Mean loss over each step range:

run	11-31	32-62	51-100	101-150
dion seed 0 / rerun	2.2213 / 2.2213	0.3069 / 0.3052	0.0089 / 0.0089	0.0015 / 0.0014
this PR seed 0 / rerun	2.2231 / 2.2232	0.3072 / 0.3070	0.0093 / 0.0092	0.0013 / 0.0014
dion / this PR, seed 1	2.2310 / 2.2290	0.3113 / 0.3127	0.0093 / 0.0096	0.0015 / 0.0014
this PR resumed from a dion step-50 checkpoint	-	-	0.0088	0.0015
this PR, DeepSeek-V4 coefficients (update_rms 0.2)	2.0576	0.2192	0.0052	0.0013

After the resume, the step-51 loss equals dion's own continuation (0.0737). The SFT runs came before the commit that compiles the momentum and update ops. That commit is numerically equivalent: unpartitioned matrices are bitwise equal on CPU, and the update A/B above was rerun after it.

Tests: tests/unit/train/test_muon.py. It checks that Newton-Schulz singular values stay in range for the default and the DeepSeek-V4 coefficients, and that a packed parameter steps like its separate matrices. pytest -m "not gpu" tests/unit/test_configs.py tests/unit/train tests/unit/orchestrator: 556 passed. The default coefficients give bit-identical params and optimizer state to the previous preset (CPU, 3 steps). ruff is clean.

Related open PRs

• perf(optim): overlap optimizer state offload copies with the AdamW step #4008 / perf(optim): stream Muon optimizer state per batch in CPU offload #4116 (chunked state offload): perf(optim): stream Muon optimizer state per batch in CPU offload #4116's Muon units become optimizer.chunks(), stepped with optimizer.step_chunk. The dion imports and the per-chunk group-step restore can go, because AdamW's step count is only bumped by its own chunk.
• feat(optim): optional bfloat16 momentum for Muon #4117 (bf16 momentum): becomes the config field plus one dtype= in Muon._state.
• feat(optim): heads_per_block for MLA up-projections (GLM-5 Muon Split) #4119 / feat(optim): per-head Muon for attention q/k/v (heads_per_block, default off) #4114 / fix(optim): Muon routing for depthwise conv1d and grouped linear weights #4120 (head blocks, MLA split, conv1d routing, GroupedLinear): unchanged. They only produce matrix_partitions and routing, which Muon takes the same way as dion did. Equal head blocks now run as one batched Newton-Schulz instead of a Python loop per partition.
• feat(optim): support Muon with full CPU optimizer offload #4124 (Muon with offload=full): the zero-grad / unit-momentum call into dion.muon_update_batch_async becomes optimizer.orthogonalize_chunk(params, bf16 momentum) plus optimizer._update_scale.
• feat(trainer): Nesterov momentum for Muon by default #4113 (Nesterov): already on main and kept.

---

Note

Medium Risk
Changes the distributed Muon optimizer step and NCCL all-to-all behavior on training-critical paths; checkpoint layout is preserved but numerical updates may differ slightly from dion (~1% in reported A/B tests).

Overview
Replaces the external dion Muon optimizer with an in-tree implementation in src/prime_rl/trainer/optim/muon.py and removes the dion git dependency from pyproject.toml / uv.lock.

The new Muon keeps dion-compatible param groups and optimizer state keys so existing checkpoints resume. It runs bf16 Newton–Schulz (configurable via optim.newton_schulz_coefficients and optim.update_rms), AdamW on non-matrix params, and matrix_partitions for fused weights. FSDP2 matrix shards are reassembled with paired all-to-alls derived from each parameter’s DTensor placement (shard_group), with up to three chunks in flight and NCCL warm-up on every shard process group.

setup_optimizer wires the new class, drops explicit mesh wiring for Muon, and exposes chunks(), step_chunk(), and orthogonalize_chunk() for chunked/offload paths. Unit tests cover Newton–Schulz singular values and packed-parameter equivalence.

^{Reviewed by Cursor Bugbot for commit a7c914e. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit c6717a3. Configure here.}