feat(checkpoint): split checkpoints into decomposition | training items (S22/S33)

CLAUDE.md

@@ -204,7 +204,7 @@ Every artifact for a decomposition lives under one dir per run:
 ```
 PARAM_DECOMP_OUT_DIR/runs/<run_id>/
   launch_config.yaml         # the single self-contained run config (the trainer reads it; resume byte-compares). NOT config.yaml: that basename collides with wandb's reserved run-config file, which wandb.save would symlink onto and clobber
-  ckpts/<step>/...           # orbax sharded checkpoints (JAX trainer)
+  ckpts/<step>/{decomposition,training}/  # orbax sharded checkpoints (JAX trainer): the trained product vs the trainer-only tail (pre-split default/-item runs need an ad-hoc migration, no in-code compat)
   metrics.jsonl              # local logs
   harvest/h-*/...            # pd-harvest output
   autointerp/a-*/...         # pd-autointerp output

param_decomp/CLAUDE.md

@@ -210,7 +210,10 @@ is a pre-tokenized parquet artifact under
 for Llama-8B, `pile_neox_tok_512` for `LlamaSimpleMLP`) — NEVER stream/tokenize from
 HF at run time (the 80-rank thunderherd lesson). The batch schedule is a pure
 function of `(seed, step)` (O(1) resume, no replay); checkpoints are orbax sharded
-saves (no on-loop full-gather); SIGTERM → save → SLURM requeue → resume from latest.
+saves (no on-loop full-gather), TWO items per step — `decomposition` (V/U + ci_fn, the
+product every consumer restores alone) and `training` (opt states + adversaries + step,
+trainer-only) — a clean break, no in-code compat for pre-split `default`-item runs (SPEC
+S22); SIGTERM → save → SLURM requeue → resume from latest.
 Resume with a changed config is refused (byte-compare). Smokes before a long run
 MUST exercise save AND resume at the production per-rank shape.
 
@@ -235,8 +238,8 @@ resume_provenance:
   parent_step: 175000
 ```
 
-On the FIRST entry (own `ckpts/` empty) the trainer loads `parent_run_dir/ckpts/175000`
-onto the fresh reference and keeps ONLY the components + ci_fn; the optimizer states,
+On the FIRST entry (own `ckpts/` empty) the trainer restores the `decomposition` item of
+`parent_run_dir/ckpts/175000` onto the fresh reference; the optimizer states,
 persistent sources, and `step` are FRESH (`step = 0`, no faith warmup) so the new LR /
 p-anneal schedule recomputes over the new `cfg.steps` from 0. A subsequent SLURM requeue
 (own `ckpts/` now non-empty) resumes from the run's own dir and ignores provenance.

param_decomp/SPEC.md

@@ -291,7 +291,7 @@ faithfulness, sources/PPGD, the squashings (S5/S6), the imp-min reduction, the g
 | S19 | Components gradients are global-norm-clipped at `0.01`, before the optimizer step. CI fn is unclipped (production). The clip coefficient uses torch's eps convention: `clip_coef = min(1, max_norm / (total_norm + 1e-6))` (`torch.nn.utils.clip_grad_norm_`). This `+1e-6` is canonical and matched JAX-side (#643); plain `optax.clip_by_global_norm` divides by `max(total_norm, max_norm)` with NO eps, which at `clip=0.01` (clip fires almost every step) gives a ~1e-4 relative component-grad difference each step — a real per-step deviation the JAX clip must avoid by reproducing the `+1e-6`. |
 | S20 | Both main optimizers are AdamW, `wd=0`, betas `(0.9, 0.999)`, eps `1e-8`; LR cosine to `0.1×` start, no warmup, stepped per training step. Cosine convention is canonical-torch: `progress = (step − warmup) / (decay_steps − 1)`, so LR reaches `0.1×` at `step = steps − 1` (`schedule.py`). This is matched JAX-side (#642); plain `optax.cosine_decay_schedule(lr, steps, alpha=0.1)` divides by `steps` and reaches `0.1×` one update LATER (at `count = steps`), a genuine formula divergence of ~O(1/steps) ≈ 2.5e-6/step at 400k that the JAX schedule must avoid by using the `decay_steps − 1` denominator. Both main LRs are evaluated by `scheduled_value_traced` honoring the full `ScheduleConfig`; the canonical cosine-to-0.1-no-warmup shape is pinned by the lab conversion gate (`assert_canonical_algorithm_config`). |
 | S21 | Faithfulness warmup (400 × AdamW lr `1e-3` on `faithfulness_loss` alone) precedes step 0; its optimizer is discarded. |
-| S22 | Checkpoints round-trip ALL trajectory state of §3 — including every persistent term's sources + SRC_STEP moments + step/schedule counters — such that a resumed run continues the same trajectory (modulo RNG streams and kernel nondeterminism, cf. D4). **SIGTERM lifecycle (decision):** a SLURM-delivered SIGTERM sets a flag that is serviced at the main train-step boundary (synchronous save of the completed step, then exit for requeue), inside the faith-warmup loop (clean exit WITHOUT a save — no valid checkpoint exists pre-step-0, and resume skips warmup whenever any checkpoint is present, so a partial step-0 save would resume as if fully warmed; the requeue redoes warmup), and inside the in-loop eval pass (abandon the partial pass unlogged, fall through to the step-boundary save). The **first-jit-compile window remains unserviced** — a SIGTERM there is honored only once compilation finishes and control reaches the next servicing point. Periodic `save_every` is the BACKSTOP guarantee for every window; the SIGTERM→save path is the low-latency fast path, not the sole guarantee (lore `jsp_sigterm_save_never_fired`: 6/6 historical preemptions fell back to periodic ckpts; warmup/eval servicing closes the two largest unserviced windows). |
+| S22 | Checkpoints round-trip ALL trajectory state of §3 — including every persistent term's sources + SRC_STEP moments + step/schedule counters — such that a resumed run continues the same trajectory (modulo RNG streams and kernel nondeterminism, cf. D4). **AMENDED 2026-07-06** (checkpoint split, pending Oli sign-off): each step stores TWO orbax items — `decomposition` (the trained product: V/U `components` + `ci_fn`) and `training` (the process: both optimizer states, persistent adversaries, `step`). `TrainState` COMPOSES exactly these two (`TrainState{decomposition, training}`) — one representation shared by the trainer and the checkpoint, so save/restore map onto its own fields with no regrouping. Trainer save/restore round-trips both (S22 semantics unchanged); fine-tune init (S33) and every consumer (harvest/autointerp/clustering/app) restore ONLY `decomposition`, with zero knowledge of training's optimizer/adversary structure. Both items save sharded (no on-loop full-gather). No in-code compat for pre-split single-`default`-item checkpoints: this is a clean format break — pre-split runs must be migrated by an ad-hoc script before a tip trainer/consumer can read them, not carried by a dual-read fallback. **SIGTERM lifecycle (decision):** a SLURM-delivered SIGTERM sets a flag that is serviced at the main train-step boundary (synchronous save of the completed step, then exit for requeue), inside the faith-warmup loop (clean exit WITHOUT a save — no valid checkpoint exists pre-step-0, and resume skips warmup whenever any checkpoint is present, so a partial step-0 save would resume as if fully warmed; the requeue redoes warmup), and inside the in-loop eval pass (abandon the partial pass unlogged, fall through to the step-boundary save). The **first-jit-compile window remains unserviced** — a SIGTERM there is honored only once compilation finishes and control reaches the next servicing point. Periodic `save_every` is the BACKSTOP guarantee for every window; the SIGTERM→save path is the low-latency fast path, not the sole guarantee (lore `jsp_sigterm_save_never_fired`: 6/6 historical preemptions fell back to periodic ckpts; warmup/eval servicing closes the two largest unserviced windows). |
 | S23 | A persistent source bundle feeds exactly ONE loss term. (The fused-backward S14′ unscaling divides that term's coeff out of the source gradient; a bundle shared across terms would make the division wrong.) |
 | S24 | A persistent term's WARMUP ascents forward all sites, routed everywhere — regardless of the term's loss plan (torch parity: `persistent_pgd_state.warmup` hardcodes route-all). A fresh-PGD entry draws its routing ONCE per step, shared by all its ascents and its main loss forward (torch parity: `pgd_masked_recon_loss_update`). |
 | S25 | Recon KL direction is `KL(softmax(clean_output) ‖ softmax(masked_output))` — `P = clean`, `Q = masked`. Equivalently `Σ p_clean · (log p_clean − log p_masked)`. Torch `recon_loss_kl` realizes this as `F.kl_div(log_softmax(pred=masked), softmax(target=clean), reduction='sum')` (`batch_and_loss_fns.py::recon_loss_kl`); reversing the arguments is a silent, plausible bug, so the direction is semantic, not incidental. |
@@ -302,7 +302,7 @@ faithfulness, sources/PPGD, the squashings (S5/S6), the imp-min reduction, the g
 | S29 | JAX `EvalConfig` carries `batch_size`, `every`, `n_steps`, `slow_every`, `slow_on_first_step`, `slow_n_batches_accum`, `density_heatmap_n_bins` (+ `rounding_threshold`, `ci_alive_threshold`). `slow_every` / `slow_on_first_step` are read from the canonical schema and drive the in-loop slow tier (S28); `slow_n_batches_accum` is the `CIHistograms.n_batches_accum` histogram-sample cap (None = uncapped); `density_heatmap_n_bins` is the `CIHistograms.density_heatmap_n_bins` opt-in for the per-token CI density heatmap (None = off — the add-on shares the slow-eval forward's `lower`, adds only an on-device per-component bincount `(C, n_bins + 1)` — column 0 = underflow (CI < 1e-9, incl. exact-0), columns 1..n_bins = log-spaced `[1e-9, 1]` bands — accumulated over EVERY batch, and emits `figures/ci_density_heatmap` on a log y-axis; no raw-value host transfer). `eval` is atomic-optional (`EvalConfig | None`): `None` disables in-loop eval (both tiers). **AMENDED 2026-06-18**: re-adds `slow_every` / `slow_on_first_step`, deliberately dropped in the original S29 when the slow tier was offline-only. **AMENDED 2026-07-01**: adds `density_heatmap_n_bins` (opt-in per-token CI density heatmap sharing the CIHistograms forward). |
 | 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 (`<ClassName>/<site>` + a combined `<ClassName>` = Σ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/<parent_step>` 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. |
+| 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 restores the parent's `ckpts/<parent_step>` `decomposition` item (the trained `components` (V/U) + `ci_fn`, cf. S22) onto the fresh reference's decomposition; 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. |
 
 ## 6. Variation points