feat(optim): Muon for the chunkwise ci-fn optimizer (completes the {adamw,muon} x {UV,ci-fn} matrix)

[claude-spd1]

Stacked on #931 (draft) — diff here is only the ci-fn delta. Requested by Oli in the muon pilot slack thread; bridge task muon-cifn-matrix.

What

type: muon on ci_fn_optimizer now works correctly for the chunkwise CI fn. Its trainable tree is per-chunk stacks — 3D [n_chunks, d_in, d_out] matrix stacks and 2D [n_chunks, d] bias stacks — and optax's default muon rule (2D → muon, rest → Adam) labels that tree exactly backwards (it would NS-orthogonalize the bias stacks and Adam the matrices). build_optimizers now passes explicit MuonDimensionNumbers for the ci-fn group when the arch is chunkwise: 3D leaves are orthogonalized over the trailing two axes (chunk axis batched), everything else takes the Adam fallback. V/U and the MLP CI fns keep the default 2D rule (correct for their trees). Default adamw stays bit-identical canonical. SPEC S20 amendment extended (2026-07-11) — needs Oli sign-off.

Also: the two matrix arm configs (byte-derived from the try-muon control / muon-1x stored launch configs, only ci_fn_optimizer flipped), plus JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES=none in their launch_env (the shared autotune-cache footgun, see task xla-autotune-cache-group-unwritable / #974; also matches the 07-03 arms' autotune behavior).

Validation

• New unit test: 3D-stack per-chunk NS orthogonalization + 2D-bias-stack Adam fallback (test_optim_torch_parity)
• New checkpoint test: muon on BOTH groups roundtrips + exact-resumes (the scavenge-preemption path)
• Full suite green at default devices (464 passed) AND 4 sim devices (236 passed); make check clean

Experiment (full 2x2 at lr 5e-5, pile4l ppgd bsc 40k, seed 0, wandb group muon-cifn-matrix)

UV \ ci-fn	adamw	muon
adamw	p-d0f66d3b (try-muon control)	p-282f5206
muon	p-cedbbf7c (try-muon 1x)	p-c01f5833

View: https://wandb.ai/goodfire/param-decomp?nw=r3mmipzgpne

Mid-training window (10k–35k, per the try-muon methodology; endpoint is schedule-compression-artifacted):

• muon/muon is Pareto-best: kl_stoch 0.31–0.50 vs adamw/adamw 0.60–0.84, at LOWER CI-L0 from 20k on (332 vs 447 @30k); best kl_ci, best hidden-side train PPGD recon (0.83 vs 2.43 @30k, ~3x), best eval 20-step PGD, slightly best unmasked KL.
• ci-fn muon helps on BOTH UV rows (adamw/muon beats adamw/adamw everywhere mid-window; muon/muon beats muon/adamw everywhere) — the effect is additive, not redundant with UV muon.
• Cost is the catch: ~3.3x step time at this scale (0.40–0.44 s/step vs 0.12–0.15). The per-chunk NS orthogonalization (vmapped over 8 chunk stacks x per-block matrices) dwarfs the +19% UV-muon overhead. Relative cost should shrink at production scale, but at 4L the quality win is not free.
• Single seed, 4L, compressed 40k schedule — standard caveats.

Crew-Address: slack/C08T7UV4449/1783730484.936959

🤖 Generated with Claude Code

[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