diff --git a/param_decomp/CLAUDE.md b/param_decomp/CLAUDE.md index 0914b2a4d..741455671 100644 --- a/param_decomp/CLAUDE.md +++ b/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 diff --git a/param_decomp/README.md b/param_decomp/README.md index 4d4a11de3..81c0d1c66 100644 --- a/param_decomp/README.md +++ b/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.`) onto the current layout (site-keyed `components.vu`, `sources..`) 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`) | diff --git a/param_decomp/arithmetic_eval.py b/param_decomp/arithmetic_eval.py index ef3ab0fba..02bd2983d 100644 --- a/param_decomp/arithmetic_eval.py +++ b/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) diff --git a/param_decomp/compile_cache.py b/param_decomp/compile_cache.py new file mode 100644 index 000000000..6ceeb969b --- /dev/null +++ b/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 diff --git a/param_decomp_lab/clustering/scripts/run_worker.py b/param_decomp_lab/clustering/scripts/run_worker.py index 15539a8d0..c6a7e0a83 100644 --- a/param_decomp_lab/clustering/scripts/run_worker.py +++ b/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) diff --git a/param_decomp_lab/experiments/lm/load_run.py b/param_decomp_lab/experiments/lm/load_run.py index c064d8591..31513b9f8 100644 --- a/param_decomp_lab/experiments/lm/load_run.py +++ b/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), ) diff --git a/param_decomp_lab/experiments/lm/run.py b/param_decomp_lab/experiments/lm/run.py index 167d11a8b..2fcaaccd6 100644 --- a/param_decomp_lab/experiments/lm/run.py +++ b/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 `/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: diff --git a/param_decomp_lab/harvest/scripts/run_worker.py b/param_decomp_lab/harvest/scripts/run_worker.py index ba012bbd4..85ee22e05 100644 --- a/param_decomp_lab/harvest/scripts/run_worker.py +++ b/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( diff --git a/pretrain/train.py b/pretrain/train.py index e53d5cbab..649a94b3f 100644 --- a/pretrain/train.py +++ b/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)