From bab52b59433341331cdb5132f66e2181f170fc64 Mon Sep 17 00:00:00 2001 From: Dan Braun Date: Mon, 6 Jul 2026 12:05:08 +0100 Subject: [PATCH 1/2] =?UTF-8?q?feat(train):=20stale-CI=20replay=20(S34)=20?= =?UTF-8?q?=E2=80=94=20repeat=20steps=20reuse=20the=20window-first=20CI=20?= =?UTF-8?q?envelope?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Under free-AT batch replay, a window's taps are identical across its steps, and the (~10x-the-target) CI fn fwd+vjp is the dominant step slab at 4L. New config-gated data.replay_stale_ci: repeat steps consume the window-first step's CI envelope as a constant — no taps read, no CI-fn forward/backward, ci_fn + optimizer state pass through untouched (it trains only on window-first steps, matching the 1-per-window unique-data rate). V/U + both adversaries see identical gradients either way (CI live vs constant doesn't change their backward), pinned by test_stale_ci_repeat_step_semantics. - train.py: make_train_step internals -> shared _build_step_impl(cached=None | CIEnvelope); flag-off traces the pre-S34 step unchanged (op order preserved). make_stale_ci_train_steps returns fresh/repeat/compute_ci; repeat's first arg bundles (model, envelope) so all-except-first donation spares both. - ci_fn LR under stale CI: unit-LR optimizer (run_state.unit_lr_ci_fn_optimizer) + in-step scheduled_value_traced at the GLOBAL step — optax's count advances only on window-first steps, so a count-driven cosine would stretch by the replay factor; Adam bias correction correctly tracks the update count. - run.py: do_step dispatches fresh/repeat on step % train_batch_replay; a mid-window requeue-resume rebuilds the envelope via compute_ci (one ci_fn update fresher than the interrupted window's — preemption-only wobble, noted in S34). Grad-norm window stats take the key union (repeats carry no ci_fns family). - configs: staleci 400k arms (replay 2/4/8/16, group-twin the running replay-only sweep) + a replay1 flag-off code baseline from this branch. SPEC S34 added (experimental, off by default; flag-off unamended). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01QKw3iikejknWeF4ZQDFZbL --- param_decomp/CLAUDE.md | 7 + param_decomp/SPEC.md | 1 + param_decomp/built_run.py | 5 + ..._bsc_nw0_lr0p01_replay16_staleci_400k.yaml | 179 ++++++++++ ..._nw0_lr0p01_replay1_staleci_base_400k.yaml | 174 ++++++++++ ...d_bsc_nw0_lr0p01_replay2_staleci_400k.yaml | 179 ++++++++++ ...d_bsc_nw0_lr0p01_replay4_staleci_400k.yaml | 179 ++++++++++ ...d_bsc_nw0_lr0p01_replay8_staleci_400k.yaml | 179 ++++++++++ param_decomp/run.py | 134 ++++++-- param_decomp/run_state.py | 11 + .../tests/equivalence/test_equivalence.py | 2 +- param_decomp/tests/test_llama_simple_mlp.py | 128 ++++++- param_decomp/train.py | 318 +++++++++++++++--- param_decomp_lab/experiments/lm/config.py | 11 + 14 files changed, 1420 insertions(+), 87 deletions(-) create mode 100644 param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay16_staleci_400k.yaml create mode 100644 param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay1_staleci_base_400k.yaml create mode 100644 param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay2_staleci_400k.yaml create mode 100644 param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay4_staleci_400k.yaml create mode 100644 param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay8_staleci_400k.yaml diff --git a/param_decomp/CLAUDE.md b/param_decomp/CLAUDE.md index 7b5a6a180..8c8986b56 100644 --- a/param_decomp/CLAUDE.md +++ b/param_decomp/CLAUDE.md @@ -24,6 +24,13 @@ never silently diverge. Cite IDs (`S14`, `N1`, …) in commit messages and revie ## Architecture in one breath +**Stale-CI replay** (SPEC S34, experimental, `data.replay_stale_ci`): under free-AT batch +replay (`train_batch_replay > 1`), repeat steps reuse the window-first CI envelope as a +constant — no taps / CI-fn forward / CI-fn backward on repeats, ci_fn updates only on +window-first steps (unit-LR optimizer + in-step global-step LR schedule). Built by +`train.make_stale_ci_train_steps` (fresh/repeat/compute_ci), dispatched in +`run.py::do_step`. Off by default; flag-off emits the unchanged step. + `lm.py` defines `DecomposedModel` — a `@runtime_checkable Protocol`: ordered `sites` + `leading_axes` + the methods `clean_output`, `read_activations`, `masked_output`, `masked_site_outputs`, `weight_deltas`, and a `recon_loss_fn` (LM: `kl_per_position`). The diff --git a/param_decomp/SPEC.md b/param_decomp/SPEC.md index 23d0844c9..f9fea75dd 100644 --- a/param_decomp/SPEC.md +++ b/param_decomp/SPEC.md @@ -303,6 +303,7 @@ faithfulness, sources/PPGD, the squashings (S5/S6), the imp-min reduction, the g | S30 | `cfg.cadence.log_every` divides `cfg.eval.every` (`eval.every % log_every == 0`, asserted at `run.py:255`) so every eval step is also a train-log step. | | S31 | The two hidden-acts recon metrics (`CIHiddenActsReconLoss`, `StochasticHiddenActsReconLoss`) are STANDALONE OFFLINE EVAL metrics in JAX (`hidden_acts_eval.py`, wired into `jsp-slow-eval`) — **NOT** recon-grid training terms. `build_loss_terms` still refuses them as training losses (the parameterized recon loss stays KL-on-final-logits only, §2.3–2.5); their objective is per-ELEMENT MSE on each decomposed site's OUTPUT activations, which as a *training* loss is exactly the site-local recon the trainer treats as a conceptual no-no (LOSS_PARITY_DESIGN §4c). The port adds a fifth per-target seam `masked_site_outputs(vu, batch, masks, delta_masks, routes, live, has_delta) -> dict[site, (B,T,d_out)]` (`lm.py`), factored out of `masked_output` (the shared masked forward with a per-site `collect` — the masked per-site output is an intermediate of that forward, so no logic is duplicated). The clean (target) per-site output is the frozen `x @ W`, obtained from the same seam by routing FALSE everywhere (`_site_out`'s frozen branch). Per site, `MSE(masked_site_output, clean_site_output)` with `reduction="sum"` accumulated host-side as `(Σ sum_mse, Σ n_elements)` (token-weighted, exact under micro-batching), divided once at the end; log keys mirror torch exactly (`/` + a combined `` = Σmse/Σn over all sites). `CIHiddenActsReconLoss` is the deterministic `lower_leaky` CI mask, no delta, one forward (tight torch parity); `StochasticHiddenActsReconLoss` draws `n_mask_samples` stochastic CI masks (`mask = ci + (1−ci)·s`) WITH weight deltas (the delta component is always built in JAX runs) — its draws are NOT seed-aligned to torch, so exact bitwise parity is impossible there (expected). Masked + clean run in COMPUTE_DT (bf16, matching the trained model, mirroring `load_run.py`); the MSE reduction is fp32. **AMENDED 2026-06-16** (Oli-approved): superseded the prior "keep-on-bridge / seam refused" decision — these now have a native JAX eval path; the `pd-offline-eval` torch bridge remains available for cross-framework parity checks. | | S33 | Fine-tune init (`ExperimentConfig.resume_provenance`, LM-only): a fresh run whose own `ckpts/` is EMPTY and whose `resume_provenance is not None` initializes from a PARENT checkpoint — it loads the parent's `ckpts/` onto the fresh reference `TrainState` and keeps ONLY the trained `components` (V/U) + `ci_fn`; the optimizer states, persistent sources, and `step` are the FRESH reference's (`step = 0`). Rationale: a fine-tune runs a NEW LR/p-anneal schedule computed over the new `cfg.steps` from 0, so carrying stale Adam momentum / a stale adversary would mis-scale the restart; faith warmup is also skipped (the parent's V/U is already faithful). The run records lineage via `resume_provenance` in `config.yaml` + `wandb.config`. The parent's decomposition STRUCTURE (sites names + C, ci-fn arch) must match the new config's — asserted from the parent's pinned `config.yaml` (`run.py::assert_finetune_structural_compat`) before the orbax restore; only LR / coeffs / eps / seq / batch / steps may change. This is distinct from same-config requeue-resume (S22): on a subsequent requeue the run's own `ckpts/` is non-empty, so `restore_latest` from its own dir wins and provenance is ignored. | +| S34 | **Stale-CI replay** (config-gated, OFF by default: `DataConfig.replay_stale_ci`, requires `train_batch_replay > 1`; **ADDED 2026-07-06**, experimental — an S4-cadence amendment riding the free-AT batch replay): on the REPEAT steps of a replay window the CI envelope is the window-FIRST step's, reused as a CONSTANT — no taps read, no CI-fn forward or backward, no ci_fn/optimizer update (the ci_fn trains only on window-first steps, matching the 1-per-window unique-data rate). V/U, faithfulness, and both adversaries behave exactly as S12′–S14′ with `ci_lower` a constant (their gradients are unchanged by CI being live vs constant); the imp-min term is computed from the cached upper for logging and contributes zero gradient on repeats. The ci_fn LR follows the run's GLOBAL-step schedule via a unit-LR optimizer + in-step `scheduled_value_traced` (`run_state.unit_lr_ci_fn_optimizer`; optax's update count advances only on window-first steps, so a count-driven schedule would stretch by the replay factor); Adam bias correction correctly tracks the update count. A requeue-resume landing MID-window rebuilds the envelope from the current ci_fn (`StaleCITrainSteps.compute_ci`) — one ci_fn update fresher than the interrupted window's envelope, a tiny preemption-only trajectory wobble (weaker than S22's exactness; accepted for the experimental flag). Flag-off emits the pre-S34 step unchanged. | ## 6. Variation points diff --git a/param_decomp/built_run.py b/param_decomp/built_run.py index a12801fae..db1cb308b 100644 --- a/param_decomp/built_run.py +++ b/param_decomp/built_run.py @@ -49,6 +49,11 @@ class DataConfig: """Consecutive train steps that share one token batch (free-AT-style replay: the persistent adversary's once-per-step ascent becomes batch-matched on the repeated steps). 1 = a fresh batch every step.""" + replay_stale_ci: bool + """Stale-CI replay (SPEC S34): repeat steps of a replay window reuse the window-first + step's CI envelope as a constant — no taps / CI-fn forward / CI-fn backward on + repeats; the ci_fn updates only on window-first steps. Requires + `train_batch_replay > 1`.""" @dataclass(frozen=True) diff --git a/param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay16_staleci_400k.yaml b/param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay16_staleci_400k.yaml new file mode 100644 index 000000000..f7cb5a931 --- /dev/null +++ b/param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay16_staleci_400k.yaml @@ -0,0 +1,179 @@ +# Stale-CI replay-scaling arm (SPEC S34): free-AT batch replay x16 with +# replay_stale_ci — repeat steps reuse the window-first CI envelope as a constant +# (no taps / CI-fn fwd/bwd on repeats; ci_fn updates only on window-first steps, LR +# still on the global-step schedule). Derived from +# pile_ppgd_bsc_nw0_lr0p01_replay16_400k.yaml; only run_name / replay_stale_ci differ, +# so each arm A/Bs stale CI against its running replay-only twin. replay=1 has no +# repeat steps (no stale-CI arm); the replay-only sweep's replay1 is the shared +# unique-data baseline. `steps` counts optimizer steps (replay repeats ARE steps): +# unique batches = 400k / replay. +run_name: pile4l-ppgd-nw0-lr0p01-replay16-staleci-400k +# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of +# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot +# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the +# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile +# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15. +# Target: LlamaSimpleMLP 4L pile model +# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope +# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/ +# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml +# (p-af354eb1, fresh PGD) except this one recon term. 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 + train_batch_replay: 16 + replay_stale_ci: 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: shared_across_batch + 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 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_warmup_steps: 0 + optimizer: + beta1: 0.5 + 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: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +runtime: + remat_recon_forwards: false + dp: 8 +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/configs/pile_ppgd_bsc_nw0_lr0p01_replay1_staleci_base_400k.yaml b/param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay1_staleci_base_400k.yaml new file mode 100644 index 000000000..fa0df4e3c --- /dev/null +++ b/param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay1_staleci_base_400k.yaml @@ -0,0 +1,174 @@ +# Code baseline for the stale-CI replay sweep: replay=1, replay_stale_ci OFF, run from +# the feature/replay-stale-ci branch. Identical config to the running +# pile_ppgd_bsc_nw0_lr0p01_replay1_400k.yaml arm (only run_name differs) — it pins that +# the stale-CI refactor's flag-off step path reproduces the parent branch's trajectory +# at scale before attributing any stale-CI arm differences to the algorithm change. +run_name: pile4l-ppgd-nw0-lr0p01-replay1-staleci-base-400k +# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of +# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot +# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the +# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile +# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15. +# Target: LlamaSimpleMLP 4L pile model +# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope +# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/ +# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml +# (p-af354eb1, fresh PGD) except this one recon term. 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 + train_batch_replay: 1 + 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: shared_across_batch + 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 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_warmup_steps: 0 + optimizer: + beta1: 0.5 + 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: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +runtime: + remat_recon_forwards: false + dp: 8 +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/configs/pile_ppgd_bsc_nw0_lr0p01_replay2_staleci_400k.yaml b/param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay2_staleci_400k.yaml new file mode 100644 index 000000000..81552f0d8 --- /dev/null +++ b/param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay2_staleci_400k.yaml @@ -0,0 +1,179 @@ +# Stale-CI replay-scaling arm (SPEC S34): free-AT batch replay x2 with +# replay_stale_ci — repeat steps reuse the window-first CI envelope as a constant +# (no taps / CI-fn fwd/bwd on repeats; ci_fn updates only on window-first steps, LR +# still on the global-step schedule). Derived from +# pile_ppgd_bsc_nw0_lr0p01_replay2_400k.yaml; only run_name / replay_stale_ci differ, +# so each arm A/Bs stale CI against its running replay-only twin. replay=1 has no +# repeat steps (no stale-CI arm); the replay-only sweep's replay1 is the shared +# unique-data baseline. `steps` counts optimizer steps (replay repeats ARE steps): +# unique batches = 400k / replay. +run_name: pile4l-ppgd-nw0-lr0p01-replay2-staleci-400k +# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of +# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot +# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the +# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile +# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15. +# Target: LlamaSimpleMLP 4L pile model +# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope +# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/ +# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml +# (p-af354eb1, fresh PGD) except this one recon term. 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 + train_batch_replay: 2 + replay_stale_ci: 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: shared_across_batch + 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 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_warmup_steps: 0 + optimizer: + beta1: 0.5 + 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: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +runtime: + remat_recon_forwards: false + dp: 8 +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/configs/pile_ppgd_bsc_nw0_lr0p01_replay4_staleci_400k.yaml b/param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay4_staleci_400k.yaml new file mode 100644 index 000000000..f879ed105 --- /dev/null +++ b/param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay4_staleci_400k.yaml @@ -0,0 +1,179 @@ +# Stale-CI replay-scaling arm (SPEC S34): free-AT batch replay x4 with +# replay_stale_ci — repeat steps reuse the window-first CI envelope as a constant +# (no taps / CI-fn fwd/bwd on repeats; ci_fn updates only on window-first steps, LR +# still on the global-step schedule). Derived from +# pile_ppgd_bsc_nw0_lr0p01_replay4_400k.yaml; only run_name / replay_stale_ci differ, +# so each arm A/Bs stale CI against its running replay-only twin. replay=1 has no +# repeat steps (no stale-CI arm); the replay-only sweep's replay1 is the shared +# unique-data baseline. `steps` counts optimizer steps (replay repeats ARE steps): +# unique batches = 400k / replay. +run_name: pile4l-ppgd-nw0-lr0p01-replay4-staleci-400k +# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of +# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot +# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the +# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile +# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15. +# Target: LlamaSimpleMLP 4L pile model +# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope +# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/ +# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml +# (p-af354eb1, fresh PGD) except this one recon term. 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 + train_batch_replay: 4 + replay_stale_ci: 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: shared_across_batch + 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 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_warmup_steps: 0 + optimizer: + beta1: 0.5 + 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: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +runtime: + remat_recon_forwards: false + dp: 8 +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/configs/pile_ppgd_bsc_nw0_lr0p01_replay8_staleci_400k.yaml b/param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay8_staleci_400k.yaml new file mode 100644 index 000000000..61b54b78c --- /dev/null +++ b/param_decomp/configs/pile_ppgd_bsc_nw0_lr0p01_replay8_staleci_400k.yaml @@ -0,0 +1,179 @@ +# Stale-CI replay-scaling arm (SPEC S34): free-AT batch replay x8 with +# replay_stale_ci — repeat steps reuse the window-first CI envelope as a constant +# (no taps / CI-fn fwd/bwd on repeats; ci_fn updates only on window-first steps, LR +# still on the global-step schedule). Derived from +# pile_ppgd_bsc_nw0_lr0p01_replay8_400k.yaml; only run_name / replay_stale_ci differ, +# so each arm A/Bs stale CI against its running replay-only twin. replay=1 has no +# repeat steps (no stale-CI arm); the replay-only sweep's replay1 is the shared +# unique-data baseline. `steps` counts optimizer steps (replay repeats ARE steps): +# unique batches = 400k / replay. +run_name: pile4l-ppgd-nw0-lr0p01-replay8-staleci-400k +# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of +# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot +# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the +# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile +# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15. +# Target: LlamaSimpleMLP 4L pile model +# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope +# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/ +# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml +# (p-af354eb1, fresh PGD) except this one recon term. 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 + train_batch_replay: 8 + replay_stale_ci: 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: shared_across_batch + 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 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_warmup_steps: 0 + optimizer: + beta1: 0.5 + 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: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +runtime: + remat_recon_forwards: false + dp: 8 +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/run.py b/param_decomp/run.py index b62cf163b..587bfd975 100644 --- a/param_decomp/run.py +++ b/param_decomp/run.py @@ -55,7 +55,7 @@ from param_decomp.configs import Cadence, PDConfig, ProfileConfig, flatten_typed_lists from param_decomp.lm import DecomposedModel from param_decomp.recon import build_loss_terms -from param_decomp.run_state import build_optimizers, init_train_state +from param_decomp.run_state import build_optimizers, init_train_state, unit_lr_ci_fn_optimizer from param_decomp.slow_eval import ( PermutationMetricSpec, PositionCI, @@ -63,7 +63,14 @@ render_permutation_figures, render_slow_eval_figures, ) -from param_decomp.train import TrainState, make_faith_warmup_step, make_train_step +from param_decomp.train import ( + CIEnvelope, + StaleCITrainSteps, + TrainState, + make_faith_warmup_step, + make_stale_ci_train_steps, + make_train_step, +) _sigterm_received = False @@ -159,16 +166,20 @@ def _grad_norm_summary_window_stats(window: list[dict[str, jax.Array]]) -> dict[ so the whole window reduces in a single host transfer here — the loop stays unsynced between logs rather than subsampling grad norms at the log step.""" assert window, "grad-norm summary window is empty at a log boundary" - keys = list(window[0].keys()) - stacked = jnp.stack([jnp.stack([snap[k] for snap in window]) for k in keys]) # [keys, steps] - mins = np.asarray(jnp.min(stacked, axis=1)) - maxs = np.asarray(jnp.max(stacked, axis=1)) - medians = np.asarray(jnp.median(stacked, axis=1)) + # Union of keys: under stale-CI replay the repeat steps carry no `ci_fns` family, so a + # key's stats reduce over the (fresh) steps that emitted it. + keys = sorted({k for snap in window for k in snap}) + stats = { + key: (jnp.min(vals), jnp.max(vals), jnp.median(vals)) + for key in keys + for vals in (jnp.stack([snap[key] for snap in window if key in snap]),) + } + host: dict[str, tuple[np.ndarray, np.ndarray, np.ndarray]] = jax.device_get(stats) out: dict[str, float] = {} - for i, key in enumerate(keys): - out[f"{key}/min"] = float(mins[i]) - out[f"{key}/max"] = float(maxs[i]) - out[f"{key}/median"] = float(medians[i]) + for key, (mn, mx, med) in host.items(): + out[f"{key}/min"] = float(mn) + out[f"{key}/max"] = float(mx) + out[f"{key}/median"] = float(med) return out @@ -482,6 +493,13 @@ def run_decomposition_training( run.run_dir.mkdir(parents=True, exist_ok=True) opt_vu, opt_ci, (sched_vu, sched_ci) = build_optimizers(pd) + # Stale-CI replay (SPEC S34): repeat steps of a `train_batch_replay` window reuse the + # window-first CI envelope as a constant; the ci_fn updates only on window-first steps, + # under a unit-LR optimizer whose schedule the step applies at the global step. + stale_ci = data is not None and data.replay_stale_ci + if stale_ci: + assert data is not None and data.train_batch_replay > 1, data + opt_ci = unit_lr_ci_fn_optimizer(pd) key = random.PRNGKey(pd.seed) init_key, src_key, run_key = random.split(key, 3) @@ -497,18 +515,72 @@ def run_decomposition_training( return # SIGTERM mid-warmup: clean exit for requeue state, start_step = init - step_fn = make_train_step( - lm=lm, - losses=build_loss_terms(pd.loss_metrics, lm.site_names), - components_optimizer=opt_vu, - ci_fn_optimizer=opt_ci, - total_steps=pd.steps, - remat_recon_forwards=remat_recon_forwards, - remat_ci_fn=remat_ci_fn, - ascend_replicate=ascend_replicate, - compiler_options=compiler_options, - mesh=mesh, - ) + stale_steps: StaleCITrainSteps | None = None + step_fn = None + if stale_ci: + assert data is not None + stale_steps = make_stale_ci_train_steps( + lm=lm, + losses=build_loss_terms(pd.loss_metrics, lm.site_names), + components_optimizer=opt_vu, + ci_fn_optimizer=opt_ci, + ci_fn_lr_schedule=pd.ci_fn_optimizer.lr_schedule, + total_steps=pd.steps, + remat_recon_forwards=remat_recon_forwards, + remat_ci_fn=remat_ci_fn, + ascend_replicate=ascend_replicate, + compiler_options=compiler_options, + mesh=mesh, + ) + else: + step_fn = make_train_step( + lm=lm, + losses=build_loss_terms(pd.loss_metrics, lm.site_names), + components_optimizer=opt_vu, + ci_fn_optimizer=opt_ci, + total_steps=pd.steps, + remat_recon_forwards=remat_recon_forwards, + remat_ci_fn=remat_ci_fn, + ascend_replicate=ascend_replicate, + compiler_options=compiler_options, + mesh=mesh, + ) + + replay_window = data.train_batch_replay if data is not None else 1 + ci_envelope: CIEnvelope | None = None + + def do_step( + step: int, state: TrainState, batch: Any, step_key: PRNGKeyArray + ) -> tuple[TrainState, dict[str, jax.Array]]: + """One train step, dispatching fresh vs repeat under stale-CI replay. The envelope + rides in a loop-local (never checkpointed): a requeue-resume landing mid-window + rebuilds it from the current ci_fn (see `StaleCITrainSteps.compute_ci`).""" + nonlocal ci_envelope + if stale_steps is None: + assert step_fn is not None + return step_fn(lm, state, batch, step_key) + if step % replay_window == 0: + new_state, metrics, ci_envelope = stale_steps.fresh(lm, state, batch, step_key) + return new_state, metrics + if ci_envelope is None: + ci_envelope = stale_steps.compute_ci(lm, state.ci_fn, batch) + return stale_steps.repeat((lm, ci_envelope), state, batch, step_key) + + # The profile probes below want a bare `(model, state, batch, key) -> (state, metrics)` + # with `.lower()` available; under stale-CI they probe the fresh (superset) step. + if stale_steps is None: + assert step_fn is not None + _probe_jit: Any = step_fn + _probe_step: Callable[..., Any] = step_fn + else: + _fresh = stale_steps.fresh + _probe_jit = _fresh + + def _probe_fresh_step(m: Any, s: Any, b: Any, k: Any) -> Any: + new_state, metrics, _ = _fresh(m, s, b, k) + return new_state, metrics + + _probe_step = _probe_fresh_step # record what this run actually executes on so wandb never lies about topology. # flatten the metric lists into the same flat keys torch logs (E14) so cross-impl @@ -577,11 +649,11 @@ def _bench_dispatch(tree: object, n: int = 15) -> float: _atk = random.fold_in(run_key, start_step) _ab = sample_batch(start_step) _aa0 = time.perf_counter() - _as1, _am1 = step_fn(lm, state, _ab, _atk) + _as1, _am1 = _probe_step(lm, state, _ab, _atk) _aa1 = time.perf_counter() - _as2, _am2 = step_fn(lm, _as1, _ab, _atk) + _as2, _am2 = _probe_step(lm, _as1, _ab, _atk) _aa2 = time.perf_counter() - _as3, _am3 = step_fn(lm, _as2, _ab, _atk) + _as3, _am3 = _probe_step(lm, _as2, _ab, _atk) _aa3 = time.perf_counter() jax.block_until_ready((_as3, _am3["total"])) _aa4 = time.perf_counter() @@ -599,7 +671,7 @@ def _bench_dispatch(tree: object, n: int = 15) -> float: _mb = sample_batch(start_step) _mk = random.fold_in(run_key, start_step) _step_jit: Any = ( - step_fn # eqx.filter_jit object exposes .lower(); the Callable alias hides it + _probe_jit # eqx.filter_jit object exposes .lower(); the Callable alias hides it ) _compiled = _step_jit.lower(lm, state, _mb, _mk).compile() _ma = getattr(_compiled, "compiled", _compiled).memory_analysis() @@ -614,8 +686,8 @@ def _bench_dispatch(tree: object, n: int = 15) -> float: ) _dev = jax.local_devices()[0] _dev.memory_stats() # reset peak baseline read - _ms_state, _ms_metrics = step_fn(lm, state, _mb, _mk) - _ms_state, _ms_metrics = step_fn(lm, _ms_state, _mb, _mk) + _ms_state, _ms_metrics = _probe_step(lm, state, _mb, _mk) + _ms_state, _ms_metrics = _probe_step(lm, _ms_state, _mb, _mk) jax.block_until_ready((_ms_state, _ms_metrics["total"])) _ms = _dev.memory_stats() if is_main and _ms is not None: @@ -672,7 +744,7 @@ def _bench_dispatch(tree: object, n: int = 15) -> float: _ts0 = time.perf_counter() batch = sample_batch(step) _ts1 = time.perf_counter() - state, metrics = step_fn(lm, state, batch, random.fold_in(run_key, step)) + state, metrics = do_step(step, state, batch, random.fold_in(run_key, step)) _ts2 = time.perf_counter() jax.block_until_ready((state, metrics["total"])) _ts3 = time.perf_counter() @@ -683,7 +755,7 @@ def _bench_dispatch(tree: object, n: int = 15) -> float: ) else: batch = sample_batch(step) - state, metrics = step_fn(lm, state, batch, random.fold_in(run_key, step)) + state, metrics = do_step(step, state, batch, random.fold_in(run_key, step)) grad_norm_summary_window.append( {k: v for k, v in metrics.items() if k.startswith("grad_norms/summary/")} diff --git a/param_decomp/run_state.py b/param_decomp/run_state.py index a20bfc436..87abff284 100644 --- a/param_decomp/run_state.py +++ b/param_decomp/run_state.py @@ -98,6 +98,17 @@ def build_optimizers(pd: PDConfig): return opt_vu, opt_ci, (sched_vu, sched_ci) +def unit_lr_ci_fn_optimizer(pd: PDConfig) -> optax.GradientTransformation: + """The ci_fn AdamW with LR pinned to 1.0, for stale-CI replay (SPEC S34): the ci_fn + updates only on window-first steps, so its optax update count lags the global step by + the replay factor and a count-driven schedule would stretch the cosine. The step + applies the global-step schedule itself (`make_stale_ci_train_steps`). A callable + constant schedule keeps the optimizer-state structure identical to + `build_optimizers`'s (a `scale_by_schedule` count leaf), so checkpoint reference + trees are unaffected.""" + return _adamw_with_clip(pd.ci_fn_optimizer, lambda count: jnp.asarray(1.0, jnp.float32)) + + def init_train_state( pd: PDConfig, lm: DecomposedModel, diff --git a/param_decomp/tests/equivalence/test_equivalence.py b/param_decomp/tests/equivalence/test_equivalence.py index 51873da41..7bd72be14 100644 --- a/param_decomp/tests/equivalence/test_equivalence.py +++ b/param_decomp/tests/equivalence/test_equivalence.py @@ -100,7 +100,7 @@ def test_chunk_plan_static_live_gate() -> None: def test_structure_stoch_is_per_chunk() -> None: """SPEC S10: one forward per (chunk, sample), normalized by `n_chunks · n_samples` — matching the torch chunkwise pool, not one fused forward over all sites.""" - src = inspect.getsource(train_mod.make_train_step) + src = inspect.getsource(train_mod._build_step_impl) assert "for entry_idx, entry in enumerate(term.plan)" in src, ( "each recon term must loop its plan's entries" ) diff --git a/param_decomp/tests/test_llama_simple_mlp.py b/param_decomp/tests/test_llama_simple_mlp.py index 567e524ea..de607d6f6 100644 --- a/param_decomp/tests/test_llama_simple_mlp.py +++ b/param_decomp/tests/test_llama_simple_mlp.py @@ -49,7 +49,12 @@ site_name, site_specs, ) -from param_decomp.train import TrainState, make_faith_warmup_step, make_train_step +from param_decomp.train import ( + TrainState, + make_faith_warmup_step, + make_stale_ci_train_steps, + make_train_step, +) def _tiny_cfg() -> LlamaSimpleMLPConfig: @@ -440,6 +445,127 @@ def test_step_trains_and_has_vpd_signature(): assert state.ci_fn.chunks.in_proj_w.dtype == jnp.float32 +def test_stale_ci_repeat_step_semantics(): + """SPEC S34: a repeat step given the fresh step's envelope trains the components and + ascends the adversary like the fresh step (CI live vs constant doesn't change their + gradients), while leaving the ci_fn and its optimizer state untouched.""" + cfg = _tiny_cfg() + seq = 16 + sites = site_specs(cfg, _MIXED_SITE_CS) + lm = _tiny_decomposed_model(cfg, sites, jax.random.PRNGKey(0)) + vu = init_decomp_vu(sites, jax.random.PRNGKey(1)) + ci_fn = _build_chunkwise_ci_fn(lm, jax.random.PRNGKey(2)) + opt_vu = optax.chain(optax.clip_by_global_norm(0.01), optax.adamw(1e-3, weight_decay=0.0)) + # Unit-lr ci_fn optimizer: the stale-CI fresh step applies the schedule at the global step. + opt_ci = optax.adamw(1.0, weight_decay=0.0) + + src = init_persistent_sources( + lm.site_names, tuple(s.C for s in lm.sites), (1, seq), jnp.float32, jax.random.PRNGKey(3) + ) + ppgd_cfg = PersistentPGDReconLossConfig( + coeff=0.5, + scope=SCScope(), + optimizer=AdamPGDConfig( + beta1=0.5, + beta2=0.99, + lr_schedule=ScheduleConfig(start_val=0.01, warmup_pct=0.025), + ), + n_warmup_steps=0, + ) + ppgd_coeff = ppgd_cfg.coeff + assert ppgd_coeff is not None + + def copied[T](tree: T) -> T: + """Fresh buffers per step call — the jitted steps donate their state/batch args.""" + return jax.tree.map(lambda x: jnp.copy(x) if eqx.is_array(x) else x, tree) + + def make_state() -> TrainState: + vu_c, ci_fn_c, src_c = copied(vu), copied(ci_fn), copied(src) + return TrainState( + components=vu_c, ci_fn=ci_fn_c, + components_opt_state=opt_vu.init(eqx.filter(vu_c, eqx.is_array)), + ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn_c, eqx.is_array)), + adversaries={ + ppgd_cfg.type: PersistentAdversary( + sources=src_c, + opt_state=init_sources_adam_state(src_c), + state_key=ppgd_cfg.type, + coeff=ppgd_coeff, + adam=ppgd_cfg.optimizer, + n_warmup=ppgd_cfg.n_warmup_steps, + ) + }, + step=jnp.zeros((), jnp.int32), + ) # fmt: skip + + loss_terms = build_loss_terms( + ( + FaithfulnessLossConfig(coeff=1e5), + ImportanceMinimalityLossConfig( + coeff=5e-6, + pnorm=ScheduleConfig(start_val=2.0, fn_type="linear", final_val_frac=0.2), + ), + ChunkwiseSubsetReconLossConfig( + routing=UniformKSubsetRoutingConfig(), coeff=0.5, sites_per_chunk=2, n_samples=1 + ), + ppgd_cfg, + ), + lm.site_names, + ) + steps = make_stale_ci_train_steps( + lm=lm, + losses=loss_terms, + components_optimizer=opt_vu, + ci_fn_optimizer=opt_ci, + ci_fn_lr_schedule=ScheduleConfig(start_val=1e-3), + total_steps=100, + remat_recon_forwards=True, + remat_ci_fn=False, + mesh=None, + ) + + tokens = jax.random.randint(jax.random.PRNGKey(4), (2, seq), 0, cfg.vocab_size) + key = jax.random.PRNGKey(100) + + envelope = steps.compute_ci(lm, ci_fn, tokens) + state_fresh, m_fresh, env_fresh = steps.fresh(lm, make_state(), jnp.copy(tokens), jnp.copy(key)) + state_rep, m_rep = steps.repeat((lm, envelope), make_state(), jnp.copy(tokens), jnp.copy(key)) + + # compute_ci reproduces the fresh step's own envelope (same ci_fn, same batch). + for site in envelope.lower: + assert jnp.allclose(envelope.lower[site], env_fresh.lower[site], atol=1e-6) + assert jnp.allclose(envelope.upper[site], env_fresh.upper[site], atol=1e-6) + + # The repeat step leaves the ci_fn and its optimizer state bit-untouched... + for got, want in zip( + jax.tree.leaves(eqx.filter(state_rep.ci_fn, eqx.is_array)), + jax.tree.leaves(eqx.filter(ci_fn, eqx.is_array)), + strict=True, + ): + assert (got == want).all() + # ...while the fresh step trained it. + fresh_ci_leaves = jax.tree.leaves(eqx.filter(state_fresh.ci_fn, eqx.is_array)) + init_ci_leaves = jax.tree.leaves(eqx.filter(ci_fn, eqx.is_array)) + assert any((a != b).any() for a, b in zip(fresh_ci_leaves, init_ci_leaves, strict=True)) + + # Components + adversary see the same gradients either way (CI enters their backward + # as the same values, live or constant) — updates match to numeric tolerance. + for (fv, fu), (rv, ru) in zip( + state_fresh.components.vu.values(), state_rep.components.vu.values(), strict=True + ): + assert jnp.allclose(fv, rv, rtol=1e-4, atol=1e-6) + assert jnp.allclose(fu, ru, rtol=1e-4, atol=1e-6) + adv_f = state_fresh.adversaries[ppgd_cfg.type] + adv_r = state_rep.adversaries[ppgd_cfg.type] + for site in adv_f.sources: + assert jnp.allclose(adv_f.sources[site], adv_r.sources[site], rtol=1e-4, atol=1e-6) + + # Metric families: the repeat step carries no ci_fn grad norms; totals stay finite. + assert "grad_norms/summary/ci_fns" in m_fresh and "grad_norms/summary/ci_fns" not in m_rep + assert jnp.isfinite(m_rep["total"]) and jnp.isfinite(m_fresh["total"]) + assert int(state_rep.step) == 1 + + def test_faith_warmup_decreases_faith(): cfg = _tiny_cfg() sites = site_specs(cfg, canonical_site_cs(_MIXED_SITE_CS)) diff --git a/param_decomp/train.py b/param_decomp/train.py index 6694d2157..0602a6503 100644 --- a/param_decomp/train.py +++ b/param_decomp/train.py @@ -42,6 +42,7 @@ annealed_imp_min_param, faithfulness_loss, imp_min_terms, + scheduled_value_traced, ) from param_decomp.recon import ( ConstantSources, @@ -52,6 +53,7 @@ Routes, StochasticSources, ) +from param_decomp.schedule import ScheduleConfig from param_decomp.sharding import batch_shard_leading COMPUTE_DT = jnp.bfloat16 @@ -75,12 +77,53 @@ class TrainState: step: Array -def _grad_norm_metrics(components_grad: DecompVU, ci_fn_grad: Any) -> dict[str, Array]: +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class CIEnvelope: + """A step's detached CI squashings (stop-gradient'd, compute-dtype, sharding-pinned). + Under stale-CI replay (SPEC S34) the window-first step returns its envelope and the + window's repeat steps consume it as a CONSTANT.""" + + lower: dict[str, Array] + upper: dict[str, Array] + + +@dataclass(frozen=True) +class StaleCITrainSteps: + """The jitted step pair for stale-CI replay (SPEC S34), plus the mid-window resume + rebuild. + + `fresh` is the full step (taps → CI-fn forward/vjp → ci_fn update), additionally + returning its `CIEnvelope` for the window's repeats. `repeat` consumes that envelope + as a constant — no taps, no CI-fn forward/backward; the ci_fn and its optimizer state + pass through untouched. Its first arg bundles `(model, envelope)` so donation + (`all-except-first`) spares both across the window's repeats. + + `compute_ci(model, ci_fn, batch) -> CIEnvelope` rebuilds the envelope when a + requeue-resume lands mid-window. The ci_fn updates only on window-first steps, so the + rebuilt envelope reflects the last fresh step's POST-update ci_fn — one ci_fn Adam + step fresher than the envelope the interrupted window was using (plus vjp-primal vs + plain-forward compilation differences). A tiny, preemption-only trajectory wobble.""" + + fresh: Callable[ + [DecomposedModel, TrainState, Any, PRNGKeyArray], + tuple[TrainState, dict[str, Array], CIEnvelope], + ] + repeat: Callable[ + [tuple[DecomposedModel, CIEnvelope], TrainState, Any, PRNGKeyArray], + tuple[TrainState, dict[str, Array]], + ] + compute_ci: Callable[[DecomposedModel, Any, Any], CIEnvelope] + + +def _grad_norm_metrics(components_grad: DecompVU, ci_fn_grad: Any | None) -> dict[str, Array]: """Pre-clip gradient L2 norms, matching the torch `component_grad_norms` families: per-leaf `grad_norms/components` / `grad_norms/ci_fns` (paths are this pytree's own — e.g. `.vu['layers.18.mlp.gate_proj'][0]` for the per-site Llama layout, vs torch's per-site names) and the overlay-critical - `grad_norms/summary/{components,ci_fns,total}`.""" + `grad_norms/summary/{components,ci_fns,total}`. `ci_fn_grad` is None on a stale-CI + repeat step (no CI-fn backward): the `ci_fns` family is absent and `total` covers + the components only.""" out: dict[str, Array] = {} def family(grad_tree: Any, prefix: str) -> Array: @@ -92,7 +135,9 @@ def family(grad_tree: Any, prefix: str) -> Array: out[f"grad_norms/summary/{prefix}"] = jnp.sqrt(sum_sq) return sum_sq - total_sq = family(components_grad, "components") + family(ci_fn_grad, "ci_fns") + total_sq = family(components_grad, "components") + if ci_fn_grad is not None: + total_sq = total_sq + family(ci_fn_grad, "ci_fns") out["grad_norms/summary/total"] = jnp.sqrt(total_sq) return out @@ -100,7 +145,7 @@ def family(grad_tree: Any, prefix: str) -> Array: # ───────────────────────────── the step factory ───────────────────────────── -def make_train_step( +def _build_step_impl( lm: DecomposedModel, *, losses: LossSurface, @@ -110,18 +155,21 @@ def make_train_step( remat_recon_forwards: bool, remat_ci_fn: bool, mesh: Mesh | None, - ascend_replicate: bool = False, - compiler_options: dict[str, bool | int | str] | None = None, + ascend_replicate: bool, + stale_ci_fn_lr: ScheduleConfig | None, ): - """Build the `eqx.filter_jit`'d `step(model, state, batch, key) -> (state, metrics)`. - - `model` is the jit ARG (frozen 8B weights traced as array leaves, never baked); the - factory closes over only static config (`site_names`, `recon_loss_fn`, term wiring) read - off `lm` here. `losses` (from `build_loss_terms`) is the `LossSurface` record — the - faithfulness + importance-minimality singletons and the recon Σ, read by name. `mesh` - (when given) pins every batch-leading activation over the full mesh - (`P(('replicate', 'fsdp'), ...)`) so the masked re-forwards stay on per-rank sub-batches - (activation memory 1/N).""" + """The un-jitted step machinery shared by `make_train_step` and + `make_stale_ci_train_steps`: returns `(step_impl, compute_ci_impl)` where + `step_impl(model, state, batch, key, cached)` runs the full step when `cached is None` + (computing taps + the CI envelope, training the ci_fn) and the stale-CI repeat step + when `cached` is a `CIEnvelope` (SPEC S34: CI a constant, ci_fn untouched). The + `cached` switch is a Python-level trace decision — each variant jits to its own + program, and the `cached is None` program is the pre-S34 step unchanged. + + `stale_ci_fn_lr` (stale-CI mode only) is the ci_fn LR schedule applied IN-STEP at the + global step: the ci_fn optimizer must then carry unit LR (`unit_lr_ci_fn_optimizer`), + since its optax update count advances only on window-first steps and a count-driven + schedule would stretch by the replay factor.""" site_names = lm.site_names sites = lm.sites recon_loss_fn = lm.recon_loss_fn # static method: pure, holds no arrays — safe to close @@ -252,24 +300,29 @@ def entry_loss_for_sources( return total / len(routes_per_draw) @jaxtyped(typechecker=beartype) - def step( + def step_impl( model: DecomposedModel, state: TrainState, batch: Any, key: PRNGKeyArray, - ) -> tuple[TrainState, dict[str, Array]]: + cached: CIEnvelope | None, + ) -> tuple[TrainState, dict[str, Array], CIEnvelope]: step_f32 = state.step.astype(jnp.float32) imp_min_param = annealed_imp_min_param(step_f32, total_steps, imp_min) batch = batch_sharded(batch) with jax.named_scope("pd_clean_fwd"): clean_output = jax.lax.stop_gradient(batch_sharded(model.clean_output(batch))) - with jax.named_scope("pd_read_taps"): - taps = model.read_activations(batch, state.ci_fn.input_names) - # `leading` (batch, *positions) — the shape masks/sources/routes live in. Sourced - # from a tap (always `[*leading, d_tap]`), not the opaque batch, so the engine never - # assumes the batch's rank/feature dim. - leading = next(iter(taps.values())).shape[:-1] + if cached is None: + with jax.named_scope("pd_read_taps"): + taps = model.read_activations(batch, state.ci_fn.input_names) + # `leading` (batch, *positions) — the shape masks/sources/routes live in. Sourced + # from a tap (always `[*leading, d_tap]`), not the opaque batch, so the engine never + # assumes the batch's rank/feature dim. + leading = next(iter(taps.values())).shape[:-1] + else: + taps = None + leading = next(iter(cached.lower.values())).shape[:-1] # ── adversary ascents: params + CI detached (SPEC §4.5) ── prepared, recon_vjp = jax.vjp( @@ -282,13 +335,23 @@ def step( # its vjp, mirroring `prepared`/`recon_vjp`. The ascend uses the stop_gradient'd value; # `loss_fn` takes the live value and its ci-fn grad is pulled back through `ci_vjp`. So the # (≈10x-the-target) CI fn is forward-evaluated ONCE, not once detached for the ascend + - # once inside the main backward. - with jax.named_scope("pd_ci_fn_fwd"): - ci, ci_vjp = eqx.filter_vjp( - lambda cf: ci_batch_sharded(cast_floating(cf, COMPUTE_DT)(taps, remat=remat_ci_fn)), - state.ci_fn, - ) - ci_lower_detached = jax.lax.stop_gradient(ci).lower + # once inside the main backward. On a stale-CI repeat step it is not evaluated at all: + # the window-first envelope is the constant CI everywhere. + if cached is None: + assert taps is not None + with jax.named_scope("pd_ci_fn_fwd"): + ci, ci_vjp = eqx.filter_vjp( + lambda cf: ci_batch_sharded( + cast_floating(cf, COMPUTE_DT)(taps, remat=remat_ci_fn) + ), + state.ci_fn, + ) + ci_detached = jax.lax.stop_gradient(ci) + envelope = CIEnvelope(lower=ci_detached.lower, upper=ci_detached.upper) + else: + ci, ci_vjp = None, None + envelope = cached + ci_lower_detached = envelope.lower # ── persistent adversaries: each runs its supplemental ascents vs the route-ALL # all-sites forward (SPEC S24 — torch warmup parity, NOT the term's loss plan), @@ -365,21 +428,25 @@ def sign_ascend_body( ) fresh_sources[(term_idx, entry_idx)] = jax.lax.stop_gradient(ascended) - # ── main losses: live components/ci; the PERSISTENT sources participate in + # ── main losses: live components (+ live ci when fresh — on a stale-CI repeat the + # envelope enters as a CONSTANT); the PERSISTENT sources participate in # the graph so their gradient comes from the SAME backward (SPEC S14'); they # are NOT detached here, but components/ci grads through them are what torch # gets too (sources are leaves). ── - def loss_fn( - trainable: tuple[Any, DecompVU, CI, dict[str, dict[str, Array]]], + def losses_body( + prepared: Any, + components: DecompVU, + ci_lower: dict[str, Array], + ci_upper: dict[str, Array], + persistent_sources: dict[str, dict[str, Array]], ) -> tuple[Array, tuple[Array, Array, Array, tuple[Array, ...]]]: - prepared, components, ci, persistent_sources = trainable # Stochastic recon builds its masks INSIDE the target's `masked_output_stochastic` # from this once-per-step shared CI form — a scan target recomputes them in its # checkpointed block (mask never held, the memory win); others fall back to building # masks then `masked_output`. Either way the engine holds no per-forward mask stacks. - ci_stacked = model.stack_ci(ci.lower) + ci_stacked = model.stack_ci(ci_lower) faith_loss = faithfulness_loss(model.weight_deltas(components)) - imp_lp, imp_freq = imp_min_terms(ci.upper, imp_min, imp_min_param) + imp_lp, imp_freq = imp_min_terms(ci_upper, imp_min, imp_min_param) term_losses: list[Array] = [] for term_idx, term in enumerate(recon_terms): @@ -429,12 +496,12 @@ def pre_built_fwd( ) case ConstantSources() as strategy: masked = pre_built_fwd( - constant_entry_masks(strategy, ci.lower, entry.live_sites) + constant_entry_masks(strategy, ci_lower, entry.live_sites) ) case FreshPGDSources(): masked = pre_built_fwd( source_masks( - ci.lower, + ci_lower, fresh_sources[(term_idx, entry_idx)], entry.live_sites, ) @@ -442,7 +509,7 @@ def pre_built_fwd( case PersistentSources(state_key=state_key): masked = pre_built_fwd( source_masks( - ci.lower, + ci_lower, persistent_sources[state_key], entry.live_sites, ) @@ -458,20 +525,43 @@ def pre_built_fwd( total_loss = total_loss + term.coeff * term_loss return total_loss, (faith_loss, imp_lp, imp_freq, tuple(term_losses)) - with jax.named_scope("pd_value_and_grad"): - (total_loss, (faith_loss, imp_lp, imp_freq, term_losses)), grads = ( - eqx.filter_value_and_grad(loss_fn, has_aux=True)( - ( - prepared, - state.components, - ci, - {k: a.sources for k, a in warmed_advs.items()}, + warmed_sources = {k: a.sources for k, a in warmed_advs.items()} + if cached is None: + assert ci is not None and ci_vjp is not None + + def loss_fn( + trainable: tuple[Any, DecompVU, CI, dict[str, dict[str, Array]]], + ) -> tuple[Array, tuple[Array, Array, Array, tuple[Array, ...]]]: + prepared_t, components_t, ci_t, persistent_t = trainable + return losses_body(prepared_t, components_t, ci_t.lower, ci_t.upper, persistent_t) + + with jax.named_scope("pd_value_and_grad"): + (total_loss, (faith_loss, imp_lp, imp_freq, term_losses)), grads = ( + eqx.filter_value_and_grad(loss_fn, has_aux=True)( + (prepared, state.components, ci, warmed_sources) ) ) - ) - prepared_grad, components_grad_faith, ci_grad, persistent_grads_scaled = grads + prepared_grad, components_grad_faith, ci_grad, persistent_grads_scaled = grads + ci_fn_grad = ci_vjp(ci_grad)[0] + else: + + def loss_fn_repeat( + trainable: tuple[Any, DecompVU, dict[str, dict[str, Array]]], + ) -> tuple[Array, tuple[Array, Array, Array, tuple[Array, ...]]]: + prepared_t, components_t, persistent_t = trainable + return losses_body( + prepared_t, components_t, envelope.lower, envelope.upper, persistent_t + ) + + with jax.named_scope("pd_value_and_grad"): + (total_loss, (faith_loss, imp_lp, imp_freq, term_losses)), grads = ( + eqx.filter_value_and_grad(loss_fn_repeat, has_aux=True)( + (prepared, state.components, warmed_sources) + ) + ) + prepared_grad, components_grad_faith, persistent_grads_scaled = grads + ci_fn_grad = None components_grad_recon = recon_vjp(prepared_grad)[0] - ci_fn_grad = ci_vjp(ci_grad)[0] components_grad = jax.tree.map( lambda recon_g, faith_g: recon_g + faith_g, components_grad_recon, @@ -494,11 +584,24 @@ def pre_built_fwd( state.components_opt_state, eqx.filter(state.components, eqx.is_array), ) - ci_fn_updates, new_ci_fn_opt_state = ci_fn_optimizer.update( - ci_fn_grad, state.ci_fn_opt_state, eqx.filter(state.ci_fn, eqx.is_array) - ) + if ci_fn_grad is not None: + ci_fn_updates, new_ci_fn_opt_state = ci_fn_optimizer.update( + ci_fn_grad, state.ci_fn_opt_state, eqx.filter(state.ci_fn, eqx.is_array) + ) + if stale_ci_fn_lr is not None: + # Unit-lr optimizer (see factory docstring): apply the run's ci_fn LR + # schedule here, at the GLOBAL step — matching what optax's count-driven + # schedule applies when every step updates the ci_fn. + ci_fn_lr = scheduled_value_traced(step_f32, total_steps, stale_ci_fn_lr) + ci_fn_updates = jax.tree.map(lambda u: ci_fn_lr * u, ci_fn_updates) + else: + ci_fn_updates, new_ci_fn_opt_state = None, state.ci_fn_opt_state new_components = eqx.apply_updates(state.components, components_updates) - new_ci_fn = eqx.apply_updates(state.ci_fn, ci_fn_updates) + new_ci_fn = ( + eqx.apply_updates(state.ci_fn, ci_fn_updates) + if ci_fn_updates is not None + else state.ci_fn + ) new_state = TrainState( components=new_components, @@ -524,11 +627,118 @@ def pre_built_fwd( metrics["src_lr"] = next(iter(source_lrs.values())) else: metrics |= {f"schedules/lr/src/{k}": v for k, v in source_lrs.items()} + return new_state, metrics, envelope + + def compute_ci_impl(model: DecomposedModel, ci_fn: CIFn, batch: Any) -> CIEnvelope: + batch = batch_sharded(batch) + taps = model.read_activations(batch, ci_fn.input_names) + ci = jax.lax.stop_gradient( + ci_batch_sharded(cast_floating(ci_fn, COMPUTE_DT)(taps, remat=remat_ci_fn)) + ) + return CIEnvelope(lower=ci.lower, upper=ci.upper) + + return step_impl, compute_ci_impl + + +def make_train_step( + lm: DecomposedModel, + *, + losses: LossSurface, + components_optimizer: optax.GradientTransformation, + ci_fn_optimizer: optax.GradientTransformation, + total_steps: int, + remat_recon_forwards: bool, + remat_ci_fn: bool, + mesh: Mesh | None, + ascend_replicate: bool = False, + compiler_options: dict[str, bool | int | str] | None = None, +): + """Build the `eqx.filter_jit`'d `step(model, state, batch, key) -> (state, metrics)`. + + `model` is the jit ARG (frozen 8B weights traced as array leaves, never baked); the + factory closes over only static config (`site_names`, `recon_loss_fn`, term wiring) read + off `lm` here. `losses` (from `build_loss_terms`) is the `LossSurface` record — the + faithfulness + importance-minimality singletons and the recon Σ, read by name. `mesh` + (when given) pins every batch-leading activation over the full mesh + (`P(('replicate', 'fsdp'), ...)`) so the masked re-forwards stay on per-rank sub-batches + (activation memory 1/N).""" + step_impl, _ = _build_step_impl( + lm, + losses=losses, + components_optimizer=components_optimizer, + ci_fn_optimizer=ci_fn_optimizer, + total_steps=total_steps, + remat_recon_forwards=remat_recon_forwards, + remat_ci_fn=remat_ci_fn, + mesh=mesh, + ascend_replicate=ascend_replicate, + stale_ci_fn_lr=None, + ) + + def step( + model: DecomposedModel, state: TrainState, batch: Any, key: PRNGKeyArray + ) -> tuple[TrainState, dict[str, Array]]: + new_state, metrics, _ = step_impl(model, state, batch, key, None) return new_state, metrics return filter_jit(step, donate="all-except-first", compiler_options=compiler_options) +def make_stale_ci_train_steps( + lm: DecomposedModel, + *, + losses: LossSurface, + components_optimizer: optax.GradientTransformation, + ci_fn_optimizer: optax.GradientTransformation, + ci_fn_lr_schedule: ScheduleConfig, + total_steps: int, + remat_recon_forwards: bool, + remat_ci_fn: bool, + mesh: Mesh | None, + ascend_replicate: bool = False, + compiler_options: dict[str, bool | int | str] | None = None, +) -> StaleCITrainSteps: + """Build the stale-CI replay step pair (SPEC S34; `DataConfig.replay_stale_ci`). + + `ci_fn_optimizer` MUST carry unit LR (`run_state.unit_lr_ci_fn_optimizer`): the fresh + step applies `ci_fn_lr_schedule` itself at the GLOBAL step, so the ci_fn follows the + run's LR schedule even though its optax update count only advances on window-first + steps (a count-driven schedule would stretch by the replay factor).""" + step_impl, compute_ci_impl = _build_step_impl( + lm, + losses=losses, + components_optimizer=components_optimizer, + ci_fn_optimizer=ci_fn_optimizer, + total_steps=total_steps, + remat_recon_forwards=remat_recon_forwards, + remat_ci_fn=remat_ci_fn, + mesh=mesh, + ascend_replicate=ascend_replicate, + stale_ci_fn_lr=ci_fn_lr_schedule, + ) + + def fresh( + model: DecomposedModel, state: TrainState, batch: Any, key: PRNGKeyArray + ) -> tuple[TrainState, dict[str, Array], CIEnvelope]: + return step_impl(model, state, batch, key, None) + + def repeat( + model_and_envelope: tuple[DecomposedModel, CIEnvelope], + state: TrainState, + batch: Any, + key: PRNGKeyArray, + ) -> tuple[TrainState, dict[str, Array]]: + model, cached_envelope = model_and_envelope + new_state, metrics, _ = step_impl(model, state, batch, key, cached_envelope) + return new_state, metrics + + return StaleCITrainSteps( + fresh=filter_jit(fresh, donate="all-except-first", compiler_options=compiler_options), + repeat=filter_jit(repeat, donate="all-except-first", compiler_options=compiler_options), + compute_ci=filter_jit(compute_ci_impl, compiler_options=compiler_options), + ) + + # ───────────────────────────── faithfulness warmup (SPEC S21) ───────────────────────────── diff --git a/param_decomp_lab/experiments/lm/config.py b/param_decomp_lab/experiments/lm/config.py index c3deafbda..1f9d42645 100644 --- a/param_decomp_lab/experiments/lm/config.py +++ b/param_decomp_lab/experiments/lm/config.py @@ -141,6 +141,13 @@ class LMDataConfig(BaseConfig): default=1, description="Consecutive train steps sharing one batch (free-AT replay); 1 = fresh batch every step", ) + replay_stale_ci: bool = Field( + default=False, + description=( + "Reuse the window-first CI envelope on replay repeat steps (skip taps + CI-fn " + "fwd/bwd; ci_fn updates only on fresh batches). Requires train_batch_replay > 1" + ), + ) class LMExperimentConfig(ExperimentConfig[LMTargetConfig, LMDataConfig]): @@ -320,6 +327,9 @@ def _data(cfg: LMExperimentConfig) -> DataConfig: ) assert data.dataset_name == "parquet" and data.column_name == "input_ids", data assert data.data_files is not None + assert not data.replay_stale_ci or data.train_batch_replay > 1, ( + "replay_stale_ci needs a replay window (train_batch_replay > 1) to have repeat steps" + ) shard_glob = Path(data.data_files) assert shard_glob.name == "*.parquet", f"expected a *.parquet glob, got {data.data_files}" return DataConfig( @@ -327,6 +337,7 @@ def _data(cfg: LMExperimentConfig) -> DataConfig: seq_len=data.max_seq_len, global_batch=cfg.pd.batch_size, train_batch_replay=data.train_batch_replay, + replay_stale_ci=data.replay_stale_ci, ) From 16ce66fe375f20059ebc07e5f56143fef851d076 Mon Sep 17 00:00:00 2001 From: Dan Braun Date: Mon, 6 Jul 2026 13:53:02 +0100 Subject: [PATCH 2/2] config(sweep): nw2/replay1 warmup-ascent control for the stale-CI sweep Twin of the scaling sweep's nw2 control (only run_name differs), run from this branch's codebase with replay_stale_ci off. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01QKw3iikejknWeF4ZQDFZbL --- ..._nw2_lr0p01_replay1_staleci_base_400k.yaml | 174 ++++++++++++++++++ 1 file changed, 174 insertions(+) create mode 100644 param_decomp/configs/pile_ppgd_bsc_nw2_lr0p01_replay1_staleci_base_400k.yaml diff --git a/param_decomp/configs/pile_ppgd_bsc_nw2_lr0p01_replay1_staleci_base_400k.yaml b/param_decomp/configs/pile_ppgd_bsc_nw2_lr0p01_replay1_staleci_base_400k.yaml new file mode 100644 index 000000000..a1bda53a0 --- /dev/null +++ b/param_decomp/configs/pile_ppgd_bsc_nw2_lr0p01_replay1_staleci_base_400k.yaml @@ -0,0 +1,174 @@ +# Warmup-ascent control for the stale-CI sweep: n_warmup_steps=2, replay=1, +# replay_stale_ci OFF, run from the feature/replay-stale-ci branch. Identical to +# pile_ppgd_bsc_nw2_lr0p01_replay1_400k.yaml (the scaling sweep's nw2 control; only +# run_name differs) — the torch-production supplemental-ascent strength the stale-CI +# arms' batch-matched ascents are competing against, on this sweep's codebase. +run_name: pile4l-ppgd-nw2-lr0p01-replay1-staleci-base-400k +# JAX leg of the persistent-PGD pile experiment — byte-derived from the stored config of +# torch run p-20f9fc15 (the pile_llama_simple_mlp-4L baseline; snapshot +# runs/p-20f9fc15/experiment_config.yaml). This is the same config José used for the +# paper's parameter decomposition (the "Interpreting Language Model Parameters" 4L Pile +# run): param_decomp_lab/experiments/lm/pile_llama_simple_mlp-4L.yaml, run as p-20f9fc15. +# Target: LlamaSimpleMLP 4L pile model +# (t-9d2b8f02), all h.* sites, per-site C. Adversary: PersistentPGDReconLoss, scope +# per_batch_per_position (-> bsc), n_warmup_steps=2, Adam(0.5,0.99) lr const 0.01 w/ +# 2.5% warmup, clamp parameterization. Identical to pile_llama_simple_mlp_4l_pgd1.yaml +# (p-af354eb1, fresh PGD) except this one recon term. 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 + train_batch_replay: 1 + 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: shared_across_batch + 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 + gamma: + start_val: 1.0 + fn_type: linear + final_val_frac: 0.01 + type: SmoothL0ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + 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: per_batch_per_position + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 400000 +runtime: + remat_recon_forwards: false + dp: 8 +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