perf(compile): kill the init-compile share of the cold start (ci-fn 167s→4.4s, sources 55s→<2s) + compile-time attribution

[danbraunai-goodfire]

Description

Compile-time investigation follow-through (continues #941/#951). A controlled probe grid on and-btdr — one production-shape arm per mutation (full 32L / 224 sites / full C / dp32), each with an isolated XLA cache — attributed the cold-start wall to individual compiles. The headline: the historical "~20–24 min before step 0" was mostly the serial init compiles, not jit_step, and the graph-structure refactors previously scoped for jit_step would buy ~nothing.

Fixes (both the init_decomp_vu recipe from #951, bit-identity pinned by tests):

• init_chunkwise_transformer_ci_fn unrolled n_chunks python-side chunk inits and stacked them (32 copies of the chunk RNG body in one graph). Now eqx.filter_vmap over the same fold_in keys. jit_build_ci_fn: 167s → 4.4s measured at the production shape.
• init_sources_sharded emitted n_sites (224) sharded outputs. Now the two-stage form: vmap-stacked per C group, then a donated slice jit fans out per site. jit_init_persistent_sources: 55s → under the 2s log threshold.
• param_decomp/log.py: dictConfig was disabling jax's already-created loggers (disable_existing_loggers defaults True) — JAX_LOG_COMPILES=1 and persistent-cache diagnostics never reached slurm logs. One line.

Measured attribution (probe grid, docs added to param_decomp/CLAUDE.md):

• jit_step compile is flat across every graph-structure knob: recon chunks 1/2/4/8 = 383/262/306/365s (noise band), halved C ≈ 289s, 2-block CI fn ≈ 289s, PPGD warmup on/off ≈ −45s, PPGD removed ≈ −70s. Topology slope mild: dp32 ≈ 290–306s → dp64 ≈ 345s. Pass profile (TF_CPP hlo_pass_pipeline arm): ~83% priority-fusion. The scan-over-recon-chunks follow-up scoped in perf(compile): 13x faster V/U init compile + compile-time findings #941 is not worth doing for compile time.
• Cold start at the production dp32 shape: was ≈ model load + 211s (V/U) + 167s (ci fn) + 55s (sources) + ~300s (step); now ≈ model load + ~22s of inits + ~260–290s step compile — ~5–6 min, down from ~12–13.

No measurement tooling ships in this PR — the probe configs were one-shot artifacts (cluster-specific paths), so the method is documented in param_decomp/CLAUDE.md and the raw logs + launch configs are archived on btdr (compile_probe_grid_2026-07-06_logs.tgz). Every file in the diff is cluster-independent library code.

Stacked on #951 (contains its commit b676138cc; also duplicates the sbatch env-scrub commit from #941). Merge #951 first and this rebases to just the new work.

Related Issue

Continues the bridge task reduce-jax-compilation-time (#941, #951).

Motivation and Context

Every new-config launch paid ~7 min of pure init compiles before the step compile even started; sweeps paid it per arm. The init fixes remove that entirely, bit-identically. The attribution data stops us from spending a core-trainer refactor (chunk scan) on a cost that measurement shows it cannot move.

How Has This Been Tested?

• Bit-identity of both new inits pinned: vmapped CI init vs unrolled+stacked reference (exact array_equal), stacked sources vs per-site init_persistent_sources (exact, incl. bf16 storage dtype) — in test_sharding at 1 and 4 sim devices.
• make test (same 4 pre-existing test_slow_eval failures on clean feature/jax), tests/equivalence goldens pass, make check/make type clean.
• GPU A/B at the full production shape (btdr, dp32, isolated caches): fastinit arm vs base4 arm — ci-fn 167→4.4s, sources 55s→<2s, jit_step unchanged (~260–275s vs ~290–306s, within the arm noise band).

Does this PR introduce a breaking change?

No. Init values are bit-identical; the log change only re-enables foreign loggers.

🤖 Generated with Claude Code

https://claude.ai/code/session_01E34cK5BaVr6ZFAwRidNp6n

[danbraunai-goodfire]

I'm not sure about this, leaving it on the shelf.

1. Seems like a lot of code.
2. It does this unrelated hotifx to scrub out the env of the submitter. But I think we want to preserve the submitter's env.

[danbraunai-goodfire]

It did give some reasonable compilation speedups. Note that the increase in step time isn't due to code changes, it's just measurement noise. We didn't make code changes that can increase step time.

[danbraunai-goodfire]

Superseded by #956: this PR's content minus the already-merged #955 commit, combined with #951 into one clean 3-commit branch.