From ab1ef67cb78ee426e38b7952041292064111310f Mon Sep 17 00:00:00 2001 From: "Monocline (Goodfire crew)" Date: Fri, 31 Jul 2026 15:55:46 +0000 Subject: [PATCH] Release BF16 loading and explicit data roots Publishes the validated public release-candidate tree with BF16-safe Hugging Face weight loading and required data_root arguments at every entry edge. Crew-Address: agent/irlp --- CLAUDE.md | 22 +++---- Makefile | 20 ++++-- README.md | 4 +- conftest.py | 21 ++++++- param_decomp/autointerp/CLAUDE.md | 3 +- .../scoring/scripts/run_label_scoring.py | 3 +- .../autointerp/scripts/run_interpret.py | 3 +- param_decomp/clustering/CLAUDE.md | 9 +-- .../clustering/scripts/calc_distances.py | 3 +- param_decomp/clustering/scripts/run_merge.py | 3 +- param_decomp/clustering/scripts/run_worker.py | 3 +- param_decomp/core/CLAUDE.md | 2 +- .../core/tests/test_generic_model_io.py | 15 +++-- param_decomp/experiments/lm/load_run.py | 9 +-- param_decomp/experiments/lm/run.py | 3 +- param_decomp/experiments/lm/training.py | 3 +- param_decomp/experiments/resid_mlp/run.py | 3 +- param_decomp/experiments/tms/run.py | 3 +- param_decomp/harvest/CLAUDE.md | 2 +- param_decomp/harvest/scripts/run_intruder.py | 3 +- param_decomp/harvest/scripts/run_merge.py | 3 +- param_decomp/harvest/scripts/run_worker.py | 3 +- param_decomp/infra/paths.py | 10 +-- param_decomp/infra/run_files.py | 8 +-- param_decomp/pretrain/config.py | 3 +- param_decomp/pretrain/train.py | 13 ++-- param_decomp/targets/glu_transformer.py | 12 +++- param_decomp/targets/llama_simple_mlp.py | 11 +++- param_decomp/targets/tests/test_hf_weights.py | 63 +++++++++++++++++++ 29 files changed, 178 insertions(+), 85 deletions(-) create mode 100644 param_decomp/targets/tests/test_hf_weights.py diff --git a/CLAUDE.md b/CLAUDE.md index ce1139cff..e6c8a8d36 100644 --- a/CLAUDE.md +++ b/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 diff --git a/Makefile b/Makefile index 13ab459d6..55eb9cb19 100644 --- a/Makefile +++ b/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 diff --git a/README.md b/README.md index d2d632307..ae34d6845 100644 --- a/README.md +++ b/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 ` (likewise -`...experiments.resid_mlp.run`). The torch +`uv run python -m param_decomp.experiments.tms.run --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. diff --git a/conftest.py b/conftest.py index b4af8fec0..a25fd4a68 100644 --- a/conftest.py +++ b/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") diff --git a/param_decomp/autointerp/CLAUDE.md b/param_decomp/autointerp/CLAUDE.md index 0532ca82c..1c779345c 100644 --- a/param_decomp/autointerp/CLAUDE.md +++ b/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 \ --config_json '{...AutointerpConfig...}' \ - --harvest_subrun_id h-YYYYMMDD_HHMMSS + --harvest_subrun_id h-YYYYMMDD_HHMMSS \ + --data_root ``` `` is the decomposition's identifier — for PD runs, the wandb path diff --git a/param_decomp/autointerp/scoring/scripts/run_label_scoring.py b/param_decomp/autointerp/scoring/scripts/run_label_scoring.py index 4b20dd058..a2e9befb4 100644 --- a/param_decomp/autointerp/scoring/scripts/run_label_scoring.py +++ b/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() diff --git a/param_decomp/autointerp/scripts/run_interpret.py b/param_decomp/autointerp/scripts/run_interpret.py index 5dcc4db93..8ee41cc66 100644 --- a/param_decomp/autointerp/scripts/run_interpret.py +++ b/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) diff --git a/param_decomp/clustering/CLAUDE.md b/param_decomp/clustering/CLAUDE.md index 5ea39a4e6..23ebd1ab7 100644 --- a/param_decomp/clustering/CLAUDE.md +++ b/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 python -m param_decomp.clustering.scripts.run_merge \ - /clustering/harvests/ch-member$i merge_config.json --run-id c-member$i --seed $i + /clustering/harvests/ch-member$i merge_config.json --run-id c-member$i --seed $i \ + --data-root done # Consensus over the member runs: python -m param_decomp.clustering.scripts.calc_distances --ensemble-id e- \ --clustering-run-ids c-member0,c-member1,c-member2,c-member3 \ - --distances-method perm_invariant_hamming + --distances-method perm_invariant_hamming --data-root ``` `calc_distances` creates the ensemble dir and writes: @@ -88,7 +89,7 @@ snapshot. `run_merge` reads it unchanged. ## Data Storage -Stored under `/clustering/` (`param_decomp/clustering/paths.py`; `data_root` is the workers' explicit `--data_root` / `--data-root` arg, default `./out`): +Stored under `/clustering/` (`param_decomp/clustering/paths.py`; `data_root` is the workers' explicit `--data_root` / `--data-root` arg): ``` /clustering/ diff --git a/param_decomp/clustering/scripts/calc_distances.py b/param_decomp/clustering/scripts/calc_distances.py index d3fb180b3..3e45fad53 100644 --- a/param_decomp/clustering/scripts/calc_distances.py +++ b/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, diff --git a/param_decomp/clustering/scripts/run_merge.py b/param_decomp/clustering/scripts/run_merge.py index b480fe75e..f37bdd12d 100644 --- a/param_decomp/clustering/scripts/run_merge.py +++ b/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() diff --git a/param_decomp/clustering/scripts/run_worker.py b/param_decomp/clustering/scripts/run_worker.py index 07160e8c1..9961cc9bd 100644 --- a/param_decomp/clustering/scripts/run_worker.py +++ b/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) diff --git a/param_decomp/core/CLAUDE.md b/param_decomp/core/CLAUDE.md index e3c19145b..cb91b9ebf 100644 --- a/param_decomp/core/CLAUDE.md +++ b/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 [--data-root …]` runs HERE, in the +`python -m param_decomp.experiments.lm.run --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 `/runs//` with the pinned config, and passes `--run-id `. diff --git a/param_decomp/core/tests/test_generic_model_io.py b/param_decomp/core/tests/test_generic_model_io.py index 26497da64..958987925 100644 --- a/param_decomp/core/tests/test_generic_model_io.py +++ b/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={}, ) diff --git a/param_decomp/experiments/lm/load_run.py b/param_decomp/experiments/lm/load_run.py index 8ce04ac81..1cc050934 100644 --- a/param_decomp/experiments/lm/load_run.py +++ b/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 (`/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) diff --git a/param_decomp/experiments/lm/run.py b/param_decomp/experiments/lm/run.py index 703bcb0b4..b93ceadab 100644 --- a/param_decomp/experiments/lm/run.py +++ b/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 [--run-id ...]" + "usage: python -m param_decomp.experiments.lm.run " + " --data-root [--run-id ...]" ) raw = yaml.safe_load(Path(sys.argv[1]).read_text()) runtime = RuntimeConfig.model_validate(raw["runtime"]) diff --git a/param_decomp/experiments/lm/training.py b/param_decomp/experiments/lm/training.py index 4a79a4238..4f7a3c182 100644 --- a/param_decomp/experiments/lm/training.py +++ b/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: diff --git a/param_decomp/experiments/resid_mlp/run.py b/param_decomp/experiments/resid_mlp/run.py index f87ba87cd..5bd023a53 100644 --- a/param_decomp/experiments/resid_mlp/run.py +++ b/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) diff --git a/param_decomp/experiments/tms/run.py b/param_decomp/experiments/tms/run.py index f79176f2d..4a4699a39 100644 --- a/param_decomp/experiments/tms/run.py +++ b/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) diff --git a/param_decomp/harvest/CLAUDE.md b/param_decomp/harvest/CLAUDE.md index 20a22cc49..ea97dd5f5 100644 --- a/param_decomp/harvest/CLAUDE.md +++ b/param_decomp/harvest/CLAUDE.md @@ -112,7 +112,7 @@ python -m param_decomp.harvest.scripts.run_intruder \ The only worker. Opens a JAX run, runs its frozen forward, accumulates into the NumPy `Harvester`. Args: - `--run_dir`: the JAX run dir (`runs/`) (required) -- `--data_root`: the output root the harvest writes under (default `./out`) +- `--data_root`: the required output root the harvest writes under - `--n_batches`, `--batch_size`, `--activation_threshold` - `--rank R --world_size N`: serve `process_index=R`'s slice of every global batch; save to `worker_states/worker_.npz`. Omit both for a single-process run that writes the diff --git a/param_decomp/harvest/scripts/run_intruder.py b/param_decomp/harvest/scripts/run_intruder.py index aba465718..e3aae55a8 100644 --- a/param_decomp/harvest/scripts/run_intruder.py +++ b/param_decomp/harvest/scripts/run_intruder.py @@ -9,14 +9,13 @@ from param_decomp.harvest.db import HarvestDB from param_decomp.harvest.intruder import run_intruder_scoring 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 = DEFAULT_DATA_ROOT, + data_root: Path, ) -> None: assert isinstance(config_json, dict), f"Expected dict from fire, got {type(config_json)}" load_dotenv() diff --git a/param_decomp/harvest/scripts/run_merge.py b/param_decomp/harvest/scripts/run_merge.py index 8e0fb9b76..25b8d367a 100644 --- a/param_decomp/harvest/scripts/run_merge.py +++ b/param_decomp/harvest/scripts/run_merge.py @@ -13,10 +13,9 @@ from param_decomp.harvest.config import HarvestConfig from param_decomp.harvest.pipeline import merge_harvest from param_decomp.harvest.schemas import get_harvest_subrun_dir -from param_decomp.infra.paths import DEFAULT_DATA_ROOT -def main(subrun_id: str, config_json: dict[str, Any], data_root: Path = DEFAULT_DATA_ROOT) -> None: +def main(subrun_id: str, config_json: dict[str, Any], data_root: Path) -> None: assert isinstance(config_json, dict), f"Expected dict from fire, got {type(config_json)}" config = HarvestConfig.model_validate(config_json) output_dir = get_harvest_subrun_dir(Path(data_root), config.method_config.id, subrun_id) diff --git a/param_decomp/harvest/scripts/run_worker.py b/param_decomp/harvest/scripts/run_worker.py index 43caec9ea..a3bd53314 100644 --- a/param_decomp/harvest/scripts/run_worker.py +++ b/param_decomp/harvest/scripts/run_worker.py @@ -34,7 +34,6 @@ from param_decomp.harvest.repo import HarvestRepo from param_decomp.harvest.schemas import HarvestBatch, get_harvest_subrun_dir 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 @@ -139,7 +138,7 @@ def _activation_threshold(config: HarvestConfig) -> float: def main() -> None: ap = argparse.ArgumentParser(description=__doc__) ap.add_argument("--run_dir", type=Path, required=True) - 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("--n_batches", type=int, required=True) ap.add_argument( diff --git a/param_decomp/infra/paths.py b/param_decomp/infra/paths.py index 07e523046..bdff67ae9 100644 --- a/param_decomp/infra/paths.py +++ b/param_decomp/infra/paths.py @@ -8,16 +8,16 @@ """ from pathlib import Path -from typing import Annotated, Final +from typing import Annotated from pydantic import BeforeValidator, PlainSerializer REPO_ROOT = Path(__file__).resolve().parents[2] -"""The checkout this package runs from — a property of the install, not the environment.""" +"""The checkout this package runs from — a property of the install, not the environment. -DEFAULT_DATA_ROOT: Final[Path] = Path("out") -"""The entry-edge default for `data_root` parameters (cwd-relative `./out`). The library -reads no ambient environment for paths: every deployment passes its real root explicitly.""" +`data_root` — the one root of the library's local world — is deliberately NOT declared +here: the library reads no ambient environment for paths and ships no default root, so +every entry edge requires its deployment to pass the real root explicitly.""" def to_root_path(path: str | Path) -> Path: diff --git a/param_decomp/infra/run_files.py b/param_decomp/infra/run_files.py index 9c8142a3c..6b433bd01 100644 --- a/param_decomp/infra/run_files.py +++ b/param_decomp/infra/run_files.py @@ -10,7 +10,7 @@ from wandb.apis.public import Run as WandbRun from param_decomp.core.log import logger -from param_decomp.infra.paths import DEFAULT_DATA_ROOT, ModelPath +from param_decomp.infra.paths import ModelPath from param_decomp.infra.wandb import ( download_wandb_file, fetch_latest_checkpoint_name, @@ -63,7 +63,7 @@ def _wandb_cache_dir(data_root: Path, run_id: str) -> Path: def resolve_run_files( path: ModelPath, *, - data_root: Path = DEFAULT_DATA_ROOT, + data_root: Path, config_filename: str, checkpoint_filename: str | None = None, checkpoint_prefix: str | None = None, @@ -124,9 +124,7 @@ def resolve_run_files( ) -def resolve_config_path( - path: ModelPath, *, data_root: Path = DEFAULT_DATA_ROOT, config_filename: str -) -> Path: +def resolve_config_path(path: ModelPath, *, data_root: Path, config_filename: str) -> Path: """Locate just a run's config file, without resolving or downloading checkpoints.""" try: entity, project, run_id = parse_wandb_run_path(str(path)) diff --git a/param_decomp/pretrain/config.py b/param_decomp/pretrain/config.py index dd5d5ef5e..5e6c52c87 100644 --- a/param_decomp/pretrain/config.py +++ b/param_decomp/pretrain/config.py @@ -100,7 +100,8 @@ def block_size(self) -> int: @property def paths(self) -> PretrainRunPaths: assert self.data_root is not None and self.run_id is not None, ( - "data_root / run_id are minted by the launcher; absent in a hand-authored config" + "run identity incomplete: data_root is authored or launcher-stamped," + " run_id is minted at entry" ) return PretrainRunPaths(data_root=self.data_root, run_id=self.run_id) diff --git a/param_decomp/pretrain/train.py b/param_decomp/pretrain/train.py index 21eeebd07..7baa01b18 100644 --- a/param_decomp/pretrain/train.py +++ b/param_decomp/pretrain/train.py @@ -39,7 +39,6 @@ from param_decomp.core.sharding import hsdp_mesh, init_distributed from param_decomp.infra.dataset_store import resolve_dataset_ref -from param_decomp.infra.paths import DEFAULT_DATA_ROOT from param_decomp.pretrain.batch_data import BatchSchedule, ShardServer, scan_shards from param_decomp.pretrain.cache import ( cache_dir_for, @@ -391,16 +390,14 @@ def main(config: Path) -> None: def _stamp_local_identity(cfg: PretrainConfig) -> PretrainConfig: - """A hand-run carries no launcher stamp: mint an ephemeral identity under the default - data root.""" + """A hand-run carries no launcher stamp: mint an ephemeral run id. `data_root` has no + such fallback — the config must carry it, authored or launcher-stamped.""" import secrets - return cfg.model_copy( - update={ - "data_root": cfg.data_root or DEFAULT_DATA_ROOT, - "run_id": cfg.run_id or f"t-{secrets.token_hex(4)}", - } + assert cfg.data_root is not None, ( + "config carries no data_root: author it in the YAML or launch via a stamping launcher" ) + return cfg.model_copy(update={"run_id": cfg.run_id or f"t-{secrets.token_hex(4)}"}) def _enable_compilation_cache(paths: PretrainRunPaths) -> None: diff --git a/param_decomp/targets/glu_transformer.py b/param_decomp/targets/glu_transformer.py index 6806ee06e..daf446ae3 100644 --- a/param_decomp/targets/glu_transformer.py +++ b/param_decomp/targets/glu_transformer.py @@ -37,7 +37,6 @@ import equinox as eqx import jax import jax.numpy as jnp -import numpy as np from jax.ad_checkpoint import checkpoint_name from jax.sharding import Mesh, NamedSharding from jax.sharding import PartitionSpec as P @@ -1151,8 +1150,15 @@ def __init__(self, snapshot: Path, dtype: DTypeLike): def get(self, key: str) -> Array: fname = self._key_to_file[key] if fname not in self._open: - self._open[fname] = safe_open(str(self._snapshot / fname), framework="numpy") - return jnp.asarray(np.array(self._open[fname].get_tensor(key)), dtype=self._dtype) + # framework="flax", not "numpy": the numpy backend resolves dtypes by name + # through numpy, where bfloat16 exists only if ml_dtypes registration has + # run as an import side effect — flax requires jax, which guarantees it. + self._open[fname] = safe_open(str(self._snapshot / fname), framework="flax") + # Stage on host CPU: get_tensor materializes on the JAX default device, and a + # multi-GB checkpoint must not pass through a single accelerator before its + # placement onto the mesh (`place_via_shardings` serves shards from host). + with jax.default_device(jax.devices("cpu")[0]): + return jnp.asarray(self._open[fname].get_tensor(key), dtype=self._dtype) AttnLoader = Callable[[HFWeights, int], FrozenAttn] diff --git a/param_decomp/targets/llama_simple_mlp.py b/param_decomp/targets/llama_simple_mlp.py index ba5ae995a..2deca5f7e 100644 --- a/param_decomp/targets/llama_simple_mlp.py +++ b/param_decomp/targets/llama_simple_mlp.py @@ -32,7 +32,6 @@ import equinox as eqx import jax import jax.numpy as jnp -import numpy as np import yaml from jax.sharding import Mesh, NamedSharding from jax.sharding import PartitionSpec as P @@ -527,8 +526,14 @@ def checkpoint_safetensors_path(cache_dir: Path) -> Path: def _checkpoint_weight_getter(cache_dir: Path, dtype: DTypeLike) -> WeightGetter: - handle = safe_open(str(checkpoint_safetensors_path(cache_dir)), framework="numpy") - return lambda key: jnp.asarray(np.array(handle.get_tensor(key)), dtype=dtype) # type: ignore[attr-defined] + # framework="flax" + host staging: see `glu_transformer.HFWeights.get`. + handle = safe_open(str(checkpoint_safetensors_path(cache_dir)), framework="flax") + + def get(key: str) -> Array: + with jax.default_device(jax.devices("cpu")[0]): + return jnp.asarray(handle.get_tensor(key), dtype=dtype) + + return get def _layer_from_weights( diff --git a/param_decomp/targets/tests/test_hf_weights.py b/param_decomp/targets/tests/test_hf_weights.py new file mode 100644 index 000000000..abca28b40 --- /dev/null +++ b/param_decomp/targets/tests/test_hf_weights.py @@ -0,0 +1,63 @@ +"""`HFWeights` over a synthetic sharded snapshot — no network, no cached HF checkpoint. + +Real HF checkpoints (SmolLM2, Llama-3.1-8B, Qwen3-8B) are bf16 on disk, so the read path +must handle bf16 without relying on numpy resolving the "bfloat16" dtype name (stock +numpy only can when ml_dtypes registration has run as an import side effect). +""" + +import json +from pathlib import Path + +import jax.numpy as jnp +from safetensors.flax import save_file + +from param_decomp.targets.glu_transformer import HFWeights + +_EMBED_KEY = "model.embed_tokens.weight" +_NORM_KEY = "model.norm.weight" + + +def _write_snapshot(snapshot: Path) -> tuple[jnp.ndarray, jnp.ndarray]: + """A minimal two-shard snapshot: a bf16 tensor and an fp32 tensor, plus the + `model.safetensors.index.json` weight_map `HFWeights.__init__` reads.""" + embed = (jnp.arange(12, dtype=jnp.float32).reshape(3, 4) / 7).astype(jnp.bfloat16) + norm = jnp.linspace(0.5, 1.5, 4, dtype=jnp.float32) + shard_of = { + _EMBED_KEY: "model-00001-of-00002.safetensors", + _NORM_KEY: "model-00002-of-00002.safetensors", + } + save_file({_EMBED_KEY: embed}, snapshot / shard_of[_EMBED_KEY]) + save_file({_NORM_KEY: norm}, snapshot / shard_of[_NORM_KEY]) + (snapshot / "model.safetensors.index.json").write_text(json.dumps({"weight_map": shard_of})) + return embed, norm + + +def test_hf_weights_reads_bf16_shards(tmp_path: Path): + embed, norm = _write_snapshot(tmp_path) + + w = HFWeights(tmp_path, jnp.bfloat16) + got_embed = w.get(_EMBED_KEY) + assert got_embed.dtype == jnp.bfloat16 + assert jnp.array_equal(got_embed, embed) + got_norm = w.get(_NORM_KEY) + assert got_norm.dtype == jnp.bfloat16 + assert jnp.array_equal(got_norm, norm.astype(jnp.bfloat16)) + + +def test_hf_weights_casts_bf16_up_to_fp32(tmp_path: Path): + embed, norm = _write_snapshot(tmp_path) + + w = HFWeights(tmp_path, jnp.float32) + got_embed = w.get(_EMBED_KEY) + assert got_embed.dtype == jnp.float32 + assert jnp.array_equal(got_embed, embed.astype(jnp.float32)) + assert jnp.array_equal(w.get(_NORM_KEY), norm) + + +def test_hf_weights_stages_reads_on_host_cpu(tmp_path: Path): + """Loaded leaves must be host-resident: `place_via_shardings` serves device shards + from the loaded copy, and a multi-GB checkpoint must not pass through a single + accelerator first.""" + _write_snapshot(tmp_path) + got = HFWeights(tmp_path, jnp.bfloat16).get(_EMBED_KEY) + assert all(d.platform == "cpu" for d in got.devices())