perf(compile): 13x faster V/U init compile + compile-time findings

[ocg-goodfire]

Description

Compile-time pass for the full-model runs (bridge task reduce-jax-compilation-time). Every new-config launch of the 32L Llama-8B decomposition paid a 3m31s jit_init_decomp_vu compile plus ~24 min of jit_step compile; this PR kills the first, gates incidental compile costs, and pins down (with measurements) where the second does and does not go.

Changes:

• Stacked V/U init — 210.5s → 16.3s (12.9×) measured on GPU at the full 224-site/32L shape set. init_decomp_vu_placed now compiles two cheap stages: the RNG vmap-STACKED per V/U shape (2×n_shapes sharded jit outputs instead of 2×n_sites — the SPMD pass over a 448-output RNG graph was the multi-minute cost), then a trivial slice jit (stacked input donated) fans out to the per-site ZeRO-1 layout. Bit-identical to the previous jitted per-site init (vmap over the same per-site keys); test_sharding now pins exact equality against the jitted per-site reference. (Found in passing: the unjitted eager init_decomp_vu differs from ANY jitted init by ~1 ULP of XLA fusion — pre-existing, previously hidden by the test's allclose.)
• HLO dump is now opt-in (ProfileConfig.hlo_dump, default off) — it cost compile time + ~100s of MB on every rank-0 run.
• sbatch submissions scrub SLURM_*/SBATCH_* from the env — submitting from inside a SLURM allocation (agent sessions, dev pods) leaked SLURM_CPUS_PER_TASK into the new job, whose srun died instantly with "cpus-per-task set by two different environment variables" (the leak the btdr wrapper works around by hand).
• Compile-probe generator regenerated against the current schema (it had bit-rotted: pre-discriminator ci_config, removed beta) and now derives from the production 32L yaml with per-arm cache isolation.
• Docs: measured findings + smoke-config cache practice in param_decomp/CLAUDE.md.

Measured nulls / negative results (probe arms A–F, 8L 1-node slices of the production config, per-arm isolated caches):

• xla_gpu_autotune_level 0/2/4: jit_step compile flat (71.3/71.8/74.4s vs baseline 71.3s) — with triton_gemm off, autotuning is a negligible slice.
• Persistent autotune cache: already on by default — jax ≥0.10 auto-enables the XLA per-fusion autotune cache alongside the persistent compilation cache (arm A, baseline code, had 306 entries). A hand-wired duplicate was added and then reverted within this branch; the built-in also gates writes on rank 0 correctly.
• jax_persistent_cache_enable_xla_caches='all' (kernel cache + parallel LLVM codegen): crashes at runtime on jax 0.10.1/B200 (CUDA_ERROR_NOT_FOUND: named symbol not found in jit_step). Documented as do-not-use.

Not addressed (the remaining ~24 min jit_step compile): it is graph size (the recon plan unrolls n_chunks full-model masked forwards+backwards python-side) × topology, not autotuning. The structural lever is a lax.scan over uniform recon chunks — scoped as a follow-up board task, core-trainer change.

Related Issue

Bridge task reduce-jax-compilation-time (see its log for the full probe data).

Motivation and Context

20–30 min before step 0 on every new full-model config made iteration painful. The persistent executable cache only hits on identical HLO, so every real config change paid full price.

How Has This Been Tested?

• make test (458 passed) + make type clean, test_sharding at 1 and 4 sim devices.
• Bit-identity of the new init vs the old jitted init pinned by test (exact array_equal, sharded + unsharded).
• GPU A/B at full production site count (224 sites, 1 node, tp8): old 210.5s vs new 16.3s.
• Probe arms A–D ran the production compile path end-to-end (10 steps, checkpoints, identical losses/step times across arms: 0.95–0.99 s/step, 119GB peak/rank).

Does this PR introduce a breaking change?

No. Init values are bit-identical; HLO dump moves behind runtime.launch_env.profile.hlo_dump: true.

🤖 Generated with Claude Code

https://claude.ai/code/session_01CLGxUn9wtE3ce9XTbEJo9s