Release BF16 loading and explicit data roots

CLAUDE.md

@@ -95,11 +95,11 @@ submission or code-shipping, no cluster paths, mounts, partitions, or team names
 Deployment fit belongs to whatever launcher invokes the library. A launcher composes
 library entrypoints and may import the library; the library may never import a launcher
 (enforced fail-closed by `param_decomp/core/tests/test_runtime_standalone.py`). The
-library reads NO ambient environment for paths: output roots are explicit parameters
-(default `./out`) threaded from entry points, so a launcher passes its own resolved root
-as an argument. Credentials and third-party-tool conventions (`WANDB_*`, `*_API_KEY`,
-`HF_*`, `CUDA_*`) may follow their ecosystem's own env contract, resolved at entry
-points. Everything else takes typed values. A launcher is not privileged: if it needs
+library reads NO ambient environment for paths: output roots are explicit, required
+parameters threaded from entry points, so a launcher passes its own resolved root as an
+argument. Credentials and third-party-tool conventions (`WANDB_*`, `*_API_KEY`, `HF_*`,
+`CUDA_*`) may follow their ecosystem's own env contract, resolved at entry points.
+Everything else takes typed values. A launcher is not privileged: if it needs
 something the library doesn't publicly expose, that is a library bug — never a reason
 for a private hook. A SLURM *mention* in library prose is legitimate only when it
 documents a generic contract (e.g. SIGTERM→save semantics), never a dependency.
@@ -257,12 +257,12 @@ Both training output and the W&B download cache write here. Per-stage subdirs ar
 populated by their respective pipelines.
 
 `data_root` is the ONE root of the library's local world — runs (outputs), the dataset
-store, the pretrain and compilation caches all hang under it. It is an explicit
-parameter, never an env var: every entry edge (composition roots, worker mains, consumer
-functions like `open_jax_run`) takes it with the default `./out`
-(`param_decomp.infra.paths.DEFAULT_DATA_ROOT`, cwd-relative). A launcher with its own
-notion of shared storage passes that root in explicitly — the `--data_root` /
-`--data-root` flag, or the pretrainer's stamped `data_root`.
+store, the pretrain and compilation caches all hang under it. It is an explicit,
+required parameter, never an env var and never defaulted: every entry edge (composition
+roots, worker mains, consumer functions like `open_jax_run`) refuses to run without it,
+so a deployment cannot silently write into a cwd-relative directory. A launcher passes
+its resolved root explicitly — the `--data_root` / `--data-root` flag, or the
+pretrainer's stamped `data_root`.
 
 ## Development commands
 

Makefile

@@ -85,13 +85,23 @@ test-ci-lab-multidevice:
 	uv run pytest param_decomp/tests/ param_decomp/experiments/ --runslow --durations 10 --numprocesses $(NUM_PROCESSES) --dist worksteal
 	$(MAKE) test-multidevice
 
-# Tests needing >1 device (sharding / checkpoint topology). They hang at the default 1
-# device, so they're skipped in the 1-device passes and run here under SIMULATED CPU
-# devices (XLA_FLAGS). `make test-all` runs this automatically as a second pass; invoke it
-# directly only to run the subset alone (e.g. iterating on sharding/checkpoint).
+# Tests needing >1 device hang at the default 1, so run them under four simulated CPU
+# devices. On small runners, the checkpoint round-trip can starve XLA's CPU worker pool;
+# isolate it and retry only its SIGABRT, while every other failure remains immediate.
+MULTIDEVICE_DEADLOCK_TEST = param_decomp/core/tests/test_checkpoint_production_topology.py::test_sharded_roundtrip_persists_source_moments
+MULTIDEVICE_PYTEST = XLA_FLAGS="--xla_force_host_platform_device_count=4" uv run pytest
+MULTIDEVICE_ARGS = -m multidevice --runmultidevice --durations 10 --capture=tee-sys
+
 .PHONY: test-multidevice
 test-multidevice:
-	XLA_FLAGS="--xla_force_host_platform_device_count=4" uv run pytest $(TEST_PATHS) -m multidevice --runmultidevice --durations 10
+	$(MULTIDEVICE_PYTEST) $(TEST_PATHS) $(MULTIDEVICE_ARGS) --deselect=$(MULTIDEVICE_DEADLOCK_TEST)
+	@run() { $(MULTIDEVICE_PYTEST) $(MULTIDEVICE_DEADLOCK_TEST) $(MULTIDEVICE_ARGS); }; \
+	for attempt in 1 2 3; do \
+		set -x; run; rc=$$?; set +x; \
+		[ $$rc -ne 134 ] && exit $$rc; \
+		[ $$attempt -lt 3 ] && echo "test-multidevice: SIGABRT rc=134 (XLA rendezvous deadlock) -- retry $$attempt/2"; \
+	done; \
+	exit 134
 
 COVERAGE_DIR=docs/coverage
 

README.md

@@ -76,8 +76,8 @@ network access. `python -m param_decomp.pretrain.train` writes the same layout d
 when training a target locally.
 
 TMS and ResidualMLP run the same way — in-process module mains, on CPU:
-`uv run python -m param_decomp.experiments.tms.run <config.yaml>` (likewise
-`...experiments.resid_mlp.run`). The torch
+`uv run python -m param_decomp.experiments.tms.run <config.yaml> --data-root <data-root>`
+(likewise `...experiments.resid_mlp.run`). The torch
 trainer is preserved only at git tag `torch-oracle`; current training uses the JAX
 single-pool engine. See `param_decomp/core/SPEC.md` for its numerical contract and
 `param_decomp/experiments/CLAUDE.md` for the complete LM config schema.

conftest.py

@@ -4,7 +4,8 @@
 """
 
 import os
-from collections.abc import Iterable
+import sys
+from collections.abc import Iterable, Iterator
 from netrc import netrc
 from pathlib import Path
 from urllib.parse import urlparse
@@ -37,6 +38,24 @@
 jax.config.update("jax_persistent_cache_min_compile_time_secs", 1.0)
 jax.config.update("jax_persistent_cache_min_entry_size_bytes", 0)
 
+# Each live XLA CPU executable consumes several non-mergeable VMAs. A serial suite can
+# exhaust Linux's `vm.max_map_count`, whereupon XLA aborts instead of raising. Clear at
+# module boundaries once half full, leaving room for `test_checkpoint.py`'s ~19k mappings.
+_MAPPING_CEILING = (
+    int(Path("/proc/sys/vm/max_map_count").read_text()) if sys.platform == "linux" else None
+)
+
+
+def _mapping_count() -> int:
+    return Path("/proc/self/maps").read_bytes().count(b"\n")
+
+
+@pytest.fixture(autouse=True, scope="module")
+def _bounded_jax_executable_mappings() -> Iterator[None]:
+    yield
+    if _MAPPING_CEILING is not None and _mapping_count() > _MAPPING_CEILING // 2:
+        jax.clear_caches()
+
 
 def pytest_addoption(parser: Parser) -> None:
     parser.addoption("--runslow", action="store_true", default=False, help="run slow tests")

param_decomp/autointerp/CLAUDE.md

@@ -14,7 +14,8 @@ every CLI takes a `--harvest_subrun_id`.
 # One process, inline JSON config
 python -m param_decomp.autointerp.scripts.run_interpret <decomposition_id> \
     --config_json '{...AutointerpConfig...}' \
-    --harvest_subrun_id h-YYYYMMDD_HHMMSS
+    --harvest_subrun_id h-YYYYMMDD_HHMMSS \
+    --data_root <data-root>
 ```
 
 `<decomposition_id>` is the decomposition's identifier — for PD runs, the wandb path

param_decomp/autointerp/scoring/scripts/run_label_scoring.py

@@ -21,7 +21,6 @@
     load_component_keys_file,
 )
 from param_decomp.harvest.repo import HarvestRepo
-from param_decomp.infra.paths import DEFAULT_DATA_ROOT
 
 LabelScorerType = Literal["detection", "fuzzing"]
 
@@ -31,8 +30,8 @@ def main(
     scorer_type: LabelScorerType,
     config_json: dict[str, Any],
     harvest_subrun_id: str,
+    data_root: Path,
     autointerp_subrun_id: str | None = None,
-    data_root: Path = DEFAULT_DATA_ROOT,
 ) -> None:
     assert isinstance(config_json, dict), f"Expected dict from fire, got {type(config_json)}"
     load_dotenv()

param_decomp/autointerp/scripts/run_interpret.py

@@ -28,15 +28,14 @@
 )
 from param_decomp.core.log import logger
 from param_decomp.harvest.repo import HarvestRepo
-from param_decomp.infra.paths import DEFAULT_DATA_ROOT
 
 
 def main(
     decomposition_id: str,
     config_json: dict[str, Any],
     harvest_subrun_id: str,
+    data_root: Path,
     autointerp_subrun_id: str | None = None,
-    data_root: Path = DEFAULT_DATA_ROOT,
 ) -> None:
     assert isinstance(config_json, dict), f"Expected dict from fire, got {type(config_json)}"
     interp_config = AutointerpConfig.model_validate(config_json)

param_decomp/clustering/CLAUDE.md

@@ -49,15 +49,16 @@ N seeded harvests (1 GPU each)      run_worker --dataset_seed i
 for i in 0 1 2 3; do
   python -m param_decomp.clustering.scripts.run_worker \
       --run_dir runs/p-xxxxxxxx --n_tokens 50000 --batch_size 16 --n_tokens_per_seq 16 \
-      --dataset_seed $i --harvest_id ch-member$i
+      --dataset_seed $i --harvest_id ch-member$i --data_root <data-root>
   python -m param_decomp.clustering.scripts.run_merge \
-      <data-root>/clustering/harvests/ch-member$i merge_config.json --run-id c-member$i --seed $i
+      <data-root>/clustering/harvests/ch-member$i merge_config.json --run-id c-member$i --seed $i \
+      --data-root <data-root>
 done
 
 # Consensus over the member runs:
 python -m param_decomp.clustering.scripts.calc_distances --ensemble-id e-<id> \
     --clustering-run-ids c-member0,c-member1,c-member2,c-member3 \
-    --distances-method perm_invariant_hamming
+    --distances-method perm_invariant_hamming --data-root <data-root>
 ```
 
 `calc_distances` creates the ensemble dir and writes:
@@ -88,7 +89,7 @@ snapshot. `run_merge` reads it unchanged.
 
 ## Data Storage
 
-Stored under `<data_root>/clustering/` (`param_decomp/clustering/paths.py`; `data_root` is the workers' explicit `--data_root` / `--data-root` arg, default `./out`):
+Stored under `<data_root>/clustering/` (`param_decomp/clustering/paths.py`; `data_root` is the workers' explicit `--data_root` / `--data-root` arg):
 
 ```
 <data_root>/clustering/

param_decomp/clustering/scripts/calc_distances.py

@@ -26,7 +26,6 @@
 from param_decomp.clustering.plotting.merge import plot_dists_distribution
 from param_decomp.clustering.types import DistancesArray, DistancesMethod
 from param_decomp.core.log import logger
-from param_decomp.infra.paths import DEFAULT_DATA_ROOT
 
 
 def calc_distances(
@@ -112,7 +111,7 @@ def cli() -> None:
         choices=DistancesMethod.__args__,
         default="perm_invariant_hamming",
     )
-    parser.add_argument("--data-root", type=Path, default=DEFAULT_DATA_ROOT)
+    parser.add_argument("--data-root", type=Path, required=True)
     args = parser.parse_args()
     calc_distances(
         ensemble_id=args.ensemble_id,

param_decomp/clustering/scripts/run_merge.py

@@ -36,7 +36,6 @@
     ComponentLabels,
 )
 from param_decomp.core.log import logger
-from param_decomp.infra.paths import DEFAULT_DATA_ROOT
 
 
 def _make_iteration_plot_callback(plot_dir: Path, plot_every: int) -> LogCallback:
@@ -159,7 +158,7 @@ def cli() -> None:
     parser.add_argument("merge_config", type=Path)
     parser.add_argument("--run-id", type=str, default=None)
     parser.add_argument("--seed", type=int, default=0)
-    parser.add_argument("--data-root", type=Path, default=DEFAULT_DATA_ROOT)
+    parser.add_argument("--data-root", type=Path, required=True)
     parser.add_argument("--plot", action="store_true", help="emit per-run diagnostic plots")
     args = parser.parse_args()
     run_id = args.run_id or new_run_id()

param_decomp/clustering/scripts/run_worker.py

@@ -27,7 +27,6 @@
 from param_decomp.core.log import logger
 from param_decomp.experiments.lm.load_run import LoadedJaxRun, open_jax_run
 from param_decomp.infra.dataset_store import read_dataset_meta
-from param_decomp.infra.paths import DEFAULT_DATA_ROOT
 from param_decomp.pretrain.batch_data import BatchSchedule, ShardServer, scan_shards
 
 
@@ -149,7 +148,7 @@ def main() -> None:
         help="HarvestConfig JSON. If given, --run_dir/--n_tokens/etc. are ignored.",
     )
     ap.add_argument("--harvest_id", type=str, default=None, help="pre-assigned harvest id")
-    ap.add_argument("--data_root", type=Path, default=DEFAULT_DATA_ROOT)
+    ap.add_argument("--data_root", type=Path, required=True)
     ap.add_argument("--step", type=int, default=None, help="checkpoint step (default: latest)")
     ap.add_argument("--run_dir", type=Path, default=None)
     ap.add_argument("--n_tokens", type=int, default=None)

param_decomp/core/CLAUDE.md

@@ -339,7 +339,7 @@ fully determine process bring-up — there is no launch field:
   up via `jax.distributed`'s own cluster auto-detection (`init_distributed` — the jax
   ecosystem's contract). Multiple processes on one node is deliberately unrepresentable.
 
-`python -m param_decomp.experiments.lm.run <config> [--data-root …]` runs HERE, in the
+`python -m param_decomp.experiments.lm.run <config> --data-root …` runs HERE, in the
 current allocation, minting and pinning its own identity when `--run-id` is absent. A
 launcher that wants to own the identity mints the `p-` run id itself, stages
 `<data_root>/runs/<id>/` with the pinned config, and passes `--run-id <id>`.

param_decomp/core/tests/test_generic_model_io.py

@@ -27,6 +27,7 @@
 import equinox as eqx
 import jax
 import jax.numpy as jnp
+import numpy as np
 import optax
 import pytest
 from jax import random
@@ -49,7 +50,6 @@
 from param_decomp.core.objective import build_objective
 from param_decomp.core.recon_eval import FreshPGDReconEval, make_fresh_pgd_eval_step
 from param_decomp.core.schedule import Knot, ScheduleConfig
-from param_decomp.core.sharding import hsdp_mesh
 from param_decomp.core.train import Decomposition, TrainingItem, TrainState, make_train_step
 
 B, T, D, C = 2, 3, 8, 5
@@ -218,6 +218,11 @@ def _synthetic_inputs(key: jax.Array) -> dict[str, Array]:
     }
 
 
+def _one_device_mesh() -> jax.sharding.Mesh:
+    devices = np.asarray(jax.devices()[:1]).reshape(1, 1, 1)
+    return jax.sharding.Mesh(devices, ("replicate", "fsdp", "tp"))
+
+
 def test_dict_input_tuple_output_and_geometric_loss_flow():
     """The model consumes the loader's native DICT batch (not token ids);
     `clean_output`/`masked_output` emit a tuple; `recon_loss_fn` (MSE) contracts it."""
@@ -266,9 +271,9 @@ def test_train_step_runs_through_generic_target(with_mesh: bool):
     """End-to-end: the real `make_train_step` drives the synthetic dict-in/tuple-out/MSE
     target for two steps; the loss stays finite and the trainable V/U actually move.
 
-    Run with AND without a mesh: the sharding constraints are no-ops off-mesh, so a
-    meshless run cannot see whether the batch/output edges survive being sharded — which
-    is how an array-only `batch_sharded` passed CI while dying on every real run."""
+    Run meshless and with an explicit one-device mesh: only the latter exercises
+    `NamedSharding`, while pinning one device keeps this generic-I/O contract independent
+    of the ambient device count. Multi-device partitioning is tested separately."""
     key = random.PRNGKey(2)
     model = _synthetic_lm(key)
     components = _synthetic_vu(key)
@@ -297,7 +302,7 @@ def test_train_step_runs_through_generic_target(with_mesh: bool):
         total_steps=10,
         remat_recon_forwards=False,
         remat_ci_fn=False,
-        mesh=hsdp_mesh(gpus_per_node=1) if with_mesh else None,
+        mesh=_one_device_mesh() if with_mesh else None,
         compiler_options={},
     )
 

param_decomp/experiments/lm/load_run.py

@@ -7,7 +7,7 @@
 `decomposition` item (the trained V/U + ci_fn — optimizer/adversary state is training's
 business and is never touched), and exposes the pure forward a consumer needs:
 
-    run = open_jax_run(run_dir)                 # latest checkpoint
+    run = open_jax_run(run_dir, data_root=data_root)   # latest checkpoint
     fwd = run.forward(token_ids)                # one frozen, forward-only pass
     fwd.lower_leaky_ci[site]                    # (B, T, C) leaky CI per site
     fwd.component_acts[site]                    # (B, T, C) ‖U_c‖ · (x @ V) per site
@@ -54,7 +54,6 @@
 )
 from param_decomp.experiments.lm.resolved import LMRun, weights_jnp_dtype
 from param_decomp.infra import pretrain_cache
-from param_decomp.infra.paths import DEFAULT_DATA_ROOT
 from param_decomp.targets import llama_simple_mlp
 from param_decomp.targets.glu_transformer import glu_site_specs
 
@@ -190,9 +189,7 @@ def _restore_decomposition(
     return decomposition, resolved_step
 
 
-def open_jax_run(
-    run_dir: Path, step: int | None = None, *, data_root: Path = DEFAULT_DATA_ROOT
-) -> LoadedJaxRun:
+def open_jax_run(run_dir: Path, step: int | None = None, *, data_root: Path) -> LoadedJaxRun:
     """Open the run at `run_dir`; restore checkpoint `step` (latest if None). Restores
     only the trained decomposition (see `_restore_decomposition`). `data_root` resolves a
     `kind: pretrained` target's cache (`<data_root>/pretrain_cache/...`)."""
@@ -263,7 +260,7 @@ class RunMetadata:
     layer_activation_sizes: list[tuple[str, int]]
 
 
-def run_metadata(run_dir: Path, *, data_root: Path = DEFAULT_DATA_ROOT) -> RunMetadata:
+def run_metadata(run_dir: Path, *, data_root: Path) -> RunMetadata:
     """Target topology for `run_dir`, derived from the pinned config (+ the SimpleMLP
     pretrain cache's `model_config.yaml` for `n_layer`/`vocab_size`). No orbax restore."""
     cfg, _ = load_config(run_dir / LAUNCH_CONFIG_FILENAME, run_dir.name, data_root)

param_decomp/experiments/lm/run.py

@@ -11,7 +11,8 @@
 
 def main() -> None:
     assert len(sys.argv) >= 2 and not sys.argv[1].startswith("-"), (
-        "usage: python -m param_decomp.experiments.lm.run <config.yaml> [--run-id ...]"
+        "usage: python -m param_decomp.experiments.lm.run <config.yaml>"
+        " --data-root <path> [--run-id ...]"
     )
     raw = yaml.safe_load(Path(sys.argv[1]).read_text())
     runtime = RuntimeConfig.model_validate(raw["runtime"])

param_decomp/experiments/lm/training.py

@@ -46,7 +46,6 @@
 from param_decomp.experiments.lm.resolved import LMRun
 from param_decomp.experiments.lm.runtime import RuntimeConfig
 from param_decomp.infra.dataset_store import read_dataset_meta
-from param_decomp.infra.paths import DEFAULT_DATA_ROOT
 from param_decomp.infra.run_files import generate_run_id
 from param_decomp.pretrain.batch_data import BatchSchedule, ShardServer, scan_shards
 
@@ -198,7 +197,7 @@ def _pin_config_copy(run_dir: Path, name: str, source: Path) -> None:
         copy.write_text(source.read_text())
 
 
-def main(config: Path, run_id: str | None = None, data_root: Path = DEFAULT_DATA_ROOT) -> None:
+def main(config: Path, data_root: Path, run_id: str | None = None) -> None:
     config = Path(config)
     data_root = Path(data_root)
     if run_id is None:

param_decomp/experiments/resid_mlp/run.py

@@ -50,7 +50,6 @@
 from param_decomp.experiments.resid_mlp.config import ResidMLPExperimentConfig
 from param_decomp.experiments.toy_config import build_toy_ci_arch
 from param_decomp.experiments.toy_eval import ToyRun, make_toy_evaluation_operations
-from param_decomp.infra.paths import DEFAULT_DATA_ROOT
 from param_decomp.infra.run_files import generate_run_id
 from param_decomp.targets import resid_mlp
 
@@ -271,10 +270,10 @@ def ground_truth_eval(context: EvalInvocation) -> LogRecord:
 
 def main(
     config: str,
+    data_root: Path,
     run_id: str | None = None,
     group: str | None = None,
     tags: str | tuple[str, ...] | None = None,
-    data_root: Path = DEFAULT_DATA_ROOT,
 ) -> None:
     schema_raw = yaml.safe_load(Path(config).read_text())
     data_root = Path(data_root)

param_decomp/experiments/tms/run.py

@@ -49,7 +49,6 @@
 from param_decomp.experiments.tms.config import TMSExperimentConfig
 from param_decomp.experiments.toy_config import build_toy_ci_arch
 from param_decomp.experiments.toy_eval import ToyRun, make_toy_evaluation_operations
-from param_decomp.infra.paths import DEFAULT_DATA_ROOT
 from param_decomp.infra.run_files import generate_run_id
 from param_decomp.targets import tms
 
@@ -235,10 +234,10 @@ def ground_truth_eval(context: EvalInvocation) -> LogRecord:
 
 def main(
     config: str,
+    data_root: Path,
     run_id: str | None = None,
     group: str | None = None,
     tags: str | tuple[str, ...] | None = None,
-    data_root: Path = DEFAULT_DATA_ROOT,
 ) -> None:
     schema_raw = yaml.safe_load(Path(config).read_text())
     data_root = Path(data_root)