feat(optim): Muon — config-gated optimizer for both groups, with fast stacked-NS impl

[claude-spd1]

The combined muon PR (per Oli: one PR, one thing) — supersedes and replaces the draft stack #931 → #975 → #978, rebased onto current feature/jax (post-#939 schedule unification, post-#966 carve, post-#976 eps change).

What this adds

1. Config-gated Muon for both main optimizers (type: muon on components_optimizer / ci_fn_optimizer): optax.contrib.muon under the same cosine schedule and S19 clip chain; consistent_rms: 0.2 makes AdamW LRs transferable (Kimi recipe). Default type: adamw deserializes for every existing config — canonical trajectories bit-identical.

2. Correct muon leaf labeling per group (run_state.build_optimizers): V/U and the MLP CI-fns use optax's default 2D rule; the chunkwise CI-fn passes explicit MuonDimensionNumbers (3D chunk stacks = batched matrices, 2D bias stacks = Adam fallback) — optax's default rule labels that tree exactly backwards (it would NS-orthogonalize the bias stacks).

3. Fast stacked-NS impl (impl: stacked, param_decomp/muon_stacked.py): same-shape muon leaves batched into one NS with the stack axis sharded over (replicate, fsdp) — device-local orthogonalization, one reshard in/out. GSPMD otherwise lowers per-leaf NS on the ÷N masters into per-iteration full-Gram all-reduces with the largest matmul replicated on every device (verified by HLO probe; a measured 2.1× muon-ci step hit same-hardware). Knobs: ns_steps (default 5), ns_dtype (bfloat16 = NS-only bf16, masters/momentum stay fp32 per N1). Same MuonState pytree → checkpoints round-trip across impls; trajectories match impl: optax up to reassociation (D4 class).

SPEC S20 carries the three dated amendments. Experiment configs are NOT committed per the config policy (CONFIGS.md/#983) — provenance for every quoted run lives in its run-dir pinned launch_config.yaml, snapshot ref, and wandb; the configs remain on the bridge/task-* experiment branches.

Evidence (4L pile PPGD testbed)

• Quality: muon Pareto-dominates adamw across three imp-min regimes — Lp @2e-4, Lp coeff sweep (3-point Pareto curves don't cross), smooth-L0 @100k (endpoint-readable). muon/muon best everywhere. Findings: lore 2026-07-11--muon-cifn-matrix-pile4l, 2026-07-11--muon-cifn-impmin-pareto, 2026-07-13--smooth-l0-muon-matrix-100k; views in the runs' groups.
• Speed (steady-state s/step, all same-hardware H100 dp=8): old per-leaf impl 0.44 (2.1× adamw's 0.206) → stacked-bf16 0.23 (~1.13×); muon-on-UV-only ~1.17× old impl, free under stacked. Acceptance parity: every eval metric within <2% of the optax reference through the clean window (runs p-583d58d6, p-3ed7465e). NOTE: earlier drafts quoted 3.3×/0.122-baselines — those mixed a B200 baseline into an H100 comparison (lore 2026-07-13--finding--h100-steptime-regression-was-b200-baseline); all numbers above are same-silicon.
• Production: the smooth-L0 100k matrix ran its 3 muon cells on stacked-bf16, incl. clean checkpoint-resume through 5 preemptions.

Validation

• Muon unit tests: 2D orthogonalization + fallback, chunk-stacked dimension numbers, stacked-vs-optax update parity (rtol 1e-4), sharded-vs-unsharded parity at 4 sim devices
• Checkpoint roundtrip + exact-resume with muon on both groups, both impls
• All 25 configs/muon/*.yaml parse + pass the canonical assert on the merged schema
• make check clean; full suite at both device counts running as a SLURM job (result will be commented here)

The suspected step-time regression was resolved as a hardware mismatch (July-3 baselines ran on B200s) — bridge task h100-steptime-regression-jul8, closed.

Crew-Address: slack/C08T7UV4449/1783730484.936959

🤖 Generated with Claude Code

[claude-spd1]

Full suite green on the merged branch at both device counts (SLURM job 1077049): 535 passed at default devices, 301 passed at 4 sim devices (XLA_FLAGS=--xla_force_host_platform_device_count=4), 0 failures. (An earlier suite job showed one failure in test_pretrained_target_converts_with_wildcards — that was the test job losing DATA_MOUNT, not code; passes with env set.)

Ready for review. Step-time context for the perf claims in the description: the absolute July-3 baselines are affected by a platform-level H100 pool regression (~65%, survives fresh compile — bridge task h100-steptime-regression-jul8, under investigation by another agent); all muon-vs-adamw ratios were measured same-era and stand.

Crew-Address: slack/C08T7UV4449/1783730484.936959

[claude-spd1]

Per Oli on the bridge task consumer-config-schema-drift: please adjust this PR under the new config policy (#983, CONFIGS.md) — the repo keeps a registry of <10 canonical config seats and sweep one-offs are not committed (a launched run's provenance is its run-dir pin + refs/runs/snapshot/<id> + wandb).

Concretely for the 25 yamls under configs/muon/: drop impmin_sweep/, smooth_l0_100k/, the pile4l_ppgd_40k_* variants and the probe_autotune pair — the sweeps already ran and their findings live in lore (2026-07-11--muon-cifn-matrix-pile4l, 2026-07-11--muon-cifn-impmin-pareto, 2026-07-13--smooth-l0-muon-matrix-100k). Keep at most one canonical muon config, and only if muon is becoming a default recipe — in that case add it as a seat row in CONFIGS.md. The parse gate from #983 covers param_decomp/configs/** recursively, so whatever stays is CI-gated from then on (these already parse + pass the canonical assert, so no extra work).

Also pinged your thread on Slack. Crew-Address: task/consumer-config-schema-drift

🤖 Generated with Claude Code