fix(infra): persistent XLA compile cache in harvest/clustering workers; one policy definition

param_decomp/CLAUDE.md

@@ -288,13 +288,16 @@ Requeues rebuild the node workspaces and re-read the pinned launch config, never
 checkout. `--run_id` resubmits an existing run. Don't hand-write sbatch files.
 
 `main` enables JAX's persistent compilation cache
-(`_enable_persistent_compilation_cache`) at `$PARAM_DECOMP_OUT_DIR/xla_compilation_cache`
-— a SIBLING of `runs/` (derived from `out_dir.parent`), shared across all runs and all
-8N ranks, NOT per-run. The multi-minute chunkwise-step compile is keyed by HLO + backend +
-topology + jax/xla version, so a requeue/resume or a fresh run at the same config+topology
-loads the executable from disk in seconds. Set after `init_distributed` (the write gate
-reads the distributed state) and before the first compile; threshold 60s
-(`jax_persistent_cache_min_compile_time_secs`) so only the big compiles cache. Multi-host
+(`param_decomp.compile_cache.enable_persistent_compilation_cache` — the ONE policy
+definition, also called at startup by `pretrain.train` and the harvest/clustering
+workers) at `$PARAM_DECOMP_OUT_DIR/xla_compilation_cache` — a SIBLING of `runs/`
+(derived from `out_dir.parent`), shared across all runs and all 8N ranks, NOT per-run.
+The multi-minute chunkwise-step compile is keyed by HLO + backend + topology + jax/xla
+version, so a requeue/resume or a fresh run at the same config+topology loads the
+executable from disk in seconds. Set before the first compile (the cache initializes
+lazily, so before/after `init_distributed` both work); threshold 5s
+(`jax_persistent_cache_min_compile_time_secs`) so step-level compiles cache but trivial
+slice/utility jits don't churn the dir. Multi-host
 safe on jax 0.10.1: jax gates the cache WRITE on `process_id == 0` (`compiler.py` — "Only
 write cache entries from the first process … contention for writes on some filesystems"),
 so all ranks read but only rank 0 writes — no shared-FS race. Requires the cache dir on a

param_decomp/README.md

@@ -31,6 +31,7 @@ root with sibling packages `pretrain/` (the in-house target-LM pretrainer) and
 | `run_state.py` | optimizer + initial-`TrainState` construction from an `ExperimentConfig` (orbax restores onto this reference) |
 | `tools/` | `convert_llama_simple_mlp_checkpoint.py` (torch venv) — one-off `.pt` → safetensors conversion of the pile pretrain checkpoint; `migrate_c49k_checkpoint.py` — one-off remap of the frozen C49k clone's orbax `TrainState` (legacy `components.{Vg..Ud}` `(1,*,*)` + flat `sources.<site>`) onto the current layout (site-keyed `components.vu`, `sources.<state_key>.<site>`) so a fine-tune can `restore_latest` it |
 | `sharding.py` | generic GSPMD helpers (`init_distributed`, `dp_mesh`, `replicate`, `shard_batch`) |
+| `compile_cache.py` | `enable_persistent_compilation_cache` — the one definition of the persistent XLA compile-cache policy (dir sibling of `runs/`, thresholds); called at startup by the trainer, pretrainer, and harvest/clustering workers |
 | `targets/llama8b.py` | Llama-3.1-8B target (`LlamaDecomposedModel`, full-model token-input forward, embed internal), arbitrary per-layer matrix sites (`q/k/v/o/gate/up/down`, per-site C; q/k/v decomposed before RoPE/SDPA), per-site `DecompVU`, HF safetensors loader (`build_decomposed_lm` / `load_decomposed_lm_from_hf`) |
 | `targets/llama_simple_mlp.py` | `LlamaSimpleMLP` pile-pretrained target (`goodfire/spd/runs/t-9d2b8f02`: 4L, d768, GELU MLP, plain rotate-half RoPE, tied head): sites `h.{i}.attn.{q,k,v,o}_proj` / `h.{i}.mlp.{c_fc,down_proj}` with `h.*` wildcard expansion, pretrain-cache safetensors loader (one-off `.pt` conversion: `tools/convert_llama_simple_mlp_checkpoint.py`), `llama_simple_mlp_decomposed_lm(cfg, sites)`; frozen weights small enough to replicate (`replicate_frozen`), V/U/CI/source placement reuses the generic per-site plan |
 | `run.py` | the generic ENGINE `run_decomposition_training` (pure library, no `main`/YAML): faith warmup, loop, metrics jsonl/wandb, in-loop slow renderer, orbax checkpoints, SIGTERM-save + requeue-resume. The LM composition root that reads YAML + builds the target lives lab-side (`param_decomp_lab/experiments/lm/run.py`) |

param_decomp/arithmetic_eval.py

@@ -18,6 +18,7 @@
 
 from param_decomp.ci_fn import lower_leaky_hard_sigmoid
 from param_decomp.components import DecompVU
+from param_decomp.jit_util import filter_jit
 from param_decomp.lm import DecomposedModel
 from param_decomp.train import COMPUTE_DT, cast_floating
 
@@ -80,7 +81,10 @@ def to_grid(self, per_prompt: np.ndarray) -> np.ndarray:
 
 
 def make_arithmetic_grid_step(
-    lm: ComponentActivationModel, answer_position: int, n_valid_rows: int
+    lm: ComponentActivationModel,
+    answer_position: int,
+    n_valid_rows: int,
+    compiler_options: dict[str, bool | int | str] | None = None,
 ) -> ArithmeticGridStep:
     """Build the jit'd step returning, at `answer_position` with the batch axis KEPT as the
     grid, BOTH per-component lower-leaky CI (from the CI fn) and the pre-mask activation `x@V`
@@ -93,7 +97,6 @@ def make_arithmetic_grid_step(
 
     # HLO-baking rule: read STATIC config (site_names, Cs) off the closed-over `lm`; all array
     # access goes through the traced `model` arg.
-    @eqx.filter_jit
     def step(
         model: ComponentActivationModel,
         components: DecompVU,
@@ -123,9 +126,11 @@ def step(
         max_ci = {s: jnp.where(valid, ci[s], -jnp.inf).max(axis=0) for s in site_names}
         return ci, xv, max_ci
 
-    return step
+    return filter_jit(step, compiler_options=compiler_options)
 
 
+# Trivial gather; deliberately compiled WITHOUT `compiler_options` (module-level, no config
+# in scope — GPU compiler flags don't move a `jnp.take`).
 @eqx.filter_jit
 def _take_columns(per_prompt: Float[Array, "n_pad C"], idx: Int[Array, " k"]) -> Array:
     return jnp.take(per_prompt, idx, axis=1)

param_decomp/compile_cache.py

@@ -0,0 +1,38 @@
+"""The one definition of the persistent XLA compilation-cache policy."""
+
+from pathlib import Path
+
+import jax
+
+
+def enable_persistent_compilation_cache(runs_dir: Path) -> Path:
+    """Cache compiled XLA executables to a shared-FS dir reused across runs/requeues.
+
+    Every jax entrypoint (trainer, pretrainer, harvest/clustering workers) calls this at
+    startup so a compile is keyed by HLO + backend + topology + jax/xla version and a
+    matching re-compile (requeue, another worker of the same fan-out, a fresh run at the
+    same config+topology) loads from disk in seconds. The dir is a SIBLING of `runs_dir`
+    (not per-run, not inside the immutable per-run workspace) so every run and consumer
+    shares it. Multi-host safe: jax gates writes on `process_id == 0`; every rank reads.
+    Independent single-process workers all write, but entries land via atomic rename, so
+    same-key writers just overwrite each other with identical bytes.
+
+    Call before the first compile that should cache; the cache initializes lazily at
+    first compile (before/after `init_distributed` both work — the write gate reads the
+    distributed state at write time, not here). Rank-0's HLO dump (`XLA_FLAGS`
+    `--xla_dump_*`, see the LM composition root's `_enable_hlo_dump`) does not fork
+    rank 0's cache key: jax strips debug/dump options when computing it.
+
+    Do NOT add `jax_persistent_cache_enable_xla_caches='all'` — it crashes at runtime on
+    jax 0.10.1/B200 (`CUDA_ERROR_NOT_FOUND`); the default already persists the
+    per-fusion autotune cache on jax >= 0.10.
+    """
+    cache_dir = runs_dir.parent / "xla_compilation_cache"
+    jax.config.update("jax_compilation_cache_dir", str(cache_dir))
+    # 5s (jax default: 1s): high enough that trivial slice/utility jits don't churn the
+    # dir, low enough that every step-level compile caches — eval tiers and init fans
+    # (1-60s) recompiled on every requeue under the old 60s threshold, and the worker
+    # forward must not miss the cache just because it compiles in under a minute.
+    jax.config.update("jax_persistent_cache_min_compile_time_secs", 5.0)
+    jax.config.update("jax_persistent_cache_min_entry_size_bytes", 0)
+    return cache_dir

param_decomp_lab/clustering/scripts/run_worker.py

@@ -22,6 +22,7 @@
 import numpy as np
 
 from param_decomp.built_run import DataConfig
+from param_decomp.compile_cache import enable_persistent_compilation_cache
 from param_decomp.data import BatchSchedule, ShardServer, scan_shards
 from param_decomp.log import logger
 from param_decomp_lab.clustering.harvest_config import HarvestConfig
@@ -106,6 +107,8 @@ def run_harvest_to_dir(config: HarvestConfig, harvest_id: str, step: int | None)
     dir for `harvest_id`, and return that dir. Used by the standalone CLI and the ensemble
     pipeline (which pre-assigns the harvest id so its dependent merge can find the snapshot).
     """
+    cache_dir = enable_persistent_compilation_cache(Path(config.model_path).parent)
+    logger.info(f"persistent XLA compilation cache: {cache_dir}")
     run = open_jax_run(Path(config.model_path), step)
     output_dir = clustering_harvest_dir(harvest_id)
     output_dir.mkdir(parents=True, exist_ok=True)

param_decomp_lab/experiments/lm/load_run.py

@@ -27,7 +27,6 @@
 from dataclasses import dataclass
 from pathlib import Path
 
-import equinox as eqx
 import jax
 import jax.numpy as jnp
 from jaxtyping import Array, Float, Int
@@ -36,6 +35,7 @@
 from param_decomp.checkpoint import make_checkpoint_manager, restore_latest, restore_step
 from param_decomp.ci_fn import CIFn
 from param_decomp.components import DecompVU
+from param_decomp.jit_util import filter_jit
 from param_decomp.lm import DecomposedModel
 from param_decomp.run_state import build_optimizers, init_train_state
 from param_decomp.sharding import hsdp_mesh, place_via_shardings
@@ -176,7 +176,6 @@ def open_jax_run(run_dir: Path, step: int | None = None) -> LoadedJaxRun:
 
     # `model` is the filter_jit ARG (frozen weights traced, not baked). It embeds the token
     # ids internally — the harvest forward feeds tokens straight in.
-    @eqx.filter_jit
     def forward(
         model: DecomposedModel,
         components: DecompVU,
@@ -210,7 +209,9 @@ def forward(
         config=cfg,
         vocab_size=vocab_size,
         _state=state,
-        _forward=forward,
+        # The run's pinned `compiler_options` so the consumer forward compiles with the
+        # same XLA flags as training (they enter the compile-cache key).
+        _forward=filter_jit(forward, compiler_options=cfg.runtime.compiler_options),
     )
 
 

param_decomp_lab/experiments/lm/run.py

@@ -61,6 +61,7 @@
     EvalConfig,
     TargetSites,
 )
+from param_decomp.compile_cache import enable_persistent_compilation_cache
 from param_decomp.configs import ResumeProvenance
 from param_decomp.data import BatchSchedule, ShardServer, scan_shards
 from param_decomp.eval import make_eval_step
@@ -100,24 +101,6 @@
 from param_decomp_lab.experiments.lm.load_run import build_target
 
 
-def _enable_persistent_compilation_cache(out_dir: Path) -> Path:
-    """Cache compiled XLA executables to a shared-FS dir reused across runs/requeues.
-
-    The ~24-min compile of the chunkwise step is keyed by HLO + backend + topology +
-    jax/xla version, so a matching re-compile (requeue, or a fresh run at the same
-    config+topology) loads from disk in seconds. The dir is a SIBLING of `runs/` (not
-    per-run, not inside the immutable per-run workspace) so every run shares it; all 8N
-    ranks point at the same shared-FS path. Only process 0 writes (jax gates the write on
-    `process_id == 0` to avoid shared-FS write contention); every rank reads. Must run
-    after `init_distributed` (the rank gate reads the distributed state) and before the
-    first compile."""
-    cache_dir = out_dir.parent / "xla_compilation_cache"
-    jax.config.update("jax_compilation_cache_dir", str(cache_dir))
-    jax.config.update("jax_persistent_cache_min_compile_time_secs", 60.0)
-    jax.config.update("jax_persistent_cache_min_entry_size_bytes", 0)
-    return cache_dir
-
-
 def _enable_hlo_dump(run_dir: Path) -> None:
     """Dump the step modules' optimized HLO + buffer assignment to `<run_dir>/hlo` (rank 0).
 
@@ -234,7 +217,9 @@ def _make_arithmetic_eval(
     probe = build_arithmetic_probe(arith.operation, arith.a_range, arith.b_range, tokenizer)
     n_prompts = probe.tokens.shape[0]
     return _ArithmeticEval(
-        step=make_arithmetic_grid_step(lm, probe.answer_position, n_prompts),
+        step=make_arithmetic_grid_step(
+            lm, probe.answer_position, n_prompts, compiler_options=compiler_options
+        ),
         probe_eval_step=make_eval_step(
             lm,
             eval_cfg.rounding_threshold,
@@ -548,7 +533,7 @@ def main(config: Path, run_id: str) -> None:
     if built.run.resume_provenance is not None:
         assert_finetune_structural_compat(built, built.run.resume_provenance)
 
-    cache_dir = _enable_persistent_compilation_cache(built.run.out_dir)
+    cache_dir = enable_persistent_compilation_cache(built.run.out_dir)
 
     is_main = jax.process_index() == 0
     if is_main:

param_decomp_lab/harvest/scripts/run_worker.py

@@ -28,6 +28,7 @@
 import numpy as np
 
 from param_decomp.built_run import DataConfig
+from param_decomp.compile_cache import enable_persistent_compilation_cache
 from param_decomp.data import BatchSchedule, ShardServer, scan_shards
 from param_decomp.log import logger
 from param_decomp_lab.experiments.lm.load_run import HarvestForward, LoadedJaxRun, open_jax_run
@@ -152,6 +153,8 @@ def main() -> None:
     args = ap.parse_args()
     assert (args.rank is None) == (args.world_size is None)
 
+    cache_dir = enable_persistent_compilation_cache(args.run_dir.parent)
+    logger.info(f"persistent XLA compilation cache: {cache_dir}")
     run = open_jax_run(args.run_dir, args.step)
     subrun_id = args.subrun_id or "h-" + datetime.now().strftime("%Y%m%d_%H%M%S")
     config = HarvestConfig(

pretrain/train.py

@@ -37,6 +37,7 @@
 from jaxtyping import Array, Float, Int
 from orbax.checkpoint.type_handlers import ArrayHandler, register_type_handler
 
+from param_decomp.compile_cache import enable_persistent_compilation_cache
 from param_decomp.data import BatchSchedule, ShardServer, scan_shards
 from param_decomp.sharding import hsdp_mesh, init_distributed
 from pretrain.cache import (
@@ -380,7 +381,8 @@ def main(config: Path) -> None:
     if cfg.out_dir is None or cfg.run_id is None:
         # Local hand-run without the launcher: mint an ephemeral identity under cwd.
         cfg = _mint_local_identity(cfg)
-    _maybe_enable_compilation_cache(cfg)
+    assert cfg.out_dir is not None
+    enable_persistent_compilation_cache(cfg.out_dir)
     train(cfg)
 
 
@@ -392,14 +394,6 @@ def _mint_local_identity(cfg: PretrainConfig) -> PretrainConfig:
     return cfg.model_copy(update={"out_dir": out_dir, "run_id": run_id})
 
 
-def _maybe_enable_compilation_cache(cfg: PretrainConfig) -> None:
-    if cfg.out_dir is None:
-        return
-    cache = cfg.out_dir.parent / "xla_compilation_cache"
-    jax.config.update("jax_compilation_cache_dir", str(cache))
-    jax.config.update("jax_persistent_cache_min_compile_time_secs", 60)
-
-
 def cli() -> None:
     fire.Fire(main)