feat(optim): stacked-NS Muon — batched, stack-axis-sharded orthogonalization (fast muon)

[claude-spd1]

Stacked on #975 (which stacks on #931) — diff here is the fast-muon delta only. Bridge task muon-sharding-design; design + HLO-probe analysis on its log, Oli-greenlit direction ("work towards fast muon", 2026-07-12 slack).

Why

GSPMD lowers the per-leaf 5-step Newton-Schulz on our ÷N-sharded fp32 masters into two pathological patterns (verified by compiling the real update at real shapes/shardings and reading the partitioned HLO): per-NS-iteration full-Gram all-reduces with the iteration's LARGEST matmul (A@A) replicated on every device, and gather-everything layouts moving up to 1.5 GB per leaf per step. Net: ~24.6 GB of serialized fp32 collectives + 2.1x redundant compute per step for the 4L ci-fn group — the measured 3.3x muon-ci step-time hit (+19% for V/U). Ideal-sharded NS would cost +21% / +2% in FLOPs.

What

param_decomp/muon_stacked.py: same-shape muon leaves are batched into ONE NS with the stack axis sharded over (replicate, fsdp) — each 2D orthogonalization runs device-local, one reshard in/out, zero per-iteration collectives (Kimi's parameter-partitioned recipe, expressed GSPMD-natively). Config-gated on MuonOptimizerConfig:

• impl: optax | stacked — default optax keeps the 07-02/07-11 experiment arms' exact semantics; stacked matches up to float reassociation (SPEC D4 tolerance class)
• ns_steps (default 5, unchanged)
• ns_dtype: float32 | bfloat16 — bf16 NS halves compute+comm (Kimi/modded-nanogpt practice); masters/momentum stay fp32 per N1; stacked-only

Same MuonState pytree → checkpoints round-trip across impls. build_optimizers takes the mesh (None for toys/CPU). SPEC S20 amended 2026-07-12 — needs Oli sign-off.

Validation

• stacked-vs-optax update parity (2 steps, mixed 2D/3D/Adam-fallback tree, rtol 1e-4)
• sharded-vs-unsharded parity at 4 sim devices (the constraint is layout-only)
• checkpoint roundtrip + exact-resume with stacked muon on both groups
• make check clean; full suite at both device counts: SLURM job (result to be posted on the task)

Acceptance (in flight)

Two 40k both-muon relaunches, byte-identical to p-c01f5833 except the impl gate (group fast-muon-accept): impl: stacked and impl: stacked + ns_dtype: bfloat16. Expected step ~0.15-0.16s vs 0.44 today (baseline adamw/adamw 0.122); quality curves should track p-c01f5833 within reassociation noise. Results will be posted on the bridge task + PR.

Crew-Address: slack/C08T7UV4449/1783730484.936959

🤖 Generated with Claude Code

[claude-spd1]

Full suite green at both device counts (SLURM job 1072837): 468 passed default devices, 240 passed at 4 sim devices (XLA_FLAGS=--xla_force_host_platform_device_count=4), 0 failures. Acceptance arms (stacked-fp32 p-583d58d6, stacked-bf16 p-4a37302c) are running; timed step-time read to follow.

Crew-Address: slack/C08T7UV4449/1783730484.936959

[claude-spd1]

Acceptance verdict: PASS.

Timing (4L, steady-state s/step): adamw 0.122 · old optax muon-ci impl 0.44 · stacked-fp32 0.263 · stacked-bf16 0.232 — both-muon overhead 3.6x → 1.9x.

Quality parity vs the optax reference p-c01f5833 (both arms byte-identical configs except the impl gate): every eval metric (stoch/CI/unmasked KL, CI-L0, 20-step PGD) tracks within <2% at 10k/20k/30k/35k — well inside reassociation noise. Endpoint (40k) diverges as all runs do in the schedule-squeeze zone; that window is excluded from reads per the established methodology. Runs: stacked-fp32 p-583d58d6, stacked-bf16 p-3ed7465e, group fast-muon-accept.

Remaining ~0.11s over adamw at 4L (vs ~0.03 ideal): parked as a profile-first follow-up (candidates: reshard layout-transforms, batched-gemm utilization, concat/unstack HLO). The stacked-bf16 mode is now in live use by the smooth-L0 100k matrix relaunch.

Crew-Address: slack/C08T7UV4449/1783730484.936959

[claude-spd1]

Superseded by the combined muon PR #982 (per Oli: one PR for the whole muon feature, rebased onto current feature/jax). All content + validation carried over there.

Crew-Address: slack/C08T7UV4449/1783730484.936959