research: reduce cold JAX compile time (post-#71 follow-up)

[Jammy2211]

Overview

Follow-up to #71 (research: settings suffice → cache-by-default shipped in PyAutoConf#128): investigate ways to reduce the cold JAX compile time itself — the once-per-(machine, jax-version, shape) cost the persistent cache cannot remove. Worst measured case: ~7m30 for the GPU pixelized-gradient input_reduce_fusion; ~2 min for CPU MGE gradient fits; plus the uncacheable ~15 s–2 min Python tracing floor every process pays even warm.

Plan

1. HLO artifact (unblocked by PyAutoConf#128's XLA_FLAGS fix): sync RAL mirrors, re-run the 8-min A100 probe with dump flags, analyze what the 7m30 fusion contains; produce the upstream-XLA-report artifact.
2. XLA compile-speed knobs on that fusion (A100, fresh cache per config): compilation parallelism, fusion-heuristic flags; honest autotune A/B (the 2026-07-15 one is unproven — flags were clobbered).
3. Tracing floor (local CPU): profile the uncacheable tracing cost and its scaling with model depth.
4. Cache-entry proliferation: count distinct entries a realistic SLaM pipeline generates; identify avoidable shape/graph variation (canonicalization candidates).
5. Pre-warming: shared cluster cache-dir guidance; jax.export/AOT portability (can a warm cache be shipped?).

Deliverable: cold-compile findings + second verdict in jax_compile/README.md (which levers get productized); follow-ups via intake. Out of scope: restructuring likelihoods into piecewise jit stages (ruled out by #71).

Detailed implementation plan

Affected Repositories

• autolens_profiling (primary — extends the merged jax_compile/; parallel claim with slam-resume-profiling Profile SLaM stage-resume overhead + judge checkpointing vs speed-up #70, user-approved 2026-07-17, different subtrees)

Branch Survey

Repository	Current Branch	Dirty?
./autolens_profiling	main (post-#73 merge)	clean

Suggested branch: feature/cold-compile-reduction

Key Files / Instruments

• jax_compile/probe.py (merged) — extend with tracing-profile and flag-sweep modes as needed
• RAL: pixgrad_logs/compile_ab.sbatch pattern (fresh JAX_COMPILATION_CACHE_DIR per job; matched args or shapes miss)
• Established (do not re-derive): research: JAX compile time — jit boundaries vs settings #71 findings — XLA compiles on host CPUs (load-sensitive); slow-compile banner re-fires during one compile; cache serves AOT path cross-process

Original Prompt

PyAutoMind draft/research/workspaces/investigate_ways_to_reduce_the_cold_jax.md

[Jammy2211]

MAJOR: the 7m30 cold compile IS GPU autotuning. First run ever with --xla_gpu_autotune_level=0 actually applied (PyAutoConf#128 preserved the sbatch flags — every earlier 'autotune off' run was clobbered to defaults, including the 2026-07-15 'ruled out' A/B and #71's controlled A/B, which compared autotune-ON to autotune-ON):

probe (A100, fresh cache)	total	slow-compile alarms
330587 — flags clobbered (autotune ON)	498 s	7m34s fusion compile
330596 — autotune OFF + HLO dump	29 s	none

17× cold-compile reduction from one flag, logL bit-identical (25536.848940 both). The input_reduce_fusion alarm was autotune compiling candidate configs of that fusion as standalone modules — which is also why it never appeared as a dumped module name.

In flight: job 330597 — the full 300×16 adam fit with autotune genuinely off (vs 330537's 2081 s effectively-ON baseline) to price the steady-eval cost of unautotuned kernels. If eval holds, the cold-fit UX answer is a config default (autotune level 0 for fits, or on-by-env), stacking with the warm cache. HLO dump artifact also landed (1443 files, job 330596) for the upstream report if still needed.

Phase 3 (tracing floor) first result: cProfile attribution of a 40 s MGE vag trace — 58 % jax internals, 34 % stdlib/numpy, 7 % autoarray, ~0 % autofit/autogalaxy/autolens — there is no meaningful PyAuto-side Python to cut; the tracing floor is graph size in jax's own machinery.

[Jammy2211]

Steady-eval cost of autotune-off: none observed. Full 300×16 adam fit, fresh cache (330597, autotune OFF) vs 330537 (autotune ON):

	autotune ON (330537)	autotune OFF (330597)
wall	2081 s	1253 s (−40 %)
max logL	−32329	−28434
r_E	1.419	1.434 (truth 1.6)

Wall dropped by MORE than the single alarm's 7m24s — several paths were being autotuned. Eval portion shows no slowdown; fit outcome is parity-or-better (both in the known stochastic-adam regime; the fixed-input probe was bit-identical). Cold-fit picture now: autotune-off 29 s cold probe (17×), −40 % cold full fit, zero measured downside on this likelihood.

Next (decides the productization default): a steady-eval A/B matrix on the A100 — jax_compile/probe.py steady_s for MGE + pixelization × jit/vag × autotune on/off — one likelihood's parity is thin evidence for flipping a stack-wide default in the wrapper. Mirror pull for autolens_profiling (today's merged jax_compile/) queued.

[Jammy2211]

Steady-eval A/B matrix landed (330601/330602) — VERDICT 2: default --xla_gpu_autotune_level=0.

cell (A100, fresh cache)	steady ON	steady OFF	compile ON	compile OFF
mge / jit	0.0042 s	0.0042 s	10.1 s	9.2 s
mge / vag	0.0096 s	0.0098 s	29.4 s	29.3 s
pix / jit	0.0571 s	0.0574 s	8.5 s	5.9 s
pix / vag	0.0871 s	0.0910 s	27.3 s	19.7 s

Eval parity everywhere (worst ~4 % on one cell; the 4800-eval full fit was 40 % faster end-to-end with autotune off, so no real penalty). Combined with the pathological-cell evidence (17× cold probe, bit-identical logL), the recommendation is an env-respecting wrapper default, same pattern as the cache. Cache (#128) + autotune-off = worst-case first-fit UX ~70 min → ~30 s.

Tracing floor closed as a direction (58 % jax-internal / 34 % stdlib / 7 % autoarray — no PyAuto lever); cache-proliferation + pre-warming leads deprioritized (cold is now seconds); upstream XLA report moot (the slow fusion is explained; HLO artifact from 330596 retained). Research note updated on feature/cold-compile-reduction — ready to PR, with the wrapper-default change as the follow-up implementation prompt.

[Jammy2211]

Merged via PR #76 — cold-compile findings + verdict 2 on main. Implementation shipped and merged same day (PyAutoConf#132). Closing: the compile-time arc (#71 → #74) is complete — cache + autotune-off defaults live; no source restructuring anywhere.