Migrate VPD to JAX: train + analyze in one framework; retire torch to oracle

.github/workflows/checks.yaml

@@ -64,32 +64,10 @@ jobs:
         run: uv run ruff format .
 
       - name: Run tests
-        run: |
-          NUM_PROCESSES=$(nproc | awk '{print ($1<4?$1:4)}')
-          echo "Using $NUM_PROCESSES pytest workers"
-          uv run pytest param_decomp/tests/ param_decomp_lab/tests/ --runslow --durations 10 --numprocesses $NUM_PROCESSES --dist worksteal
+        # `make test-all` runs the 1-device suite then the multidevice subset (sim CPU devices).
+        run: make test-all
         env:
           WANDB_API_KEY: ${{ secrets.WANDB_API_KEY }}
           OMPI_ALLOW_RUN_AS_ROOT: 1
           OMPI_ALLOW_RUN_AS_ROOT_CONFIRM: 1
           HF_HOME: ${{ github.workspace }}/.hf-cache
-  build-frontend:
-    runs-on: ubuntu-latest
-    defaults:
-      run:
-        working-directory: param_decomp_lab/app/frontend
-    steps:
-      - name: Checkout
-        uses: actions/checkout@v6
-
-      - name: Install dependencies
-        run: npm install
-
-      - name: Run svelte-check
-        run: npm run check
-
-      - name: Run linter
-        run: npm run lint
-
-      - name: Build frontend
-        run: npm run build

.gitignore

@@ -1,9 +1,14 @@
 docs/coverage/**
 notebooks/**
 scratch/
+.scratch_*
 
 # Script outputs (generated files, often large)
 scripts/outputs/
+scripts/standalone_repros/*.json
+scripts/standalone_repros/*.txt
+scripts/profile_ci_fn_*.json
+scripts/profile_ci_fn_*.txt
 
 **/out/
 neuronpedia_outputs/
@@ -180,3 +185,8 @@ cython_debug/
 *.schema.json
 
 .claude/worktrees
+
+.venv-cuda/
+**/.venv-cuda/
+
+nano_param_decomp_jax/jax_single_pool/experiments/logs/

Makefile

@@ -1,4 +1,9 @@
 # setup
+# ONE venv for the whole workspace: the JAX trainer core (`param_decomp` + `pretrain` +
+# `vendored_jax`) is the root distribution and carries jax as a normal dependency, so a
+# single `uv sync --all-packages` installs core + config + lab into one `.venv`. The CPU
+# jax wheel is the base; the CUDA wheel is the `[cuda]` extra the per-run launch workspace
+# installs.
 .PHONY: install
 install:
 	uv sync --no-dev
@@ -10,41 +15,21 @@ install-lab:
 .PHONY: install-dev
 install-dev:
 	uv sync --all-packages
-	uv run pre-commit install
-
-.PHONY: install-all
-install-all: install-dev install-app
-
-
-.PHONY: app
-app:
-	@uv run --package param-decomp-lab python -m param_decomp_lab.app.run_app
-
-.PHONY: install-app
-install-app:
-	(cd param_decomp_lab/app/frontend && npm install)
-
-.PHONY: check-app
-check-app:
-	(cd param_decomp_lab/app/frontend && npm run format && npm run check && npm run lint)
+	uv run --no-sync pre-commit install
 
 # special install for CI (GitHub Actions) that reduces disk usage and install time
 # 1. create a fresh venv with `--clear` -- this is mostly only for local testing of the CI install
 # 2. install with `uv sync` but with some special options:
 #  > `--frozen` to enforce using the lock file for consistent dependency versions
 #  > `--link-mode copy` because symlinks/hardlinks dont work half the time anyway
-#  > `--extra-index-url` to get cpu-only pytorch wheels. installing with just `uv sync` will download a bunch of cuda stuff we cannot use anyway, since there are no GPUs anyways. takes up a lot of space and makes the install take 5x as long
-#  > `--index-strategy unsafe-best-match` because pytorch index won't have every version of each package we need. markupsafe is a particular pain point
 # Note: explored the `--compile-bytecode` option for test speedups, nothing came of it. see https://github.com/goodfire-ai/param-decomp/pull/187/commits/740f6a28f4d3378078c917125356b6466f155e71
 .PHONY: install-ci
 install-ci:
 	uv venv --python 3.13 --clear
 	uv sync \
 		--frozen \
 		--all-packages \
-		--link-mode copy \
-		--extra-index-url https://download.pytorch.org/whl/cpu \
-		--index-strategy unsafe-best-match
+		--link-mode copy
 
 # checks
 .PHONY: type
@@ -66,22 +51,36 @@ check-pre-commit:
 
 # tests
 
+# `param_decomp/tests/` is the JAX trainer core suite (incl. the LM equivalence goldens);
+# `param_decomp_lab/{tests,experiments}/` the lab suites (the toy TMS/ResidMLP tests live
+# beside their models under experiments/).
+TEST_PATHS = param_decomp/tests/ param_decomp_lab/tests/ param_decomp_lab/experiments/
+
 .PHONY: test
 test:
-	uv run pytest param_decomp/tests/ param_decomp_lab/tests/ --testmon --durations 10
+	uv run pytest $(TEST_PATHS) --testmon --durations 10
 
 # Use min(4, nproc) for numprocesses. Any more and it slows down the tests.
 NUM_PROCESSES ?= $(shell nproc | awk '{print ($$1<4?$$1:4)}')
 
 .PHONY: test-all
 test-all:
-	uv run pytest param_decomp/tests/ param_decomp_lab/tests/ --runslow --durations 10 --numprocesses $(NUM_PROCESSES) --dist worksteal
+	uv run pytest $(TEST_PATHS) --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).
+.PHONY: test-multidevice
+test-multidevice:
+	XLA_FLAGS="--xla_force_host_platform_device_count=4" uv run pytest $(TEST_PATHS) -m multidevice --runmultidevice --durations 10
 
 COVERAGE_DIR=docs/coverage
 
 .PHONY: coverage
 coverage:
-	uv run pytest param_decomp/tests/ param_decomp_lab/tests/ --cov=param_decomp --cov=param_decomp_lab --runslow
+	uv run pytest $(TEST_PATHS) --cov=param_decomp --cov=param_decomp_lab --runslow
 	mkdir -p $(COVERAGE_DIR)
 	uv run python -m coverage report -m > $(COVERAGE_DIR)/coverage.txt
 	uv run python -m coverage html --directory=$(COVERAGE_DIR)/html/

PERF_NOTES.md

@@ -0,0 +1,22 @@
+# full32L HSDP perf — synthesized into Lore
+
+The full perf investigation (V/U reconstruction hoist → buffer donation → b128/4-seq-per-GPU
+at ~3.6× the b32 throughput) has been distilled into Lore. This file is intentionally short:
+a 1100-line chronological trail of hypotheses, corrections, and dead-ends is not worth carrying
+on the trunk.
+
+**Canonical understanding** (mechanisms, exact memory anatomy, throughput, open levers):
+lore `2026-06-29--full32l-hsdp-donation-canonical-scaling-memory`.
+
+**Closed levers / measured-null results** (so they aren't re-chased — collective-combine,
+latency-hiding, NCCL NVLS/CUMEM/PROTO, command buffers, unroll-by-K, scan-unroll,
+replicate-weights, TP-without-SP): lore `2026-06-29--full32l-mfu-closed-levers-and-measured-facts`.
+
+**Open comm lever** (CI-fn weight-grad cross-node reduce un-batched in the scan backward):
+lore `2026-06-29--full32l-ci-fn-weight-grad-reduce-unbatched-in-scan`.
+
+**Trainer-architecture canon** (JAX single-pool; torch 2/3-pool retired):
+lore `2026-06-29--state--trainer-architecture`.
+
+The full chronological trail is preserved in the git history of the `perf/hsdp-mfu` and
+`perf/gather-coalesce-unroll-k` branches (this file's prior revisions).

README.md

@@ -29,66 +29,26 @@ The `pd-*` commands are installed by `param-decomp-lab`. Each in-repo experiment
 self-contained script that reads a YAML and calls `optimize()`:
 
 ```bash
-pd-tms       param_decomp_lab/experiments/tms/tms_5-2_config.yaml
-pd-resid-mlp param_decomp_lab/experiments/resid_mlp/resid_mlp1_config.yaml
-pd-lm        param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml
+pd-lm    param_decomp_lab/experiments/lm/<wrapper>.yaml --nodes N
 ```
 
-For a brand-new experiment, write your own `run.py` that builds the target model, the
-train/eval dataloaders, the eval `Metric` list, the `PDConfig` and `RuntimeConfig`, a
-`Cadence` (when to emit), and a `RunSink` (where output goes), then calls `optimize(...)`:
-
-```python
-from param_decomp.configs import Cadence, PDConfig, RuntimeConfig
-from param_decomp.optimize import EvalLoop, optimize
-from param_decomp_lab.batch_and_loss_fns import recon_loss_mse, run_batch_first_element
-from param_decomp_lab.run_sink import RunSink
-
-optimize(
-    target_model=my_target_module,
-    train_loader=train_loader,
-    run_batch=run_batch_first_element,
-    reconstruction_loss=recon_loss_mse,
-    pd_config=PDConfig(...),
-    runtime_config=RuntimeConfig(device=device),
-    cadence=Cadence(
-        train_log_every=100,
-        save_every=5000,
-    ),
-    sink=RunSink.local(out_dir),
-    eval_loop=EvalLoop(
-        loader=eval_loader,
-        metrics=[...],  # list of pre-instantiated Metric objects
-        n_steps=10,
-        every=1000,
-        slow_every=5000,
-    ),
-)
-```
+TMS and ResidualMLP now live only as JAX targets in `param_decomp/`
+(`tms.py`, `resid_mlp.py`); the torch experiment dirs were deleted.
 
-The three in-repo `run.py` files
-([tms](param_decomp_lab/experiments/tms/run.py),
- [resid_mlp](param_decomp_lab/experiments/resid_mlp/run.py),
- [lm](param_decomp_lab/experiments/lm/run.py)) are reference examples.
+Training is the JAX single-pool trainer: the generic engine
+(`param_decomp.run.run_decomposition_training`, a pure library) driven by the lab-side
+composition root (`python -m param_decomp_lab.experiments.lm.run`), launched via `pd-lm`. A
+run is one self-contained YAML (the `param_decomp_lab.experiments.config.ExperimentConfig`
+schema over the core `param_decomp.configs` pieces). The torch trainer (`optimize()`, the
+torch `Metric` impls, `RunSink`) was retired and is preserved at git tag `torch-oracle`. See
+`param_decomp/CLAUDE.md` and `SPEC.md` for the trainer.
 
 ## Metrics
 
-Configure training losses in `pd.loss_metrics` as a list of `{type: "<ClassName>", ...}`
-entries. The `type` literal dispatches to a `Metric` subclass via
-`param_decomp.metrics.dispatch.LOSS_METRIC_CLASSES`. Loss metrics must set `coeff`; they
-are evaluated automatically alongside dedicated eval metrics. New loss metrics are added
-by defining the class in `param_decomp/metrics/`, appending the config to
-`AnyLossMetricConfig` in `configs.py`, and appending the class to `LOSS_METRIC_CLASSES`.
-
-Eval metrics are caller-supplied: instantiate `Metric` objects in your `run.py` and pass
-them via `EvalLoop(metrics=...)`. The in-repo experiments validate the YAML
-`eval.metrics` list via the `AnyEvalMetricConfig` discriminated union on `EvalConfig`,
-then dispatch each entry through `EVAL_METRIC_CLASSES` (both in
-`param_decomp_lab.eval_metrics`):
-
-```python
-eval_metrics = [EVAL_METRIC_CLASSES[m.type](m) for m in cfg.eval.metrics]
-```
+Training losses are configured in `pd.loss_metrics` as a list of `{type: "<ClassName>",
+...}` entries; eval metrics in `eval.metrics`. Both are validated by the torch-free pydantic
+schema in core (`param_decomp.configs`) and computed by the JAX trainer
+(`param_decomp/losses.py`, `slow_eval.py`).
 
 ## Packaging
 

conftest.py

@@ -15,11 +15,23 @@
 
 def pytest_addoption(parser: Parser) -> None:
     parser.addoption("--runslow", action="store_true", default=False, help="run slow tests")
+    parser.addoption(
+        "--runmultidevice",
+        action="store_true",
+        default=False,
+        help="run tests needing >1 jax device (requires XLA_FLAGS=--xla_force_host_platform_device_count>=2; "
+        "use `make test-multidevice`)",
+    )
 
 
 def pytest_configure(config: Config) -> None:
     config.addinivalue_line("markers", "slow: mark test as slow to run")
     config.addinivalue_line("markers", "requires_wandb: mark test as requiring WANDB credentials")
+    config.addinivalue_line(
+        "markers",
+        "multidevice: needs >1 jax device; hangs at the default 1 device, so skipped unless "
+        "--runmultidevice (use `make test-multidevice`, which sets XLA_FLAGS for simulated CPU devices)",
+    )
 
 
 def _wandb_host() -> str:
@@ -52,13 +64,17 @@ def _have_wandb_credentials() -> bool:
 
 def pytest_collection_modifyitems(config: Config, items: Iterable[Item]) -> None:
     runslow = config.getoption("--runslow")
+    runmultidevice = config.getoption("--runmultidevice")
     have_wandb = _have_wandb_credentials()
     skip_slow = pytest.mark.skip(reason="need --runslow option to run")
+    skip_multidevice = pytest.mark.skip(reason="needs >1 device; run via `make test-multidevice`")
     skip_wandb = pytest.mark.skip(
         reason="No WANDB credentials (set WANDB_API_KEY or login via CLI)"
     )
     for item in items:
         if "slow" in item.keywords and not runslow:
             item.add_marker(skip_slow)
+        if "multidevice" in item.keywords and not runmultidevice:
+            item.add_marker(skip_multidevice)
         if "requires_wandb" in item.keywords and not have_wandb:
             item.add_marker(skip_wandb)