diff --git a/CONFIGS.md b/CONFIGS.md
new file mode 100644
index 000000000..21e40f604
--- /dev/null
+++ b/CONFIGS.md
@@ -0,0 +1,71 @@
+# Config policy
+
+The repo maintains a small, named set of experiment configs — the **canonical
+seats** — and nothing else. Every yaml the repo carries is a maintenance
+obligation: each schema change must migrate it, forever. Seats are capped at
+**fewer than 10 LM configs**, each with a stated purpose.
+
+## Why committing sweep configs adds nothing
+
+A launched run's config provenance already lives in three places, none of them
+the repo tree:
+
+1. the run dir's pinned `launch_config.yaml` (immutable; resume byte-compares it),
+2. the git snapshot ref `refs/runs/snapshot/<id>` taken at submit,
+3. the wandb run config.
+
+So a sweep/profile/one-off yaml committed "for the record" records nothing —
+it only rots. Launch one-offs from your workspace (`pd-lm <path>` takes any
+path); if the sweep matters, its lore finding cites the run ids, and the run
+dirs carry the exact configs.
+
+## The canonical seats
+
+| seat | file | purpose |
+|---|---|---|
+| llama8b L18 | `param_decomp/configs/llama8b_l18_C49k_200k.yaml` | the L18-MLP decomposition flagship recipe |
+| llama8b full-model | `param_decomp/configs/llama8b_full32L_seq512_b128_dp128.yaml` | the full-32L production recipe |
+| pile 4L PPGD testbed | `param_decomp/configs/pile_ppgd_bsc.yaml` | the "jose-ish" 4L pile algorithm-research testbed |
+| save-path smoke | `param_decomp/configs/llama8b_full32L_HSDP_b32_dp32_SAVESMOKE.yaml` | cheap end-to-end save/resume smoke launch |
+| config-suite fixture | `param_decomp/configs/llama8b_l18_b128_cmp32.yaml` | the representative full config the core config/resume tests load (`test_config.py`, `test_finetune_resume.py`, `test_llama_simple_mlp.py`) |
+| chunkwise fixture | `param_decomp/configs/llama8b_l18-26_9layer_chunkwise.yaml` | the 27-site chunkwise CI-fn config `test_config.py` converts |
+| ss 2L SimpleMLP | `param_decomp_lab/experiments/lm/ss_llama_simple_mlp-2L.yaml` | small-LM regression archetype — **broken at tip, awaiting migration** |
+| jose (gpt2-arch 4L) | `param_decomp_lab/experiments/lm/jose.yaml` | architecture-research testbed — **broken at tip, awaiting migration** |
+| pile 4L SimpleMLP | `param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml` | pretrained-target archetype — **broken at tip, awaiting migration** |
+
+The toy testbeds (`param_decomp_lab/experiments/tms/configs/`,
+`param_decomp_lab/experiments/resid_mlp/configs/`) and the pretrain configs
+(`pretrain/configs/`) are separate small schemas, maintained with their
+experiments; they are seats too, just not LM-schema ones.
+
+## Rules
+
+1. **Every LM config yaml in the tree parses at tip** — CI-enforced by
+   `param_decomp_lab/tests/test_repo_configs_parse.py` (schema parse +
+   `assert_canonical_algorithm_config`). A schema PR that breaks one migrates
+   it **in the same PR**, with an executed in-repo migration (the #966
+   pattern) — never a script attached to a PR comment (#939 attached one; it
+   never ran, and 97 of 104 stored runs became unopenable before anyone
+   noticed).
+2. **Sweep / profile / one-off configs are not committed.** Launch them from a
+   workspace path. Deleting a config is cheap (git history keeps it; run dirs
+   pin what actually ran) — un-rotting one is not.
+3. **Adding a canonical seat is taking on a maintenance obligation**: add the
+   file, a registry row above naming its purpose, and it's covered by the CI
+   gate from that commit on. If the table is at 10 LM seats, something leaves
+   before something enters. A config a test loads is a seat by definition —
+   deleting it is a test change, so grep for the basename first.
+4. **Stored-run pins are immutable.** Never migrate a run dir's
+   `launch_config.yaml` in place (resume byte-compares it; a live old-code run
+   whose pin is rewritten refuses its next requeue). Consumers must not
+   require a full-schema parse of stored pins — they parse only the fields
+   they read (`ConsumerRunConfig`, PR #952).
+
+## Case history
+
+- **#939** (2026-07-03): ScheduleConfig unification; migration script attached
+  as a PR comment, never executed → 7/104 stored runs parseable at tip.
+- **#966**: the counter-example — carve migration shipped as an in-repo,
+  executed tool covering every live repo yaml.
+- **#982**: 25 sweep yamls in one PR — the accumulation pattern this policy
+  ends. The sweeps' findings live in lore; the run dirs pin their configs.
diff --git a/param_decomp/configs/llama8b_full32L_1node_b8_dp8_PROFILE.yaml b/param_decomp/configs/llama8b_full32L_1node_b8_dp8_PROFILE.yaml
deleted file mode 100644
index 986276ed8..000000000
--- a/param_decomp/configs/llama8b_full32L_1node_b8_dp8_PROFILE.yaml
+++ /dev/null
@@ -1,586 +0,0 @@
-run_name: jax-full32L-HSDP-b32-dp32-PROFILE
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 1
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 32
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 8
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 0
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 12
-runtime:
-  dp: 8
-  tp: 1
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_HSDP_b128_dp32.yaml b/param_decomp/configs/llama8b_full32L_HSDP_b128_dp32.yaml
deleted file mode 100644
index 7b325ef94..000000000
--- a/param_decomp/configs/llama8b_full32L_HSDP_b128_dp32.yaml
+++ /dev/null
@@ -1,611 +0,0 @@
-# First full-model decomposition of Llama-3.1-8B: ALL 32 layers x 7 matrices
-# (q/k/v/o/gate/up/down) = 224 sites. De-risking launch — get it compiling, fitting
-# memory at the 128-GPU topology, and producing a sane training signal; then dial in
-# GPU count / chunking / LRs.
-#
-# Derived from the l18-23 6L reference run p-2112148d ("jax-l18-23-6L-ablbase-200k"):
-# same algorithm (chunkwise subset recon + persistent-PGD + faith + imp-min), same CI
-# fn arch, same LRs/coeffs/eval set. DELIBERATE EDITS vs that reference:
-#   1. sites: 18 (MLP, layers 18-23) -> 224 (all matrices, all 32 layers). Per-matrix C
-#      from the 2026-06-22 full-llama8b BOTEC, rounded UP to multiples of 128 (the V/U
-#      C-axis shards over the dp=128 mesh; init asserts C % dp == 0):
-#      q/k 2048, v/o 4096, gate/up 8192, down 10240. ΣC=1,245,184; V/U=18.3B params
-#      (~1.1x the reference's 16.3B — params shard cheaply, not the binding constraint).
-#   2. runtime.dp: 64 -> 128 (16 nodes). batch 128 -> per-rank 1 (vs 2 in the reference);
-#      the botec puts per-rank-1 at ~105 GiB, safe under the 180 GiB B200 cap.
-#   3. ChunkwiseSubsetReconLoss sites_per_chunk: 9 -> 56 (8 layers/chunk -> 4 chunks).
-#      Fewer, larger chunks for the first launch: minimises compile time and per-step
-#      compute (each chunk's suffix forward now spans the full network). coeff scaled
-#      0.5 x n_chunks = 2.0 to keep per-site recon pressure chunk-count-invariant.
-#      THIS IS THE KNOB TO RE-TUNE once we know step time (finer chunks = better signal).
-#   4. remat_recon_forwards: true (224 sites' full-network chunk forwards need it).
-#   5. remat_ci_fn: true (checkpoint the ~31B CI fn; without it the step is the GPU
-#      compile wall — ~80 min + near-OOM — since the whole CI forward materialises).
-#   6. data seq 512 (fineweb_llama_tok_512), matching the reference's seq.
-# out_dir / run_id are minted by pd-lm.
-run_name: jax-full32L-HSDP-b128-dp32
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 1000
-  train_log_every: 25
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 32
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 100000
-  slow_on_first_step: false
-pd:
-  batch_size: 128
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.83e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.83e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 0
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 50000
-runtime:
-  dp: 32
-  tp: 1
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_HSDP_b128_dp32_PROFILE.yaml b/param_decomp/configs/llama8b_full32L_HSDP_b128_dp32_PROFILE.yaml
deleted file mode 100644
index 6ae585d8c..000000000
--- a/param_decomp/configs/llama8b_full32L_HSDP_b128_dp32_PROFILE.yaml
+++ /dev/null
@@ -1,586 +0,0 @@
-run_name: jax-full32L-HSDP-b128-dp32-PROFILE
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 1
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 32
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 128
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 0
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 12
-runtime:
-  dp: 32
-  tp: 1
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_HSDP_b128_dp32_SAVESMOKE.yaml b/param_decomp/configs/llama8b_full32L_HSDP_b128_dp32_SAVESMOKE.yaml
deleted file mode 100644
index 471ecbb0b..000000000
--- a/param_decomp/configs/llama8b_full32L_HSDP_b128_dp32_SAVESMOKE.yaml
+++ /dev/null
@@ -1,611 +0,0 @@
-# First full-model decomposition of Llama-3.1-8B: ALL 32 layers x 7 matrices
-# (q/k/v/o/gate/up/down) = 224 sites. De-risking launch — get it compiling, fitting
-# memory at the 128-GPU topology, and producing a sane training signal; then dial in
-# GPU count / chunking / LRs.
-#
-# Derived from the l18-23 6L reference run p-2112148d ("jax-l18-23-6L-ablbase-200k"):
-# same algorithm (chunkwise subset recon + persistent-PGD + faith + imp-min), same CI
-# fn arch, same LRs/coeffs/eval set. DELIBERATE EDITS vs that reference:
-#   1. sites: 18 (MLP, layers 18-23) -> 224 (all matrices, all 32 layers). Per-matrix C
-#      from the 2026-06-22 full-llama8b BOTEC, rounded UP to multiples of 128 (the V/U
-#      C-axis shards over the dp=128 mesh; init asserts C % dp == 0):
-#      q/k 2048, v/o 4096, gate/up 8192, down 10240. ΣC=1,245,184; V/U=18.3B params
-#      (~1.1x the reference's 16.3B — params shard cheaply, not the binding constraint).
-#   2. runtime.dp: 64 -> 128 (16 nodes). batch 128 -> per-rank 1 (vs 2 in the reference);
-#      the botec puts per-rank-1 at ~105 GiB, safe under the 180 GiB B200 cap.
-#   3. ChunkwiseSubsetReconLoss sites_per_chunk: 9 -> 56 (8 layers/chunk -> 4 chunks).
-#      Fewer, larger chunks for the first launch: minimises compile time and per-step
-#      compute (each chunk's suffix forward now spans the full network). coeff scaled
-#      0.5 x n_chunks = 2.0 to keep per-site recon pressure chunk-count-invariant.
-#      THIS IS THE KNOB TO RE-TUNE once we know step time (finer chunks = better signal).
-#   4. remat_recon_forwards: true (224 sites' full-network chunk forwards need it).
-#   5. remat_ci_fn: true (checkpoint the ~31B CI fn; without it the step is the GPU
-#      compile wall — ~80 min + near-OOM — since the whole CI forward materialises).
-#   6. data seq 512 (fineweb_llama_tok_512), matching the reference's seq.
-# out_dir / run_id are minted by pd-lm.
-run_name: jax-full32L-HSDP-b128-savesmoke
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 20
-  train_log_every: 5
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 32
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 128
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.83e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.83e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 0
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 50
-runtime:
-  dp: 32
-  tp: 1
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_HSDP_b32_dp32.yaml b/param_decomp/configs/llama8b_full32L_HSDP_b32_dp32.yaml
deleted file mode 100644
index 2449a9ca4..000000000
--- a/param_decomp/configs/llama8b_full32L_HSDP_b32_dp32.yaml
+++ /dev/null
@@ -1,611 +0,0 @@
-# First full-model decomposition of Llama-3.1-8B: ALL 32 layers x 7 matrices
-# (q/k/v/o/gate/up/down) = 224 sites. De-risking launch — get it compiling, fitting
-# memory at the 128-GPU topology, and producing a sane training signal; then dial in
-# GPU count / chunking / LRs.
-#
-# Derived from the l18-23 6L reference run p-2112148d ("jax-l18-23-6L-ablbase-200k"):
-# same algorithm (chunkwise subset recon + persistent-PGD + faith + imp-min), same CI
-# fn arch, same LRs/coeffs/eval set. DELIBERATE EDITS vs that reference:
-#   1. sites: 18 (MLP, layers 18-23) -> 224 (all matrices, all 32 layers). Per-matrix C
-#      from the 2026-06-22 full-llama8b BOTEC, rounded UP to multiples of 128 (the V/U
-#      C-axis shards over the dp=128 mesh; init asserts C % dp == 0):
-#      q/k 2048, v/o 4096, gate/up 8192, down 10240. ΣC=1,245,184; V/U=18.3B params
-#      (~1.1x the reference's 16.3B — params shard cheaply, not the binding constraint).
-#   2. runtime.dp: 64 -> 128 (16 nodes). batch 128 -> per-rank 1 (vs 2 in the reference);
-#      the botec puts per-rank-1 at ~105 GiB, safe under the 180 GiB B200 cap.
-#   3. ChunkwiseSubsetReconLoss sites_per_chunk: 9 -> 56 (8 layers/chunk -> 4 chunks).
-#      Fewer, larger chunks for the first launch: minimises compile time and per-step
-#      compute (each chunk's suffix forward now spans the full network). coeff scaled
-#      0.5 x n_chunks = 2.0 to keep per-site recon pressure chunk-count-invariant.
-#      THIS IS THE KNOB TO RE-TUNE once we know step time (finer chunks = better signal).
-#   4. remat_recon_forwards: true (224 sites' full-network chunk forwards need it).
-#   5. remat_ci_fn: true (checkpoint the ~31B CI fn; without it the step is the GPU
-#      compile wall — ~80 min + near-OOM — since the whole CI forward materialises).
-#   6. data seq 512 (fineweb_llama_tok_512), matching the reference's seq.
-# out_dir / run_id are minted by pd-lm.
-run_name: jax-full32L-HSDP-b32-dp32-smoke
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 25
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 32
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 32
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 0
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 200000
-runtime:
-  dp: 32
-  tp: 1
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_HSDP_b32_dp32_PROFILE.yaml b/param_decomp/configs/llama8b_full32L_HSDP_b32_dp32_PROFILE.yaml
deleted file mode 100644
index 740d00e28..000000000
--- a/param_decomp/configs/llama8b_full32L_HSDP_b32_dp32_PROFILE.yaml
+++ /dev/null
@@ -1,611 +0,0 @@
-# First full-model decomposition of Llama-3.1-8B: ALL 32 layers x 7 matrices
-# (q/k/v/o/gate/up/down) = 224 sites. De-risking launch — get it compiling, fitting
-# memory at the 128-GPU topology, and producing a sane training signal; then dial in
-# GPU count / chunking / LRs.
-#
-# Derived from the l18-23 6L reference run p-2112148d ("jax-l18-23-6L-ablbase-200k"):
-# same algorithm (chunkwise subset recon + persistent-PGD + faith + imp-min), same CI
-# fn arch, same LRs/coeffs/eval set. DELIBERATE EDITS vs that reference:
-#   1. sites: 18 (MLP, layers 18-23) -> 224 (all matrices, all 32 layers). Per-matrix C
-#      from the 2026-06-22 full-llama8b BOTEC, rounded UP to multiples of 128 (the V/U
-#      C-axis shards over the dp=128 mesh; init asserts C % dp == 0):
-#      q/k 2048, v/o 4096, gate/up 8192, down 10240. ΣC=1,245,184; V/U=18.3B params
-#      (~1.1x the reference's 16.3B — params shard cheaply, not the binding constraint).
-#   2. runtime.dp: 64 -> 128 (16 nodes). batch 128 -> per-rank 1 (vs 2 in the reference);
-#      the botec puts per-rank-1 at ~105 GiB, safe under the 180 GiB B200 cap.
-#   3. ChunkwiseSubsetReconLoss sites_per_chunk: 9 -> 56 (8 layers/chunk -> 4 chunks).
-#      Fewer, larger chunks for the first launch: minimises compile time and per-step
-#      compute (each chunk's suffix forward now spans the full network). coeff scaled
-#      0.5 x n_chunks = 2.0 to keep per-site recon pressure chunk-count-invariant.
-#      THIS IS THE KNOB TO RE-TUNE once we know step time (finer chunks = better signal).
-#   4. remat_recon_forwards: true (224 sites' full-network chunk forwards need it).
-#   5. remat_ci_fn: true (checkpoint the ~31B CI fn; without it the step is the GPU
-#      compile wall — ~80 min + near-OOM — since the whole CI forward materialises).
-#   6. data seq 512 (fineweb_llama_tok_512), matching the reference's seq.
-# out_dir / run_id are minted by pd-lm.
-run_name: jax-full32L-HSDP-b32-dp32-PROFILE
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 1
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 32
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 32
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 0
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 12
-runtime:
-  dp: 32
-  tp: 1
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_HSDP_b64_dp32_PROFILE.yaml b/param_decomp/configs/llama8b_full32L_HSDP_b64_dp32_PROFILE.yaml
deleted file mode 100644
index 2d3f8d924..000000000
--- a/param_decomp/configs/llama8b_full32L_HSDP_b64_dp32_PROFILE.yaml
+++ /dev/null
@@ -1,586 +0,0 @@
-run_name: jax-full32L-HSDP-b64-dp32-PROFILE
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 1
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 32
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 64
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 0
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 12
-runtime:
-  dp: 32
-  tp: 1
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_HSDP_b96_dp32_PROFILE.yaml b/param_decomp/configs/llama8b_full32L_HSDP_b96_dp32_PROFILE.yaml
deleted file mode 100644
index 1a4961db6..000000000
--- a/param_decomp/configs/llama8b_full32L_HSDP_b96_dp32_PROFILE.yaml
+++ /dev/null
@@ -1,586 +0,0 @@
-run_name: jax-full32L-HSDP-b96-dp32-PROFILE
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 1
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 32
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 96
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 0
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 12
-runtime:
-  dp: 32
-  tp: 1
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_cifn4chunk_b32_dp32_PROFILE.yaml b/param_decomp/configs/llama8b_full32L_cifn4chunk_b32_dp32_PROFILE.yaml
deleted file mode 100644
index c6eb86c9f..000000000
--- a/param_decomp/configs/llama8b_full32L_cifn4chunk_b32_dp32_PROFILE.yaml
+++ /dev/null
@@ -1,586 +0,0 @@
-run_name: jax-full32L-HSDP-b32-dp32-PROFILE
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 1
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 32
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 32
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 8
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 0
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 12
-runtime:
-  dp: 32
-  tp: 1
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_noPGD_b32_dp32_PROFILE.yaml b/param_decomp/configs/llama8b_full32L_noPGD_b32_dp32_PROFILE.yaml
deleted file mode 100644
index 6221cffcc..000000000
--- a/param_decomp/configs/llama8b_full32L_noPGD_b32_dp32_PROFILE.yaml
+++ /dev/null
@@ -1,571 +0,0 @@
-run_name: jax-full32L-HSDP-b32-dp32-PROFILE
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 1
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 32
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 32
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 0
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 12
-runtime:
-  dp: 32
-  tp: 1
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_noPGD_b64_dp32_PROFILE.yaml b/param_decomp/configs/llama8b_full32L_noPGD_b64_dp32_PROFILE.yaml
deleted file mode 100644
index 2954b2d34..000000000
--- a/param_decomp/configs/llama8b_full32L_noPGD_b64_dp32_PROFILE.yaml
+++ /dev/null
@@ -1,571 +0,0 @@
-run_name: jax-full32L-noPGD-b64-dp32-PROFILE
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 1
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 32
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 64
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 0
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 12
-runtime:
-  dp: 32
-  tp: 1
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_nw0_b32_dp32_PROFILE.yaml b/param_decomp/configs/llama8b_full32L_nw0_b32_dp32_PROFILE.yaml
deleted file mode 100644
index b52941d28..000000000
--- a/param_decomp/configs/llama8b_full32L_nw0_b32_dp32_PROFILE.yaml
+++ /dev/null
@@ -1,586 +0,0 @@
-run_name: jax-full32L-HSDP-b32-dp32-PROFILE
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 1
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 32
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 32
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 0
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 0
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 12
-runtime:
-  dp: 32
-  tp: 1
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_nw1_b32_dp32_PROFILE.yaml b/param_decomp/configs/llama8b_full32L_nw1_b32_dp32_PROFILE.yaml
deleted file mode 100644
index 642470b80..000000000
--- a/param_decomp/configs/llama8b_full32L_nw1_b32_dp32_PROFILE.yaml
+++ /dev/null
@@ -1,586 +0,0 @@
-run_name: jax-full32L-HSDP-b32-dp32-PROFILE
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 1
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 32
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 32
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 0
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 1
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 12
-runtime:
-  dp: 32
-  tp: 1
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_seq512_b128_dp128_bf16src.yaml b/param_decomp/configs/llama8b_full32L_seq512_b128_dp128_bf16src.yaml
deleted file mode 100644
index 74e58d617..000000000
--- a/param_decomp/configs/llama8b_full32L_seq512_b128_dp128_bf16src.yaml
+++ /dev/null
@@ -1,612 +0,0 @@
-# First full-model decomposition of Llama-3.1-8B: ALL 32 layers x 7 matrices
-# (q/k/v/o/gate/up/down) = 224 sites. De-risking launch — get it compiling, fitting
-# memory at the 128-GPU topology, and producing a sane training signal; then dial in
-# GPU count / chunking / LRs.
-#
-# Derived from the l18-23 6L reference run p-2112148d ("jax-l18-23-6L-ablbase-200k"):
-# same algorithm (chunkwise subset recon + persistent-PGD + faith + imp-min), same CI
-# fn arch, same LRs/coeffs/eval set. DELIBERATE EDITS vs that reference:
-#   1. sites: 18 (MLP, layers 18-23) -> 224 (all matrices, all 32 layers). Per-matrix C
-#      from the 2026-06-22 full-llama8b BOTEC, rounded UP to multiples of 128 (the V/U
-#      C-axis shards over the dp=128 mesh; init asserts C % dp == 0):
-#      q/k 2048, v/o 4096, gate/up 8192, down 10240. ΣC=1,245,184; V/U=18.3B params
-#      (~1.1x the reference's 16.3B — params shard cheaply, not the binding constraint).
-#   2. runtime.dp: 64 -> 128 (16 nodes). batch 128 -> per-rank 1 (vs 2 in the reference);
-#      the botec puts per-rank-1 at ~105 GiB, safe under the 180 GiB B200 cap.
-#   3. ChunkwiseSubsetReconLoss sites_per_chunk: 9 -> 56 (8 layers/chunk -> 4 chunks).
-#      Fewer, larger chunks for the first launch: minimises compile time and per-step
-#      compute (each chunk's suffix forward now spans the full network). coeff scaled
-#      0.5 x n_chunks = 2.0 to keep per-site recon pressure chunk-count-invariant.
-#      THIS IS THE KNOB TO RE-TUNE once we know step time (finer chunks = better signal).
-#   4. remat_recon_forwards: true (224 sites' full-network chunk forwards need it).
-#   5. remat_ci_fn: true (checkpoint the ~31B CI fn; without it the step is the GPU
-#      compile wall — ~80 min + near-OOM — since the whole CI forward materialises).
-#   6. data seq 512 (fineweb_llama_tok_512), matching the reference's seq.
-# out_dir / run_id are minted by pd-lm.
-run_name: jax-full32L-allmat-seq512-b128-dp128-bf16src
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 25
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 128
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 128
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 400
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    source_dtype: bfloat16
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 200000
-runtime:
-  dp: 128
-  tp: 8
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_seq512_b128_dp64.yaml b/param_decomp/configs/llama8b_full32L_seq512_b128_dp64.yaml
deleted file mode 100644
index c9dbdbda6..000000000
--- a/param_decomp/configs/llama8b_full32L_seq512_b128_dp64.yaml
+++ /dev/null
@@ -1,611 +0,0 @@
-# First full-model decomposition of Llama-3.1-8B: ALL 32 layers x 7 matrices
-# (q/k/v/o/gate/up/down) = 224 sites. De-risking launch — get it compiling, fitting
-# memory at the 128-GPU topology, and producing a sane training signal; then dial in
-# GPU count / chunking / LRs.
-#
-# Derived from the l18-23 6L reference run p-2112148d ("jax-l18-23-6L-ablbase-200k"):
-# same algorithm (chunkwise subset recon + persistent-PGD + faith + imp-min), same CI
-# fn arch, same LRs/coeffs/eval set. DELIBERATE EDITS vs that reference:
-#   1. sites: 18 (MLP, layers 18-23) -> 224 (all matrices, all 32 layers). Per-matrix C
-#      from the 2026-06-22 full-llama8b BOTEC, rounded UP to multiples of 128 (the V/U
-#      C-axis shards over the dp=128 mesh; init asserts C % dp == 0):
-#      q/k 2048, v/o 4096, gate/up 8192, down 10240. ΣC=1,245,184; V/U=18.3B params
-#      (~1.1x the reference's 16.3B — params shard cheaply, not the binding constraint).
-#   2. runtime.dp: 64 -> 128 (16 nodes). batch 128 -> per-rank 1 (vs 2 in the reference);
-#      the botec puts per-rank-1 at ~105 GiB, safe under the 180 GiB B200 cap.
-#   3. ChunkwiseSubsetReconLoss sites_per_chunk: 9 -> 56 (8 layers/chunk -> 4 chunks).
-#      Fewer, larger chunks for the first launch: minimises compile time and per-step
-#      compute (each chunk's suffix forward now spans the full network). coeff scaled
-#      0.5 x n_chunks = 2.0 to keep per-site recon pressure chunk-count-invariant.
-#      THIS IS THE KNOB TO RE-TUNE once we know step time (finer chunks = better signal).
-#   4. remat_recon_forwards: true (224 sites' full-network chunk forwards need it).
-#   5. remat_ci_fn: true (checkpoint the ~31B CI fn; without it the step is the GPU
-#      compile wall — ~80 min + near-OOM — since the whole CI forward materialises).
-#   6. data seq 512 (fineweb_llama_tok_512), matching the reference's seq.
-# out_dir / run_id are minted by pd-lm.
-run_name: jax-full32L-allmat-seq512-b128-dp64
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 25
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 128
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 128
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 400
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 200000
-runtime:
-  dp: 64
-  tp: 8
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_seq512_b256_dp128.yaml b/param_decomp/configs/llama8b_full32L_seq512_b256_dp128.yaml
deleted file mode 100644
index edb37f325..000000000
--- a/param_decomp/configs/llama8b_full32L_seq512_b256_dp128.yaml
+++ /dev/null
@@ -1,611 +0,0 @@
-# First full-model decomposition of Llama-3.1-8B: ALL 32 layers x 7 matrices
-# (q/k/v/o/gate/up/down) = 224 sites. De-risking launch — get it compiling, fitting
-# memory at the 128-GPU topology, and producing a sane training signal; then dial in
-# GPU count / chunking / LRs.
-#
-# Derived from the l18-23 6L reference run p-2112148d ("jax-l18-23-6L-ablbase-200k"):
-# same algorithm (chunkwise subset recon + persistent-PGD + faith + imp-min), same CI
-# fn arch, same LRs/coeffs/eval set. DELIBERATE EDITS vs that reference:
-#   1. sites: 18 (MLP, layers 18-23) -> 224 (all matrices, all 32 layers). Per-matrix C
-#      from the 2026-06-22 full-llama8b BOTEC, rounded UP to multiples of 128 (the V/U
-#      C-axis shards over the dp=128 mesh; init asserts C % dp == 0):
-#      q/k 2048, v/o 4096, gate/up 8192, down 10240. ΣC=1,245,184; V/U=18.3B params
-#      (~1.1x the reference's 16.3B — params shard cheaply, not the binding constraint).
-#   2. runtime.dp: 64 -> 128 (16 nodes). batch 128 -> per-rank 1 (vs 2 in the reference);
-#      the botec puts per-rank-1 at ~105 GiB, safe under the 180 GiB B200 cap.
-#   3. ChunkwiseSubsetReconLoss sites_per_chunk: 9 -> 56 (8 layers/chunk -> 4 chunks).
-#      Fewer, larger chunks for the first launch: minimises compile time and per-step
-#      compute (each chunk's suffix forward now spans the full network). coeff scaled
-#      0.5 x n_chunks = 2.0 to keep per-site recon pressure chunk-count-invariant.
-#      THIS IS THE KNOB TO RE-TUNE once we know step time (finer chunks = better signal).
-#   4. remat_recon_forwards: true (224 sites' full-network chunk forwards need it).
-#   5. remat_ci_fn: true (checkpoint the ~31B CI fn; without it the step is the GPU
-#      compile wall — ~80 min + near-OOM — since the whole CI forward materialises).
-#   6. data seq 512 (fineweb_llama_tok_512), matching the reference's seq.
-# out_dir / run_id are minted by pd-lm.
-run_name: jax-full32L-allmat-seq512-b256-dp128
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 25
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 256
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 256
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 400
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 200000
-runtime:
-  dp: 128
-  tp: 8
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_seq512_b32_dp128_tp4.yaml b/param_decomp/configs/llama8b_full32L_seq512_b32_dp128_tp4.yaml
deleted file mode 100644
index 2a50a50df..000000000
--- a/param_decomp/configs/llama8b_full32L_seq512_b32_dp128_tp4.yaml
+++ /dev/null
@@ -1,613 +0,0 @@
-# First full-model decomposition of Llama-3.1-8B: ALL 32 layers x 7 matrices
-# (q/k/v/o/gate/up/down) = 224 sites. De-risking launch — get it compiling, fitting
-# memory at the 128-GPU topology, and producing a sane training signal; then dial in
-# GPU count / chunking / LRs.
-#
-# Derived from the l18-23 6L reference run p-2112148d ("jax-l18-23-6L-ablbase-200k"):
-# same algorithm (chunkwise subset recon + persistent-PGD + faith + imp-min), same CI
-# fn arch, same LRs/coeffs/eval set. DELIBERATE EDITS vs that reference:
-#   1. sites: 18 (MLP, layers 18-23) -> 224 (all matrices, all 32 layers). Per-matrix C
-#      from the 2026-06-22 full-llama8b BOTEC, rounded UP to multiples of 128 (the V/U
-#      C-axis shards on the `tp` mesh axis; init asserts C % tp == 0):
-#      q/k 2048, v/o 4096, gate/up 8192, down 10240. ΣC=1,245,184; V/U=18.3B params
-#      (~1.1x the reference's 16.3B — params shard cheaply, not the binding constraint).
-#   2. HIGH-WATER-MARK topology: runtime.dp=128 (16 nodes, power-of-2), tp=4 → dp_axis=32,
-#      and batch 32 -> per-DP=1 (global_batch / dp_axis = 32/32). This is the lightest
-#      activation footprint (the botec's ~105 GiB / per-DP-1 point). NOTE: b128/tp8 OOM'd
-#      (~72 GiB alloc) because tp=8 -> dp_axis=16 -> per-DP=8, ~8x these activations.
-#   3. ChunkwiseSubsetReconLoss sites_per_chunk: 9 -> 56 (8 layers/chunk -> 4 chunks).
-#      Fewer, larger chunks for the first launch: minimises compile time and per-step
-#      compute (each chunk's suffix forward now spans the full network). coeff scaled
-#      0.5 x n_chunks = 2.0 to keep per-site recon pressure chunk-count-invariant.
-#      THIS IS THE KNOB TO RE-TUNE once we know step time (finer chunks = better signal).
-#   4. remat_recon_forwards: true (224 sites' full-network chunk forwards need it).
-#   5. remat_ci_fn: true (checkpoint the ~31B CI fn; without it the step is the GPU
-#      compile wall — ~80 min + near-OOM — since the whole CI forward materialises).
-#   6. data seq 512 (fineweb_llama_tok_512), matching the reference's seq.
-# out_dir / run_id are minted by pd-lm.
-run_name: jax-full32L-allmat-seq512-b32-dp128-tp4
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 25
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 32
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 32
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 400
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 200000
-runtime:
-  dp: 128
-  tp: 4
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_seq512_b64_dp128_tp2.yaml b/param_decomp/configs/llama8b_full32L_seq512_b64_dp128_tp2.yaml
deleted file mode 100644
index 609da00dc..000000000
--- a/param_decomp/configs/llama8b_full32L_seq512_b64_dp128_tp2.yaml
+++ /dev/null
@@ -1,611 +0,0 @@
-# First full-model decomposition of Llama-3.1-8B: ALL 32 layers x 7 matrices
-# (q/k/v/o/gate/up/down) = 224 sites. De-risking launch — get it compiling, fitting
-# memory at the 128-GPU topology, and producing a sane training signal; then dial in
-# GPU count / chunking / LRs.
-#
-# Derived from the l18-23 6L reference run p-2112148d ("jax-l18-23-6L-ablbase-200k"):
-# same algorithm (chunkwise subset recon + persistent-PGD + faith + imp-min), same CI
-# fn arch, same LRs/coeffs/eval set. DELIBERATE EDITS vs that reference:
-#   1. sites: 18 (MLP, layers 18-23) -> 224 (all matrices, all 32 layers). Per-matrix C
-#      from the 2026-06-22 full-llama8b BOTEC, rounded UP to multiples of 128 (the V/U
-#      C-axis shards over the dp=128 mesh; init asserts C % dp == 0):
-#      q/k 2048, v/o 4096, gate/up 8192, down 10240. ΣC=1,245,184; V/U=18.3B params
-#      (~1.1x the reference's 16.3B — params shard cheaply, not the binding constraint).
-#   2. runtime.dp: 64 -> 128 (16 nodes). batch 128 -> per-rank 1 (vs 2 in the reference);
-#      the botec puts per-rank-1 at ~105 GiB, safe under the 180 GiB B200 cap.
-#   3. ChunkwiseSubsetReconLoss sites_per_chunk: 9 -> 56 (8 layers/chunk -> 4 chunks).
-#      Fewer, larger chunks for the first launch: minimises compile time and per-step
-#      compute (each chunk's suffix forward now spans the full network). coeff scaled
-#      0.5 x n_chunks = 2.0 to keep per-site recon pressure chunk-count-invariant.
-#      THIS IS THE KNOB TO RE-TUNE once we know step time (finer chunks = better signal).
-#   4. remat_recon_forwards: true (224 sites' full-network chunk forwards need it).
-#   5. remat_ci_fn: true (checkpoint the ~31B CI fn; without it the step is the GPU
-#      compile wall — ~80 min + near-OOM — since the whole CI forward materialises).
-#   6. data seq 512 (fineweb_llama_tok_512), matching the reference's seq.
-# out_dir / run_id are minted by pd-lm.
-run_name: jax-full32L-allmat-seq512-b64-dp128-tp2
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 25
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 64
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 64
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 400
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 200000
-runtime:
-  dp: 128
-  tp: 2
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_seq512_b64_dp64_tp1.yaml b/param_decomp/configs/llama8b_full32L_seq512_b64_dp64_tp1.yaml
deleted file mode 100644
index 803e4d007..000000000
--- a/param_decomp/configs/llama8b_full32L_seq512_b64_dp64_tp1.yaml
+++ /dev/null
@@ -1,611 +0,0 @@
-# First full-model decomposition of Llama-3.1-8B: ALL 32 layers x 7 matrices
-# (q/k/v/o/gate/up/down) = 224 sites. De-risking launch — get it compiling, fitting
-# memory at the 128-GPU topology, and producing a sane training signal; then dial in
-# GPU count / chunking / LRs.
-#
-# Derived from the l18-23 6L reference run p-2112148d ("jax-l18-23-6L-ablbase-200k"):
-# same algorithm (chunkwise subset recon + persistent-PGD + faith + imp-min), same CI
-# fn arch, same LRs/coeffs/eval set. DELIBERATE EDITS vs that reference:
-#   1. sites: 18 (MLP, layers 18-23) -> 224 (all matrices, all 32 layers). Per-matrix C
-#      from the 2026-06-22 full-llama8b BOTEC, rounded UP to multiples of 128 (the V/U
-#      C-axis shards over the dp=128 mesh; init asserts C % dp == 0):
-#      q/k 2048, v/o 4096, gate/up 8192, down 10240. ΣC=1,245,184; V/U=18.3B params
-#      (~1.1x the reference's 16.3B — params shard cheaply, not the binding constraint).
-#   2. runtime.dp: 64 -> 128 (16 nodes). batch 128 -> per-rank 1 (vs 2 in the reference);
-#      the botec puts per-rank-1 at ~105 GiB, safe under the 180 GiB B200 cap.
-#   3. ChunkwiseSubsetReconLoss sites_per_chunk: 9 -> 56 (8 layers/chunk -> 4 chunks).
-#      Fewer, larger chunks for the first launch: minimises compile time and per-step
-#      compute (each chunk's suffix forward now spans the full network). coeff scaled
-#      0.5 x n_chunks = 2.0 to keep per-site recon pressure chunk-count-invariant.
-#      THIS IS THE KNOB TO RE-TUNE once we know step time (finer chunks = better signal).
-#   4. remat_recon_forwards: true (224 sites' full-network chunk forwards need it).
-#   5. remat_ci_fn: true (checkpoint the ~31B CI fn; without it the step is the GPU
-#      compile wall — ~80 min + near-OOM — since the whole CI forward materialises).
-#   6. data seq 512 (fineweb_llama_tok_512), matching the reference's seq.
-# out_dir / run_id are minted by pd-lm.
-run_name: jax-full32L-allmat-seq512-b64-dp64-tp1
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 25
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 64
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 64
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 400
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 200000
-runtime:
-  dp: 64
-  tp: 1
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_seq512_b64_dp64_tp2.yaml b/param_decomp/configs/llama8b_full32L_seq512_b64_dp64_tp2.yaml
deleted file mode 100644
index 27cf150e9..000000000
--- a/param_decomp/configs/llama8b_full32L_seq512_b64_dp64_tp2.yaml
+++ /dev/null
@@ -1,611 +0,0 @@
-# First full-model decomposition of Llama-3.1-8B: ALL 32 layers x 7 matrices
-# (q/k/v/o/gate/up/down) = 224 sites. De-risking launch — get it compiling, fitting
-# memory at the 128-GPU topology, and producing a sane training signal; then dial in
-# GPU count / chunking / LRs.
-#
-# Derived from the l18-23 6L reference run p-2112148d ("jax-l18-23-6L-ablbase-200k"):
-# same algorithm (chunkwise subset recon + persistent-PGD + faith + imp-min), same CI
-# fn arch, same LRs/coeffs/eval set. DELIBERATE EDITS vs that reference:
-#   1. sites: 18 (MLP, layers 18-23) -> 224 (all matrices, all 32 layers). Per-matrix C
-#      from the 2026-06-22 full-llama8b BOTEC, rounded UP to multiples of 128 (the V/U
-#      C-axis shards over the dp=128 mesh; init asserts C % dp == 0):
-#      q/k 2048, v/o 4096, gate/up 8192, down 10240. ΣC=1,245,184; V/U=18.3B params
-#      (~1.1x the reference's 16.3B — params shard cheaply, not the binding constraint).
-#   2. runtime.dp: 64 -> 128 (16 nodes). batch 128 -> per-rank 1 (vs 2 in the reference);
-#      the botec puts per-rank-1 at ~105 GiB, safe under the 180 GiB B200 cap.
-#   3. ChunkwiseSubsetReconLoss sites_per_chunk: 9 -> 56 (8 layers/chunk -> 4 chunks).
-#      Fewer, larger chunks for the first launch: minimises compile time and per-step
-#      compute (each chunk's suffix forward now spans the full network). coeff scaled
-#      0.5 x n_chunks = 2.0 to keep per-site recon pressure chunk-count-invariant.
-#      THIS IS THE KNOB TO RE-TUNE once we know step time (finer chunks = better signal).
-#   4. remat_recon_forwards: true (224 sites' full-network chunk forwards need it).
-#   5. remat_ci_fn: true (checkpoint the ~31B CI fn; without it the step is the GPU
-#      compile wall — ~80 min + near-OOM — since the whole CI forward materialises).
-#   6. data seq 512 (fineweb_llama_tok_512), matching the reference's seq.
-# out_dir / run_id are minted by pd-lm.
-run_name: jax-full32L-allmat-seq512-b64-dp64-tp2
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 25
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 64
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 64
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 400
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 200000
-runtime:
-  dp: 64
-  tp: 2
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_seq512_b64_dp64_tp4.yaml b/param_decomp/configs/llama8b_full32L_seq512_b64_dp64_tp4.yaml
deleted file mode 100644
index c62a7413c..000000000
--- a/param_decomp/configs/llama8b_full32L_seq512_b64_dp64_tp4.yaml
+++ /dev/null
@@ -1,611 +0,0 @@
-# First full-model decomposition of Llama-3.1-8B: ALL 32 layers x 7 matrices
-# (q/k/v/o/gate/up/down) = 224 sites. De-risking launch — get it compiling, fitting
-# memory at the 128-GPU topology, and producing a sane training signal; then dial in
-# GPU count / chunking / LRs.
-#
-# Derived from the l18-23 6L reference run p-2112148d ("jax-l18-23-6L-ablbase-200k"):
-# same algorithm (chunkwise subset recon + persistent-PGD + faith + imp-min), same CI
-# fn arch, same LRs/coeffs/eval set. DELIBERATE EDITS vs that reference:
-#   1. sites: 18 (MLP, layers 18-23) -> 224 (all matrices, all 32 layers). Per-matrix C
-#      from the 2026-06-22 full-llama8b BOTEC, rounded UP to multiples of 128 (the V/U
-#      C-axis shards over the dp=128 mesh; init asserts C % dp == 0):
-#      q/k 2048, v/o 4096, gate/up 8192, down 10240. ΣC=1,245,184; V/U=18.3B params
-#      (~1.1x the reference's 16.3B — params shard cheaply, not the binding constraint).
-#   2. runtime.dp: 64 -> 128 (16 nodes). batch 128 -> per-rank 1 (vs 2 in the reference);
-#      the botec puts per-rank-1 at ~105 GiB, safe under the 180 GiB B200 cap.
-#   3. ChunkwiseSubsetReconLoss sites_per_chunk: 9 -> 56 (8 layers/chunk -> 4 chunks).
-#      Fewer, larger chunks for the first launch: minimises compile time and per-step
-#      compute (each chunk's suffix forward now spans the full network). coeff scaled
-#      0.5 x n_chunks = 2.0 to keep per-site recon pressure chunk-count-invariant.
-#      THIS IS THE KNOB TO RE-TUNE once we know step time (finer chunks = better signal).
-#   4. remat_recon_forwards: true (224 sites' full-network chunk forwards need it).
-#   5. remat_ci_fn: true (checkpoint the ~31B CI fn; without it the step is the GPU
-#      compile wall — ~80 min + near-OOM — since the whole CI forward materialises).
-#   6. data seq 512 (fineweb_llama_tok_512), matching the reference's seq.
-# out_dir / run_id are minted by pd-lm.
-run_name: jax-full32L-allmat-seq512-b64-dp64-tp4
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 25
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 64
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 64
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 400
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 200000
-runtime:
-  dp: 64
-  tp: 4
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_full32L_seq512_b64_dp64_tp8.yaml b/param_decomp/configs/llama8b_full32L_seq512_b64_dp64_tp8.yaml
deleted file mode 100644
index e9ea2a0c3..000000000
--- a/param_decomp/configs/llama8b_full32L_seq512_b64_dp64_tp8.yaml
+++ /dev/null
@@ -1,611 +0,0 @@
-# First full-model decomposition of Llama-3.1-8B: ALL 32 layers x 7 matrices
-# (q/k/v/o/gate/up/down) = 224 sites. De-risking launch — get it compiling, fitting
-# memory at the 128-GPU topology, and producing a sane training signal; then dial in
-# GPU count / chunking / LRs.
-#
-# Derived from the l18-23 6L reference run p-2112148d ("jax-l18-23-6L-ablbase-200k"):
-# same algorithm (chunkwise subset recon + persistent-PGD + faith + imp-min), same CI
-# fn arch, same LRs/coeffs/eval set. DELIBERATE EDITS vs that reference:
-#   1. sites: 18 (MLP, layers 18-23) -> 224 (all matrices, all 32 layers). Per-matrix C
-#      from the 2026-06-22 full-llama8b BOTEC, rounded UP to multiples of 128 (the V/U
-#      C-axis shards over the dp=128 mesh; init asserts C % dp == 0):
-#      q/k 2048, v/o 4096, gate/up 8192, down 10240. ΣC=1,245,184; V/U=18.3B params
-#      (~1.1x the reference's 16.3B — params shard cheaply, not the binding constraint).
-#   2. runtime.dp: 64 -> 128 (16 nodes). batch 128 -> per-rank 1 (vs 2 in the reference);
-#      the botec puts per-rank-1 at ~105 GiB, safe under the 180 GiB B200 cap.
-#   3. ChunkwiseSubsetReconLoss sites_per_chunk: 9 -> 56 (8 layers/chunk -> 4 chunks).
-#      Fewer, larger chunks for the first launch: minimises compile time and per-step
-#      compute (each chunk's suffix forward now spans the full network). coeff scaled
-#      0.5 x n_chunks = 2.0 to keep per-site recon pressure chunk-count-invariant.
-#      THIS IS THE KNOB TO RE-TUNE once we know step time (finer chunks = better signal).
-#   4. remat_recon_forwards: true (224 sites' full-network chunk forwards need it).
-#   5. remat_ci_fn: true (checkpoint the ~31B CI fn; without it the step is the GPU
-#      compile wall — ~80 min + near-OOM — since the whole CI forward materialises).
-#   6. data seq 512 (fineweb_llama_tok_512), matching the reference's seq.
-# out_dir / run_id are minted by pd-lm.
-run_name: jax-full32L-allmat-seq512-b64-dp64-tp8
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 25
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 64
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 64
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.0.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.0.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.0.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.0.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.1.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.1.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.1.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.1.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.2.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.2.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.2.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.2.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.3.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.3.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.3.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.3.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.4.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.4.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.4.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.4.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.5.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.5.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.5.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.5.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.6.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.6.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.6.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.6.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.7.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.7.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.7.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.7.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.8.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.8.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.8.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.8.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.9.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.9.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.9.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.9.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.10.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.10.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.10.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.10.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.11.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.11.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.11.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.11.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.12.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.12.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.12.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.12.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.13.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.13.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.13.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.13.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.14.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.14.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.14.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.14.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.15.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.15.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.15.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.15.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.16.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.16.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.16.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.16.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.17.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.17.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.17.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.17.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.18.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.19.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.19.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.19.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.19.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.20.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.20.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.20.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.20.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.21.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.21.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.21.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.21.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.22.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.22.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.22.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.22.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.23.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.23.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.23.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.23.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.24.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.24.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.24.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.24.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.25.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.25.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.25.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.25.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.26.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.26.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.26.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.26.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.27.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.27.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.27.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.27.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.28.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.28.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.28.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.28.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.29.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.29.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.29.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.29.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.30.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.30.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.30.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.30.mlp.down_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.q_proj
-  - C: 2048
-    module_pattern: model.layers.31.self_attn.k_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.v_proj
-  - C: 4096
-    module_pattern: model.layers.31.self_attn.o_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.gate_proj
-  - C: 8192
-    module_pattern: model.layers.31.mlp.up_proj
-  - C: 10240
-    module_pattern: model.layers.31.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 400
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 2.0
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 56
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 200000
-runtime:
-  dp: 64
-  tp: 8
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: full32L
-  tags:
-  - full-model
-  - 32L
-  - allmat
-  - dp128
-  - seq512
diff --git a/param_decomp/configs/llama8b_l18only_doubleC_b128_dp32.yaml b/param_decomp/configs/llama8b_l18only_doubleC_b128_dp32.yaml
deleted file mode 100644
index f5c90c915..000000000
--- a/param_decomp/configs/llama8b_l18only_doubleC_b128_dp32.yaml
+++ /dev/null
@@ -1,155 +0,0 @@
-# Layer-18-only decomposition of Llama-3.1-8B: JUST layer 18's 7 matrices
-# (q/k/v/o/gate/up/down), every per-matrix C DOUBLED vs the full-model config.
-# Trimmed from llama8b_full32L_HSDP_b128_dp32.yaml (224 sites -> 7); single chunk.
-# out_dir / run_id are minted by pd-lm.
-run_name: jax-l18only-doubleC-b128-dp32
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 1000
-  train_log_every: 25
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/fineweb_llama_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  train_split: train
-eval:
-  batch_size: 32
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups: null
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 100000
-  slow_on_first_step: false
-pd:
-  batch_size: 128
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 4096
-    n_blocks: 4
-    n_heads: 64
-    mlp_hidden: 16384
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.83e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 2.83e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.q_proj
-  - C: 4096
-    module_pattern: model.layers.18.self_attn.k_proj
-  - C: 8192
-    module_pattern: model.layers.18.self_attn.v_proj
-  - C: 8192
-    module_pattern: model.layers.18.self_attn.o_proj
-  - C: 16384
-    module_pattern: model.layers.18.mlp.gate_proj
-  - C: 16384
-    module_pattern: model.layers.18.mlp.up_proj
-  - C: 20480
-    module_pattern: model.layers.18.mlp.down_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 0
-  faithfulness_warmup_weight_decay: 0.0
-  identity_decomposition_targets: null
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-06
-      reference_token_count: 65536
-    coeff: 5.0e-06
-    eps: 1.0e-06
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 0.5
-    n_samples: 1
-    routing:
-      type: uniform_k_subset
-    sites_per_chunk: 7
-    type: ChunkwiseSubsetReconLoss
-  - coeff: 0.5
-    n_warmup_steps: 2
-    optimizer:
-      beta1: 0.01
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        final_val_frac: 1.0
-        fn_type: constant
-        start_val: 0.01
-        warmup_pct: 0.025
-      type: adam
-    scope:
-      type: bsc
-    type: PersistentPGDReconLoss
-  - coeff: 1000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 50000
-runtime:
-  dp: 32
-  tp: 1
-  remat_recon_forwards: true
-  remat_ci_fn: true
-target:
-  output_extract: logits
-  weights_dtype: bfloat16
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-wandb:
-  entity: null
-  project: param-decomp-llama
-  group: l18only
-  tags:
-  - l18
-  - single-layer
-  - doubleC
-  - dp32
diff --git a/param_decomp/configs/pile_pgd1.yaml b/param_decomp/configs/pile_pgd1.yaml
deleted file mode 100644
index 4c003d1fa..000000000
--- a/param_decomp/configs/pile_pgd1.yaml
+++ /dev/null
@@ -1,158 +0,0 @@
-# Fresh-PGD pile run — JAX counterpart of Dan's torch p-af354eb1 (pgd1step-ss1.0)
-# via the shared-config route. See the referenced yaml's header for the documented
-# edits. Small target (4L d768): no remat needed. Launch: pd-lm, 1 node.
-run_name: jax-pile4l-pgd1-ss1
-# JAX leg of Dan's fresh-PGD pile experiment — byte-derived from the stored config of
-# torch run p-af354eb1 ("pgd1step-ss1.0", the relaunch of p-a005ed60; snapshot
-# runs/p-af354eb1/experiment_config.yaml). Target: LlamaSimpleMLP 4L pile model
-# (t-9d2b8f02), all h.* sites, per-site C, PGDRecon n_steps=1 step_size=1
-# unique_per_datapoint (-> bsc). DOCUMENTED EDITS vs upstream:
-#   1. data: streaming HF danbraunai/pile-uncopyrighted-tok-shuffled -> the staged
-#      parquet artifact (datasets/pile_neox_tok_512; same dataset, same seq 512 —
-#      the JAX trainer reads pre-tokenized parquet shards only).
-#   2. data.eval_split: val -> train (parquet loading exposes a single "train"
-#      split; the staged _val dir exists for ad-hoc use but the loaders don't
-#      address it).
-#   3. cadence.save_every: null -> 5000, keep_last null -> 2 (upstream saves
-#      nothing; the JAX leg checkpoints for requeue-resume + offline eval).
-# Everything else (pd losses/optimizers/CI fn, 400k steps, batch 64, eval set) is
-# upstream-identical.
-cadence:
-  keep_last_n_checkpoints: 2
-  save_every: 5000
-  train_log_every: 200
-data:
-  buffer_size: 1000
-  column_name: input_ids
-  data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet
-  dataset_name: parquet
-  eval_split: train
-  is_tokenized: true
-  max_seq_len: 512
-  revision: null
-  shuffle_each_epoch: true
-  streaming: false
-  tokenizer_name: EleutherAI/gpt-neox-20b
-  train_split: train
-eval:
-  batch_size: 128
-  every: 1000
-  metrics:
-  - n_batches_accum: 1
-    type: CIHistograms
-  - ci_alive_threshold: 0.0
-    type: ComponentActivationDensity
-  - ci_alive_threshold: 0.0
-    groups:
-      layer_0:
-      - h.0.*
-      layer_1:
-      - h.1.*
-      layer_2:
-      - h.2.*
-      layer_3:
-      - h.3.*
-      total:
-      - '*'
-    type: CI_L0
-  - rounding_threshold: 0.0
-    type: CEandKLLosses
-  - type: CIMeanPerComponent
-  - coeff: null
-    type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - coeff: null
-    init: random
-    mask_scope: c
-    n_steps: 20
-    name: PGDReconLoss_20step
-    step_size: 0.1
-    type: PGDReconLoss
-  n_steps: 1
-  slow_every: 10000
-  slow_on_first_step: true
-pd:
-  batch_size: 64
-  ci_config:
-    type: chunkwise_transformer
-    blocks_per_chunk: 1
-    d_model: 2048
-    n_blocks: 8
-    n_heads: 16
-    mlp_hidden: 8192
-  ci_fn_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: null
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 5.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  components_optimizer:
-    betas:
-    - 0.9
-    - 0.999
-    grad_clip_norm: 0.01
-    lr_schedule:
-      final_val_frac: 0.1
-      fn_type: cosine
-      start_val: 5.0e-05
-      warmup_pct: 0.0
-    weight_decay: 0.0
-  decomposition_targets:
-  - C: 3072
-    module_pattern: h.*.mlp.c_fc
-  - C: 3584
-    module_pattern: h.*.mlp.down_proj
-  - C: 512
-    module_pattern: h.*.attn.q_proj
-  - C: 512
-    module_pattern: h.*.attn.k_proj
-  - C: 1024
-    module_pattern: h.*.attn.v_proj
-  - C: 1024
-    module_pattern: h.*.attn.o_proj
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_steps: 400
-  faithfulness_warmup_weight_decay: 0.0
-  loss_metrics:
-  - frequency:
-      coeff: 1.0e-04
-      reference_token_count: 65536
-    coeff: 0.0002
-    eps: 1.0e-12
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    type: ImportanceMinimalityLoss
-  - coeff: 0.5
-    routing:
-      type: uniform_k_subset
-    type: StochasticReconSubsetLoss
-  - coeff: 0.5
-    init: random
-    mask_scope: bsc
-    n_steps: 1
-    step_size: 1.0
-    type: PGDReconLoss
-  - coeff: 10000000.0
-    type: FaithfulnessLoss
-  seed: 0
-  steps: 400000
-runtime:
-  remat_recon_forwards: false
-  dp: 16
-target:
-  output_extract: 0
-  weights_dtype: bfloat16
-  spec:
-    kind: pretrained
-    model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP
-    run_path: goodfire/spd/runs/t-9d2b8f02
-wandb:
-  entity: null
-  project: param-decomp
diff --git a/param_decomp_lab/experiments/lm/gpt2-xl.yaml b/param_decomp_lab/experiments/lm/gpt2-xl.yaml
deleted file mode 100644
index b1c089ff3..000000000
--- a/param_decomp_lab/experiments/lm/gpt2-xl.yaml
+++ /dev/null
@@ -1,102 +0,0 @@
-pd:
-  seed: 0
-  ci_config:
-    mode: global
-    fn_type: global_shared_transformer
-    hidden_dims: null
-    simple_transformer_ci_cfg:
-      d_model: 1024
-      n_blocks: 5
-      mlp_hidden_dim:
-      - 4096
-      attn_config:
-        n_heads: 8
-        max_len: 512
-        rope_base: 10000.0
-  decomposition_targets:
-  - module_pattern: transformer.h.*.mlp.c_fc
-    C: 8192
-  - module_pattern: transformer.h.*.mlp.c_proj
-    C: 8192
-  - module_pattern: transformer.h.*.attn.c_attn
-    C: 4096
-  - module_pattern: transformer.h.*.attn.c_proj
-    C: 2048
-  components_optimizer:
-    lr_schedule:
-      start_val: 5.0e-05
-      warmup_pct: 0.0
-      final_val_frac: 0.1
-      fn_type: cosine
-    grad_clip_norm: 0.01
-  ci_fn_optimizer:
-    lr_schedule:
-      start_val: 5.0e-05
-      warmup_pct: 0.0
-      final_val_frac: 0.1
-      fn_type: cosine
-  steps: 100000
-  batch_size: 64
-  faithfulness_warmup_steps: 400
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_weight_decay: 0.0
-  loss_metrics:
-  - type: ImportanceMinimalityLoss
-    coeff: 0.00001
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    frequency:
-      coeff: 5.0e-06
-      reference_token_count: 32768
-    eps: 1.0e-12
-  - type: StochasticReconLayerwiseLoss
-    coeff: 0.5
-  - type: FaithfulnessLoss
-    coeff: 10000000.0
-cadence:
-  train_log_every: 200
-eval:
-  n_steps: 1
-  batch_size: 128
-  every: 1000
-  slow_every: 10000
-  slow_on_first_step: true
-  metrics:
-  - type: CIHistograms
-    n_batches_accum: 1
-  - type: ComponentActivationDensity
-  - type: CI_L0
-    groups: null
-  - type: CEandKLLosses
-    rounding_threshold: 0.0
-  - type: CIMeanPerComponent
-  - type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - type: PGDReconLoss
-    init: random
-    step_size: 0.1
-    n_steps: 20
-    mask_scope: c
-target:
-  spec:
-    kind: hf
-    model_class: transformers.GPT2LMHeadModel
-    model_name: openai-community/gpt2-xl
-  output_extract: logits
-data:
-  tokenizer_name: openai-community/gpt2
-  max_seq_len: 512
-  buffer_size: 1000
-  dataset_name: apollo-research/Skylion007-openwebtext-tokenizer-gpt2
-  column_name: input_ids
-  train_split: train
-  eval_split: train
-  shuffle_each_epoch: true
-  is_tokenized: true
-  streaming: true
-runtime:
-  dp: null
-wandb:
-  project: param-decomp
diff --git a/param_decomp_lab/experiments/lm/gpt2_config.yaml b/param_decomp_lab/experiments/lm/gpt2_config.yaml
deleted file mode 100644
index 23b6af168..000000000
--- a/param_decomp_lab/experiments/lm/gpt2_config.yaml
+++ /dev/null
@@ -1,72 +0,0 @@
-pd:
-  seed: 0
-  ci_config:
-    mode: layerwise
-    fn_type: vector_mlp
-    hidden_dims:
-    - 12
-  decomposition_targets:
-  - module_pattern: transformer.h.1.attn.c_attn
-    C: 768
-  batch_size: 2
-  steps: 50000
-  components_optimizer:
-    lr_schedule:
-      start_val: 3e-4
-      fn_type: cosine
-      warmup_pct: 0.01
-      final_val_frac: 0.1
-  ci_fn_optimizer:
-    lr_schedule:
-      start_val: 3e-4
-      fn_type: cosine
-      warmup_pct: 0.01
-      final_val_frac: 0.1
-  faithfulness_warmup_steps: 200
-  faithfulness_warmup_lr: 0.01
-  faithfulness_warmup_weight_decay: 0.1
-  loss_metrics:
-    - type: ImportanceMinimalityLoss
-      coeff: 0.0003
-      pnorm:
-        start_val: 2.0
-        fn_type: linear
-        final_val_frac: 0.15
-    - type: StochasticReconLayerwiseLoss
-      coeff: 2.0
-cadence:
-  train_log_every: 100
-eval:
-  n_steps: 5
-  batch_size: 2
-  every: 1000
-  slow_every: 5000
-  slow_on_first_step: true
-  metrics:
-    - type: CIHistograms
-      n_batches_accum: 5
-    - type: ComponentActivationDensity
-    - type: CI_L0
-      groups: null
-    - type: CEandKLLosses
-      rounding_threshold: 0.0
-    - type: CIMeanPerComponent
-    - type: StochasticHiddenActsReconLoss
-target:
-  spec:
-    kind: hf
-    model_class: transformers.GPT2LMHeadModel
-    model_name: openai-community/gpt2
-  output_extract: logits
-data:
-  tokenizer_name: openai-community/gpt2
-  max_seq_len: 1024
-  dataset_name: apollo-research/Skylion007-openwebtext-tokenizer-gpt2
-  is_tokenized: true
-  column_name: input_ids
-  train_split: train
-  eval_split: train
-runtime:
-  dp: null
-wandb:
-  project: param-decomp
diff --git a/param_decomp_lab/experiments/lm/llama-3.1-8b.yaml b/param_decomp_lab/experiments/lm/llama-3.1-8b.yaml
deleted file mode 100644
index 961d02c82..000000000
--- a/param_decomp_lab/experiments/lm/llama-3.1-8b.yaml
+++ /dev/null
@@ -1,128 +0,0 @@
-pd:
-  seed: 0
-  ci_config:
-    mode: global
-    fn_type: global_shared_transformer
-    hidden_dims: null
-    simple_transformer_ci_cfg:
-      d_model: 4096
-      n_blocks: 8
-      mlp_hidden_dim:
-      - 16384
-      attn_config:
-        n_heads: 8
-        max_len: 512
-        rope_base: 10000.0
-  decomposition_targets:
-  - module_pattern: model.layers.18.mlp.gate_proj
-    C: 20000
-  - module_pattern: model.layers.18.mlp.up_proj
-    C: 20000
-  - module_pattern: model.layers.18.mlp.down_proj
-    C: 20000 # There are 14336 neurons
-  # - module_pattern: model.layers.*.self_attn.q_proj
-  #   C: 4096
-  # - module_pattern: model.layers.*.self_attn.k_proj
-  #   C: 4096
-  # - module_pattern: model.layers.*.self_attn.v_proj
-  #   C: 4096
-  # - module_pattern: model.layers.*.self_attn.o_proj
-  #   C: 2048
-  components_optimizer:
-    lr_schedule:
-      start_val: 5.0e-05
-      warmup_pct: 0.0
-      final_val_frac: 0.1
-      fn_type: cosine
-    grad_clip_norm: 0.01
-  ci_fn_optimizer:
-    lr_schedule:
-      start_val: 5.0e-05
-      warmup_pct: 0.0
-      final_val_frac: 0.1
-      fn_type: cosine
-  steps: 100000
-  batch_size: 64
-  faithfulness_warmup_steps: 400
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_weight_decay: 0.0
-  loss_metrics:
-  - type: ImportanceMinimalityLoss
-    coeff: 0.000001
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    frequency:
-      coeff: 5.0e-07
-      reference_token_count: 32768
-    eps: 1.0e-12
-  - type: StochasticReconSubsetLoss
-    coeff: 0.5
-    routing:
-      type: uniform_k_subset
-  - type: PersistentPGDReconLoss
-    coeff: 0.5
-    optimizer:
-      type: adam
-      beta1: 0.5
-      beta2: 0.99
-      eps: 1.0e-08
-      lr_schedule:
-        start_val: 0.01
-        warmup_pct: 0.025
-        final_val_frac: 1.0
-        fn_type: constant
-    scope:
-      type: bsc
-    n_warmup_steps: 2
-  - type: FaithfulnessLoss
-    coeff: 10000000.0
-cadence:
-  train_log_every: 200
-eval:
-  n_steps: 1
-  batch_size: 128
-  every: 1000
-  slow_every: 10000
-  slow_on_first_step: true
-  metrics:
-  - type: CIHistograms
-    n_batches_accum: 1
-  - type: ComponentActivationDensity
-  - type: CI_L0
-    groups: null
-  - type: CEandKLLosses
-    rounding_threshold: 0.0
-  - type: CIMeanPerComponent
-  - type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - type: PGDReconLoss
-    init: random
-    step_size: 0.1
-    n_steps: 20
-    mask_scope: c
-target:
-  spec:
-    kind: hf
-    model_class: transformers.LlamaForCausalLM
-    model_name: meta-llama/Llama-3.1-8B
-  output_extract: logits
-data:
-  tokenizer_name: meta-llama/Llama-3.1-8B
-  max_seq_len: 512
-  buffer_size: 1000
-  dataset_name: HuggingFaceFW/fineweb
-  # the sample-350BT subset, globbed directly to avoid a 28k-file repo tree walk
-  data_files: sample/350BT/*.parquet
-  revision: 9bb295ddab0e05d785b879661af7260fed5140fc
-  column_name: text
-  train_split: train
-  eval_split: train
-  shuffle_each_epoch: true
-  is_tokenized: false
-  streaming: true
-runtime:
-  dp: null
-wandb:
-  project: param-decomp-llama
diff --git a/param_decomp_lab/experiments/lm/pile_llama_simple_mlp-12L.yaml b/param_decomp_lab/experiments/lm/pile_llama_simple_mlp-12L.yaml
deleted file mode 100644
index cd5898a6e..000000000
--- a/param_decomp_lab/experiments/lm/pile_llama_simple_mlp-12L.yaml
+++ /dev/null
@@ -1,149 +0,0 @@
-pd:
-  seed: 0
-  ci_config:
-    mode: global
-    fn_type: global_shared_transformer
-    hidden_dims: null
-    simple_transformer_ci_cfg:
-      d_model: 2048
-      n_blocks: 8
-      mlp_hidden_dim:
-      - 8192
-      attn_config:
-        n_heads: 16
-        max_len: 512
-        rope_base: 10000.0
-  decomposition_targets:
-  - module_pattern: h.*.mlp.c_fc
-    C: 3072
-  - module_pattern: h.*.mlp.down_proj
-    C: 3584
-  - module_pattern: h.*.attn.q_proj
-    C: 512
-  - module_pattern: h.*.attn.k_proj
-    C: 512
-  - module_pattern: h.*.attn.v_proj
-    C: 1024
-  - module_pattern: h.*.attn.o_proj
-    C: 1024
-  components_optimizer:
-    lr_schedule:
-      start_val: 5.0e-05
-      warmup_pct: 0.0
-      final_val_frac: 0.1
-      fn_type: cosine
-    grad_clip_norm: 0.01
-  ci_fn_optimizer:
-    lr_schedule:
-      start_val: 5.0e-05
-      warmup_pct: 0.0
-      final_val_frac: 0.1
-      fn_type: cosine
-  steps: 400000
-  batch_size: 64
-  faithfulness_warmup_steps: 400
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_weight_decay: 0.0
-  loss_metrics:
-    - type: ImportanceMinimalityLoss
-      coeff: 0.0001
-      pnorm:
-        start_val: 2.0
-        fn_type: linear
-        final_val_frac: 0.2
-      frequency:
-        coeff: 5.0e-05
-        reference_token_count: 32768
-      eps: 1.0e-12
-    - type: StochasticReconSubsetLoss
-      coeff: 0.5
-      routing:
-        type: uniform_k_subset
-    - type: PersistentPGDReconLoss
-      coeff: 0.5
-      optimizer:
-        type: adam
-        beta1: 0.5
-        beta2: 0.99
-        eps: 1.0e-08
-        lr_schedule:
-          start_val: 0.01
-          warmup_pct: 0.025
-          final_val_frac: 1.0
-          fn_type: constant
-      scope:
-        type: bsc
-      n_warmup_steps: 2
-    - type: FaithfulnessLoss
-      coeff: 10000000.0
-cadence:
-  train_log_every: 200
-eval:
-  n_steps: 1
-  batch_size: 64
-  every: 1000
-  slow_every: 10000
-  slow_on_first_step: true
-  metrics:
-    - type: CIHistograms
-      n_batches_accum: 1
-    - type: ComponentActivationDensity
-    - type: CI_L0
-      groups:
-        layer_0:
-        - h.0.*
-        layer_1:
-        - h.1.*
-        layer_2:
-        - h.2.*
-        layer_3:
-        - h.3.*
-        layer_4:
-        - h.4.*
-        layer_5:
-        - h.5.*
-        layer_6:
-        - h.6.*
-        layer_7:
-        - h.7.*
-        layer_8:
-        - h.8.*
-        layer_9:
-        - h.9.*
-        layer_10:
-        - h.10.*
-        layer_11:
-        - h.11.*
-        total:
-        - '*'
-    - type: CEandKLLosses
-      rounding_threshold: 0.0
-    - type: CIMeanPerComponent
-    - type: StochasticHiddenActsReconLoss
-    - type: CIHiddenActsReconLoss
-    - type: PGDReconLoss
-      init: random
-      step_size: 0.1
-      n_steps: 20
-      mask_scope: c
-target:
-  spec:
-    kind: pretrained
-    model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP
-    run_path: goodfire/spd/runs/t-f99617bb
-  output_extract: 0
-data:
-  tokenizer_name: EleutherAI/gpt-neox-20b
-  max_seq_len: 512
-  buffer_size: 1000
-  dataset_name: danbraunai/pile-uncopyrighted-tok-shuffled
-  column_name: input_ids
-  train_split: train
-  eval_split: val
-  shuffle_each_epoch: true
-  is_tokenized: true
-  streaming: true
-runtime:
-  dp: null
-wandb:
-  project: param-decomp
diff --git a/param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L_layerwise.yaml b/param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L_layerwise.yaml
deleted file mode 100644
index 0e483feef..000000000
--- a/param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L_layerwise.yaml
+++ /dev/null
@@ -1,108 +0,0 @@
-pd:
-  seed: 0
-  ci_config:
-    mode: global
-    fn_type: global_shared_transformer
-    hidden_dims: null
-    simple_transformer_ci_cfg:
-      d_model: 512
-      n_blocks: 5
-      mlp_hidden_dim:
-      - 2048
-      attn_config:
-        n_heads: 8
-        max_len: 512
-        rope_base: 10000.0
-  decomposition_targets:
-  - module_pattern: h.*.mlp.c_fc
-    C: 3072
-  - module_pattern: h.*.mlp.down_proj
-    C: 3584
-  - module_pattern: h.*.attn.q_proj
-    C: 512
-  - module_pattern: h.*.attn.k_proj
-    C: 512
-  - module_pattern: h.*.attn.v_proj
-    C: 1024
-  - module_pattern: h.*.attn.o_proj
-    C: 1024
-  components_optimizer:
-    lr_schedule:
-      start_val: 5.0e-05
-      warmup_pct: 0.0
-      final_val_frac: 0.1
-      fn_type: cosine
-    grad_clip_norm: 0.01
-  ci_fn_optimizer:
-    lr_schedule:
-      start_val: 5.0e-05
-      warmup_pct: 0.0
-      final_val_frac: 0.1
-      fn_type: cosine
-  steps: 100000
-  batch_size: 64
-  faithfulness_warmup_steps: 400
-  faithfulness_warmup_lr: 0.001
-  faithfulness_warmup_weight_decay: 0.0
-  loss_metrics:
-  - type: ImportanceMinimalityLoss
-    coeff: 0.0001
-    pnorm:
-      start_val: 2.0
-      fn_type: linear
-      final_val_frac: 0.2
-    frequency:
-      coeff: 5.0e-05
-      reference_token_count: 32768
-    eps: 1.0e-12
-  - type: StochasticReconSubsetLoss
-    coeff: 0.5
-    routing:
-      type: uniform_k_subset
-  - type: FaithfulnessLoss
-    coeff: 10000000.0
-cadence:
-  train_log_every: 200
-eval:
-  n_steps: 1
-  batch_size: 128
-  every: 1000
-  slow_every: 10000
-  slow_on_first_step: true
-  metrics:
-  - type: CIHistograms
-    n_batches_accum: 1
-  - type: ComponentActivationDensity
-  - type: CI_L0
-    groups: null
-  - type: CEandKLLosses
-    rounding_threshold: 0.0
-  - type: CIMeanPerComponent
-  - type: StochasticHiddenActsReconLoss
-  - type: CIHiddenActsReconLoss
-  - type: PGDReconLoss
-    init: random
-    step_size: 0.1
-    n_steps: 20
-    mask_scope: c
-target:
-  spec:
-    kind: pretrained
-    model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP
-    run_path: goodfire/spd/runs/t-9d2b8f02
-  output_extract: 0
-data:
-  tokenizer_name: EleutherAI/gpt-neox-20b
-  max_seq_len: 512
-  buffer_size: 1000
-  dataset_name: danbraunai/pile-uncopyrighted-tok-shuffled
-  column_name: input_ids
-  train_split: train
-  eval_split: val
-  shuffle_each_epoch: true
-  is_tokenized: true
-  streaming: true
-runtime:
-  dp: null
-wandb:
-  project: param-decomp
diff --git a/param_decomp_lab/tests/test_repo_configs_parse.py b/param_decomp_lab/tests/test_repo_configs_parse.py
new file mode 100644
index 000000000..4504ee320
--- /dev/null
+++ b/param_decomp_lab/tests/test_repo_configs_parse.py
@@ -0,0 +1,53 @@
+"""Every LM config yaml the repo maintains must parse at tip (CONFIGS.md rule 1).
+
+A schema PR that breaks a config here migrates it in the same PR, with an
+executed in-repo migration — never a script attached to a PR comment (#939
+attached one; it never ran, and 97/104 stored runs became unopenable).
+"""
+
+from pathlib import Path
+from typing import Any
+
+import pytest
+import yaml
+from pydantic import ValidationError
+
+from param_decomp_lab.experiments.config import assert_canonical_algorithm_config
+from param_decomp_lab.experiments.lm.config import LMExperimentConfig
+
+REPO = Path(__file__).resolve().parents[2]
+
+# Archetype seats predating the modern schema, awaiting migration (see
+# CONFIGS.md registry). When you fix one, remove it here — the gate then
+# covers it; this list must only ever shrink.
+KNOWN_BROKEN = {
+    "param_decomp_lab/experiments/lm/jose.yaml",
+    "param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml",
+    "param_decomp_lab/experiments/lm/ss_llama_simple_mlp-2L.yaml",
+}
+
+CONFIG_PATHS = sorted(
+    path
+    for pattern in ("param_decomp/configs/**/*.yaml", "param_decomp_lab/experiments/lm/*.yaml")
+    for path in REPO.glob(pattern)
+)
+
+
+def _load(path: Path) -> dict[str, Any]:
+    return yaml.safe_load(path.read_text())
+
+
+@pytest.mark.parametrize("path", CONFIG_PATHS, ids=lambda p: str(p.relative_to(REPO)))
+def test_config_parses_and_is_canonical(path: Path) -> None:
+    rel = str(path.relative_to(REPO))
+    if rel in KNOWN_BROKEN:
+        with pytest.raises(ValidationError):
+            LMExperimentConfig.model_validate(_load(path))
+        pytest.skip(f"{rel}: known-broken seat awaiting migration (CONFIGS.md)")
+    cfg = LMExperimentConfig.model_validate(_load(path))
+    assert_canonical_algorithm_config(cfg)
+
+
+def test_known_broken_entries_still_exist() -> None:
+    missing = [rel for rel in KNOWN_BROKEN if not (REPO / rel).exists()]
+    assert not missing, f"KNOWN_BROKEN lists deleted files, prune it: {missing}"