diff --git a/jax_compile/README.md b/jax_compile/README.md index 5adf3c6..25fc321 100644 --- a/jax_compile/README.md +++ b/jax_compile/README.md @@ -145,5 +145,66 @@ one-time-per-machine cost the cache already removes. **Do not restructure.** reuse `probe.py` and should track *warm* compile times per cell so cache regressions are caught. +## Cold-compile findings (issue #74, 2026-07-17) + +The persistent cache (above) solves repeat fits; these findings address the +**cold** cost it cannot remove. + +### 6. The pathological cold compile IS GPU autotuning — one flag removes it + +The first run in this stack's history where `--xla_gpu_autotune_level=0` +actually reached XLA (the pre-#128 wrapper clobbered every earlier attempt — +including the 2026-07-15 "ruled out" A/B and finding 3's controlled A/B, which +compared autotune-ON to autotune-ON): + +| A100, fresh cache | autotune ON | autotune OFF | +|---|---|---| +| FD probe total (kernel-CDF pixelized vag) | 498 s | **29 s** (17×) | +| full 300×16 adam fit wall | 2081 s | **1253 s** (−40 %) | +| fixed-input logL | 25536.848940 | 25536.848940 (bit-identical) | + +The `input_reduce_fusion` slow-compile alarm was autotune compiling candidate +kernel configs of that fusion as standalone modules (why it never appeared as +a dumpable module name). Autotune results cache like everything else, so with +the cache enabled this is first-fit UX; without the flag a new machine's first +pathological fit pays ~7m30. + +### 7. Steady-state eval does not need autotune (measured cells) + +A100 probe matrix, fresh caches per job (jobs 330601/330602): + +| cell | 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 | + +Worst case ~4 % on one cell — and the 4800-eval full fit ran *faster* end to +end with autotune off, so no real eval penalty is observed. (These standard +cells compile in seconds either way; the 7m30 autotune cost is specific to the +pathological kernel-CDF no-sparse-operator shape.) + +### 8. The tracing floor is jax-internal — no PyAuto lever + +cProfile attribution of a 40 s MGE `vag` trace (`trace_profile.py`): 58 % jax +internals, 34 % stdlib/numpy, 7 % autoarray, ~0 % autofit/autogalaxy/autolens. +Reducing it means emitting fewer ops (a jax-side concern), not optimizing +PyAuto Python. Documented and closed as a direction. + +## Verdict 2 (cold compile) + +**One more setting: default `--xla_gpu_autotune_level=0`.** Evidence: 17× +pathological cold-probe reduction, −40 % cold full fit, bit-identical fixed- +input likelihoods, steady-state eval parity across the measured matrix. With +the cache (#128) plus autotune-off, worst-case first-fit UX drops from ~70 min +to ~30 s. Recommended as an env-respecting wrapper default (same pattern as +the cache), so clusters can re-enable autotune where a tuned kernel matters. + +Remaining leads deprioritized by these numbers: cache-entry proliferation and +pre-warming (cold is now ~seconds-to-a-minute); upstream XLA report (the "slow +fusion" is explained — autotune candidates on a 58 GiB fusion; the HLO dump +artifact from job 330596 exists if ever needed). + Not indicated: source restructuring, jit boundaries inside likelihoods, replacing `lax.map` in MultiStartAdam, autotune flags. diff --git a/jax_compile/results/local_cpu/trace_profile.json b/jax_compile/results/local_cpu/trace_profile.json new file mode 100644 index 0000000..d2533e6 --- /dev/null +++ b/jax_compile/results/local_cpu/trace_profile.json @@ -0,0 +1,20 @@ +[ + { + "model_type": "mge", + "transform": "vag", + "ndim": 15, + "trace_s": 39.822, + "rollup_tottime_s": { + "jax": 23.046, + "other": 13.66, + "autoarray": 2.803, + "autogalaxy": 0.123, + "autofit": 0.049, + "autolens": 0.01, + "autoconf": 0.006 + }, + "jax_version": "0.10.2", + "tag": "first", + "timestamp": "2026-07-17T08:26:15" + } +] diff --git a/jax_compile/results/local_gpu_NVIDIA_A100_80GB_PCIe/mge.json b/jax_compile/results/local_gpu_NVIDIA_A100_80GB_PCIe/mge.json new file mode 100644 index 0000000..589d78a --- /dev/null +++ b/jax_compile/results/local_gpu_NVIDIA_A100_80GB_PCIe/mge.json @@ -0,0 +1,24 @@ +[ + { + "transform": "jit", + "trace_s": 4.885, + "compile_s": 9.204, + "first_s": 1.329, + "steady_s": 0.0042, + "dataset_class": "imaging", + "model_type": "mge", + "instrument": "hst", + "ndim": 15, + "n_batch": 16, + "batch_size": 4, + "hardware": "local_gpu_NVIDIA_A100_80GB_PCIe", + "jax_version": "0.10.2", + "cache_dir": "", + "mixed_precision": false, + "tag": "a100-at0", + "hostname": "euclid-ral-gpu-2", + "timestamp": "2026-07-17T08:55:59" + } +] + +] diff --git a/jax_compile/results/local_gpu_NVIDIA_A100_80GB_PCIe/pixelization.json b/jax_compile/results/local_gpu_NVIDIA_A100_80GB_PCIe/pixelization.json new file mode 100644 index 0000000..ecbb87e --- /dev/null +++ b/jax_compile/results/local_gpu_NVIDIA_A100_80GB_PCIe/pixelization.json @@ -0,0 +1,82 @@ +[ + { + "transform": "jit", + "trace_s": 4.189, + "compile_s": 5.875, + "first_s": 0.155, + "steady_s": 0.0574, + "dataset_class": "imaging", + "model_type": "pixelization", + "instrument": "hst", + "ndim": 12, + "n_batch": 16, + "batch_size": 4, + "hardware": "local_gpu_NVIDIA_A100_80GB_PCIe", + "jax_version": "0.10.2", + "cache_dir": "", + "mixed_precision": false, + "tag": "a100-at0", + "hostname": "euclid-ral-gpu-2", + "timestamp": "2026-07-17T08:57:00" + }, + { + "transform": "jit", + "trace_s": 4.254, + "compile_s": 8.497, + "first_s": 0.11, + "steady_s": 0.0571, + "dataset_class": "imaging", + "model_type": "pixelization", + "instrument": "hst", + "ndim": 12, + "n_batch": 16, + "batch_size": 4, + "hardware": "local_gpu_NVIDIA_A100_80GB_PCIe", + "jax_version": "0.10.2", + "cache_dir": "", + "mixed_precision": false, + "tag": "a100-at_on", + "hostname": "euclid-ral-gpu-2", + "timestamp": "2026-07-17T08:57:03" + }, + { + "transform": "vag", + "trace_s": 7.178, + "compile_s": 19.73, + "first_s": 0.727, + "steady_s": 0.091, + "dataset_class": "imaging", + "model_type": "pixelization", + "instrument": "hst", + "ndim": 12, + "n_batch": 16, + "batch_size": 4, + "hardware": "local_gpu_NVIDIA_A100_80GB_PCIe", + "jax_version": "0.10.2", + "cache_dir": "", + "mixed_precision": false, + "tag": "a100-at0", + "hostname": "euclid-ral-gpu-2", + "timestamp": "2026-07-17T08:57:34" + }, + { + "transform": "vag", + "trace_s": 7.272, + "compile_s": 27.256, + "first_s": 0.177, + "steady_s": 0.0871, + "dataset_class": "imaging", + "model_type": "pixelization", + "instrument": "hst", + "ndim": 12, + "n_batch": 16, + "batch_size": 4, + "hardware": "local_gpu_NVIDIA_A100_80GB_PCIe", + "jax_version": "0.10.2", + "cache_dir": "", + "mixed_precision": false, + "tag": "a100-at_on", + "hostname": "euclid-ral-gpu-2", + "timestamp": "2026-07-17T08:57:43" + } +] diff --git a/jax_compile/trace_profile.py b/jax_compile/trace_profile.py new file mode 100644 index 0000000..4a43a3b --- /dev/null +++ b/jax_compile/trace_profile.py @@ -0,0 +1,120 @@ +""" +Where does JAX *tracing* time go? (cold-compile research, issue #74) + +Tracing (`jax.jit(fn).lower(x)`) is pure Python, recurs every process, and is +the one cost the persistent compilation cache cannot remove (#71). This probe +cProfiles a single `.lower()` call and aggregates cumulative time by library, +so the uncacheable floor can be attributed (autofit model mapping? autoarray +grids? jax machinery?) before anyone tries to reduce it. + +Usage (from the ``autolens_profiling/`` root):: + + python jax_compile/trace_profile.py --model-type mge --transform vag + python jax_compile/trace_profile.py --model-type pixelization --transform jit --top 40 + +Prints the per-library rollup and the top-N cumulative functions, and appends a +JSON record under ``jax_compile/results//trace_profile.json``. +""" + +from __future__ import annotations + +import argparse +import cProfile +import json +import pstats +import sys +import time +from collections import defaultdict +from pathlib import Path + +_WORKSPACE_ROOT = Path(__file__).resolve().parents[1] +if str(_WORKSPACE_ROOT) not in sys.path: + sys.path.insert(0, str(_WORKSPACE_ROOT)) + +from jax_compile.probe import build_objective, transformed_fn_and_arg, hardware_label + +LIBRARY_KEYS = ( + "autofit", + "autoconf", + "autoarray", + "autogalaxy", + "autolens", + "jax", +) + + +def parse_args(): + p = argparse.ArgumentParser(description=__doc__.split("\n")[1]) + p.add_argument("--dataset-class", default="imaging") + p.add_argument("--model-type", default="mge") + p.add_argument("--instrument", default="hst") + p.add_argument("--transform", default="vag") + p.add_argument("--n-batch", type=int, default=16) + p.add_argument("--batch-size", type=int, default=4) + p.add_argument("--top", type=int, default=25) + p.add_argument("--mixed-precision", action="store_true") + p.add_argument("--tag", default="") + return p.parse_args() + + +def library_of(filename: str) -> str: + for key in LIBRARY_KEYS: + if f"/{key}/" in filename: + return key + return "other" + + +def main(): + args = parse_args() + + import jax + + f, x0, ndim = build_objective(args) + fn, arg = transformed_fn_and_arg(args.transform, f, x0, args.n_batch, args.batch_size) + jitted = jax.jit(fn) + + profiler = cProfile.Profile() + t0 = time.perf_counter() + profiler.enable() + jitted.lower(arg) + profiler.disable() + trace_s = time.perf_counter() - t0 + + stats = pstats.Stats(profiler) + + # Rollup: exclusive (tottime) per library — sums to the trace wall time, so + # the attribution is exact rather than double-counted cumulative time. + per_library = defaultdict(float) + for (filename, _, _), (_, _, tottime, _, _) in stats.stats.items(): + per_library[library_of(filename)] += tottime + + print(f"[trace_profile] {args.model_type} / {args.transform}: trace {trace_s:.2f}s") + print(f"[trace_profile] exclusive-time rollup (sums to ~trace time):") + rollup = dict(sorted(per_library.items(), key=lambda kv: -kv[1])) + for lib, seconds in rollup.items(): + print(f"[trace_profile] {lib:<10} {seconds:7.2f}s ({100 * seconds / trace_s:4.1f}%)") + + print(f"[trace_profile] top {args.top} functions by cumulative time:") + stats.sort_stats("cumulative").print_stats(args.top) + + record = { + "model_type": args.model_type, + "transform": args.transform, + "ndim": ndim, + "trace_s": round(trace_s, 3), + "rollup_tottime_s": {k: round(v, 3) for k, v in rollup.items()}, + "jax_version": jax.__version__, + "tag": args.tag, + "timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"), + } + out_dir = _WORKSPACE_ROOT / "jax_compile" / "results" / hardware_label(jax) + out_dir.mkdir(parents=True, exist_ok=True) + out_path = out_dir / "trace_profile.json" + existing = json.loads(out_path.read_text()) if out_path.exists() else [] + existing.append(record) + out_path.write_text(json.dumps(existing, indent=2) + "\n") + print(f"[trace_profile] wrote record -> {out_path}") + + +if __name__ == "__main__": + main()