diff --git a/CLAUDE.md b/CLAUDE.md
index f852c6063..fb1352662 100644
--- a/CLAUDE.md
+++ b/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
diff --git a/param_decomp/CLAUDE.md b/param_decomp/CLAUDE.md
index 7b5a6a180..22385d343 100644
--- a/param_decomp/CLAUDE.md
+++ b/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.
diff --git a/param_decomp/SPEC.md b/param_decomp/SPEC.md
index 23d0844c9..80aff6775 100644
--- a/param_decomp/SPEC.md
+++ b/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
 
diff --git a/param_decomp/checkpoint.py b/param_decomp/checkpoint.py
index 927fff386..d6990c2f5 100644
--- a/param_decomp/checkpoint.py
+++ b/param_decomp/checkpoint.py
@@ -1,17 +1,25 @@
 """Checkpoint / resume of the generic trainer's `TrainState` via orbax (SPEC S22).
 
-The whole trajectory — V/U + CI masters, both optimizer states, the persistent
-adversary (sources + its Adam moments), and the step counter — lives in `TrainState`
-as one pytree; orbax saves it **sharded** (every process writes its own shards, no
-full-gather on the training loop) and restores it onto the reference state's
-shardings. The frozen target is NOT saved (SPEC §3): resume rebuilds it from HF and
-loads only the trajectory.
+Each checkpoint step holds TWO orbax items, splitting the product from the process:
+
+- `decomposition` — `train.Decomposition` (V/U components + ci_fn), the trained product.
+  Every consumer (harvest/autointerp/clustering/app/fine-tune init) restores ONLY this
+  item, with zero knowledge of how training initializes its optimizers or adversaries.
+- `training` — `train.TrainingItem` (both optimizer states, the persistent adversaries,
+  the step counter), the trainer-only trajectory tail. Only trainer resume touches it.
+
+`TrainState` composes exactly these two — one representation — so save/restore map onto its
+own `.decomposition` / `.training` fields with no regrouping.
+
+Both items save **sharded** (every process writes its own shards, no full-gather on the
+training loop) and restore onto the reference state's shardings. The frozen target is
+NOT saved (SPEC §3): resume rebuilds it from HF and loads only the trajectory.
 
 Synchronous saves (no async): a SIGTERM-triggered save must be on disk before the
 process exits for SLURM requeue-resume.
 
 `init_from_parent` is the fine-tune entry (SPEC S33): a fresh run loads a PARENT
-checkpoint's V/U + ci_fn (the trained decomposition) but starts a clean schedule —
+checkpoint's `decomposition` (the trained product) but starts a clean schedule —
 fresh optimizer / sources, `step=0` — under a NEW config (changed LR / coeffs / steps,
 same component & ci-fn structure).
 """
@@ -21,10 +29,11 @@
 from typing import cast
 
 import jax
+import numpy as np
 import orbax.checkpoint as ocp
 from orbax.checkpoint.type_handlers import ArrayHandler, register_type_handler
 
-from param_decomp.train import TrainState
+from param_decomp.train import Decomposition, TrainState
 
 # Replica-parallel writes (multiple hosts cooperatively writing a REPLICATED array)
 # hit a Shard-internals incompatibility on multi-controller jax 0.10 and buy nothing
@@ -44,7 +53,13 @@ def make_checkpoint_manager(ckpt_dir: Path, keep_last: int) -> ocp.CheckpointMan
 
 
 def save_state(mgr: ocp.CheckpointManager, step: int, state: TrainState) -> None:
-    mgr.save(step, args=ocp.args.StandardSave(state))
+    mgr.save(
+        step,
+        args=ocp.args.Composite(
+            decomposition=ocp.args.StandardSave(state.decomposition),
+            training=ocp.args.StandardSave(state.training),
+        ),
+    )
     mgr.wait_until_finished()
 
 
@@ -52,7 +67,14 @@ def restore_step(mgr: ocp.CheckpointManager, reference: TrainState, step: int) -
     """Restore checkpoint `step` onto `reference`'s shapes/dtypes/shardings
     (a freshly-initialised, correctly-placed `TrainState`)."""
     abstract = jax.tree.map(ocp.utils.to_shape_dtype_struct, reference)
-    restored = mgr.restore(step, args=ocp.args.StandardRestore(abstract))
+    composite = mgr.restore(
+        step,
+        args=ocp.args.Composite(
+            decomposition=ocp.args.StandardRestore(abstract.decomposition),
+            training=ocp.args.StandardRestore(abstract.training),
+        ),
+    )
+    restored = TrainState(decomposition=composite["decomposition"], training=composite["training"])
     # Coerce the restored tree onto the reference's exact FORMAT (layout + sharding), not just its
     # sharding. StandardRestore already honors the sharding SPEC (verified), so a device_put onto
     # sharding alone is a no-op — but orbax-restored arrays carry a default memory LAYOUT that
@@ -73,28 +95,53 @@ def restore_latest(
     return restore_step(mgr, reference, step), step
 
 
+def restore_decomposition(
+    mgr: ocp.CheckpointManager, step: int, abstract: Decomposition
+) -> Decomposition:
+    """Restore ONLY the trained decomposition of checkpoint `step` onto `abstract`'s
+    shapes/dtypes/shardings (`to_shape_dtype_struct` of a correctly-placed reference)."""
+    composite = mgr.restore(
+        step, args=ocp.args.Composite(decomposition=ocp.args.StandardRestore(abstract))
+    )
+    return cast(Decomposition, composite["decomposition"])
+
+
+def restore_decomposition_to_host(
+    mgr: ocp.CheckpointManager, step: int, abstract: Decomposition
+) -> Decomposition:
+    """`restore_decomposition` for a consumer without the run's device topology: leaves
+    restore as HOST numpy (`abstract` needs no shardings — `jax.eval_shape` over
+    `run_state.init_decomposition` yields it with zero allocation), and the caller
+    `jax.device_put`s the result wherever it computes."""
+    host_args_for_leaf = lambda _: ocp.RestoreArgs(restore_type=np.ndarray)  # noqa: E731
+    restore_args = jax.tree.map(host_args_for_leaf, abstract)
+    composite = mgr.restore(
+        step,
+        args=ocp.args.Composite(
+            decomposition=ocp.args.PyTreeRestore(item=abstract, restore_args=restore_args)
+        ),
+    )
+    return cast(Decomposition, composite["decomposition"])
+
+
 def init_from_parent(parent_ckpt_dir: Path, parent_step: int, reference: TrainState) -> TrainState:
-    """Fine-tune init (SPEC S33): load the parent checkpoint's trained V/U + ci_fn ONTO
+    """Fine-tune init (SPEC S33): load the parent checkpoint's trained decomposition ONTO
     `reference` (a fresh-from-init `TrainState` built from the NEW config), and keep the
     fresh reference's optimizer states, persistent sources, and `step=0`.
 
-    Only the components and ci_fn carry over — the new schedule wants a clean Adam (no
-    stale momentum scale) and a fresh adversary; the schedule recomputes from step 0 over
-    the new `cfg.steps`. The full parent state is restored onto `reference`, which orbax
-    requires to be shape/dtype/sharding-identical to the parent's saved state: a mismatch
-    in component/ci_fn structure (different sites / C / ci-fn arch) fails the restore. The
-    config-level structural guard in `run.py` is the readable pre-check before this point."""
+    Only the decomposition carries over — the new schedule wants a clean Adam (no stale
+    momentum scale) and a fresh adversary; the schedule recomputes from step 0 over the
+    new `cfg.steps`. Orbax requires the reference's decomposition to be
+    shape/dtype/sharding-identical to the parent's saved one: a mismatch in
+    component/ci_fn structure (different sites / C / ci-fn arch) fails the restore. The
+    config-level structural guard in `run.py` is the readable pre-check before this
+    point. The parent's optimizer/adversary state is never read, so it may differ freely.
+    `reference.step` is kept as-is: it is already a GLOBAL replicated zero
+    (`_ensure_global`); re-creating it host-local would break the multi-host save."""
     parent_mgr = make_checkpoint_manager(parent_ckpt_dir, keep_last=1)
     assert parent_step in parent_mgr.all_steps(), (
         f"parent step {parent_step} not in {parent_ckpt_dir} (have {sorted(parent_mgr.all_steps())})"
     )
-    parent = restore_step(parent_mgr, reference, parent_step)
-    # Keep reference.step (already a GLOBAL replicated zero: init_train_state builds step=0,
-    # _ensure_global re-materializes it as a well-formed global array). Re-creating it as
-    # jnp.zeros((), jnp.int32) yields a HOST-LOCAL SingleDeviceSharding array that orbax
-    # refuses to serialize in a multi-host save.
-    return dataclasses.replace(
-        reference,
-        components=parent.components,
-        ci_fn=parent.ci_fn,
-    )
+    abstract = jax.tree.map(ocp.utils.to_shape_dtype_struct, reference.decomposition)
+    parent = restore_decomposition(parent_mgr, parent_step, abstract)
+    return dataclasses.replace(reference, decomposition=parent)
diff --git a/param_decomp/experiments/invariance_check.py b/param_decomp/experiments/invariance_check.py
index e30bf712f..14ce844d0 100644
--- a/param_decomp/experiments/invariance_check.py
+++ b/param_decomp/experiments/invariance_check.py
@@ -54,7 +54,7 @@
     mlp_family_site_cs,
 )
 from param_decomp.tests.test_llama8b import _tiny_cfg, _tiny_decomposed_lm
-from param_decomp.train import TrainState, make_train_step
+from param_decomp.train import Decomposition, TrainingItem, TrainState, make_train_step
 
 
 def _run(steps: int, sharded: bool) -> list[dict[str, float]]:
@@ -94,20 +94,22 @@ def _run(steps: int, sharded: bool) -> list[dict[str, float]]:
     )  # fmt: skip
     assert ppgd_cfg.coeff is not None
     state = TrainState(
-        components=vu, ci_fn=ci_fn,
-        components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
-        ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
-        adversaries={
-            ppgd_cfg.type: PersistentAdversary(
-                sources=src,
-                opt_state=init_sources_adam_state(src),
-                state_key=ppgd_cfg.type,
-                coeff=ppgd_cfg.coeff,
-                adam=ppgd_cfg.optimizer,
-                n_warmup=ppgd_cfg.n_warmup_steps,
-            )
-        },
-        step=jnp.zeros((), jnp.int32),
+        decomposition=Decomposition(components=vu, ci_fn=ci_fn),
+        training=TrainingItem(
+            components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
+            ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
+            adversaries={
+                ppgd_cfg.type: PersistentAdversary(
+                    sources=src,
+                    opt_state=init_sources_adam_state(src),
+                    state_key=ppgd_cfg.type,
+                    coeff=ppgd_cfg.coeff,
+                    adam=ppgd_cfg.optimizer,
+                    n_warmup=ppgd_cfg.n_warmup_steps,
+                )
+            },
+            step=jnp.zeros((), jnp.int32),
+        ),
     )  # fmt: skip
     loss_terms = build_loss_terms(
         (
diff --git a/param_decomp/run.py b/param_decomp/run.py
index b62cf163b..2f9f32602 100644
--- a/param_decomp/run.py
+++ b/param_decomp/run.py
@@ -365,7 +365,7 @@ def _init_or_restore_state(
     restored = restore_latest(checkpoint_manager, state)
     if restored is not None:
         state, ckpt_step = restored
-        assert int(state.step) == ckpt_step, (int(state.step), ckpt_step)
+        assert int(state.training.step) == ckpt_step, (int(state.training.step), ckpt_step)
         if is_main:
             print(f"resumed from checkpoint step {ckpt_step}", flush=True)
         return state, ckpt_step
@@ -392,10 +392,10 @@ def _init_or_restore_state(
             pd.faithfulness_warmup_lr, weight_decay=pd.faithfulness_warmup_weight_decay
         )
         faith_warmup_opt_state = faith_warmup_optimizer.init(
-            eqx.filter(state.components, eqx.is_array)
+            eqx.filter(state.decomposition.components, eqx.is_array)
         )
         faith_warmup_step = make_faith_warmup_step(faith_warmup_optimizer, compiler_options)
-        warmed_components = state.components
+        warmed_components = state.decomposition.components
         t0 = time.time()
         faith_warmup_loss = None
         for _ in range(pd.faithfulness_warmup_steps):
@@ -414,7 +414,9 @@ def _init_or_restore_state(
         jax.block_until_ready(faith_warmup_loss)
         new_opt_vu = _ensure_global(opt_vu.init(eqx.filter(warmed_components, eqx.is_array)), mesh)
         state = dataclasses.replace(
-            state, components=warmed_components, components_opt_state=new_opt_vu
+            state,
+            decomposition=dataclasses.replace(state.decomposition, components=warmed_components),
+            training=dataclasses.replace(state.training, components_opt_state=new_opt_vu),
         )
         if is_main:
             print(
@@ -554,7 +556,7 @@ def _bench_dispatch(tree: object, n: int = 15) -> float:
                 jax.block_until_ready(_ident(tree))
             return (time.perf_counter() - _b0) / n
 
-        _vu = state.components.vu
+        _vu = state.decomposition.components.vu
         _by_kind: dict[str, list[tuple[jax.Array, jax.Array]]] = _collections.defaultdict(list)
         for _name, _VU in _vu.items():
             _by_kind[_name.split(".")[-1]].append(_VU)
diff --git a/param_decomp/run_state.py b/param_decomp/run_state.py
index a20bfc436..a0a305e25 100644
--- a/param_decomp/run_state.py
+++ b/param_decomp/run_state.py
@@ -35,7 +35,7 @@
     init_decomp_vu_placed,
     init_sources_sharded,
 )
-from param_decomp.train import TrainState
+from param_decomp.train import Decomposition, TrainingItem, TrainState
 
 
 def optax_schedule(config: ScheduleConfig, total_steps: int) -> Callable[[ArrayLike], Array]:
@@ -98,6 +98,23 @@ def build_optimizers(pd: PDConfig):
     return opt_vu, opt_ci, (sched_vu, sched_ci)
 
 
+def init_decomposition(
+    lm: DecomposedModel, ci_fn_arch: CIFnArch, init_key: PRNGKeyArray, mesh: Mesh
+) -> Decomposition:
+    """The trained-product half of `init_train_state`, factored out so a consumer can
+    `jax.eval_shape` it to recover the saved `decomposition` item's tree structure
+    without building (or knowing about) the optimizers/adversaries."""
+    ci_key = random.fold_in(init_key, 1)
+    # Placement is MODEL-OWNED: V/U + CI declare their own per-leaf shardings (asserting
+    # divisibility), uniformly across mesh sizes — no scale inference, no replicate fallback.
+    components = init_decomp_vu_placed(lm.sites, init_key, mesh)
+    ci_fn = init_ci_fn_placed(ci_fn_arch, lm.sites, ci_key, mesh)
+    assert ci_fn.expects_axes == lm.leading_axes, (
+        f"CI fn expects leading axes {ci_fn.expects_axes} but model has {lm.leading_axes}"
+    )
+    return Decomposition(components=components, ci_fn=ci_fn)
+
+
 def init_train_state(
     pd: PDConfig,
     lm: DecomposedModel,
@@ -109,14 +126,8 @@ def init_train_state(
     src_key: PRNGKeyArray,
     mesh: Mesh,
 ) -> TrainState:
-    ci_key = random.fold_in(init_key, 1)
-    # Placement is MODEL-OWNED: V/U + CI declare their own per-leaf shardings (asserting
-    # divisibility), uniformly across mesh sizes — no scale inference, no replicate fallback.
-    components = init_decomp_vu_placed(lm.sites, init_key, mesh)
-    ci_fn = init_ci_fn_placed(ci_fn_arch, lm.sites, ci_key, mesh)
-    assert ci_fn.expects_axes == lm.leading_axes, (
-        f"CI fn expects leading axes {ci_fn.expects_axes} but model has {lm.leading_axes}"
-    )
+    decomposition = init_decomposition(lm, ci_fn_arch, init_key, mesh)
+    components, ci_fn = decomposition.components, decomposition.ci_fn
     losses = build_loss_terms(pd.loss_metrics, lm.site_names)
     persistent = persistent_configs(losses.recon)
     term_coeff_by_state_key = {
@@ -154,10 +165,11 @@ def init_train_state(
                 n_warmup=cfg.n_warmup_steps,
             )
     return TrainState(
-        components=components,
-        ci_fn=ci_fn,
-        components_opt_state=opt_vu.init(eqx.filter(components, eqx.is_array)),
-        ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
-        adversaries=adversaries,
-        step=jnp.zeros((), jnp.int32),
+        decomposition=decomposition,
+        training=TrainingItem(
+            components_opt_state=opt_vu.init(eqx.filter(components, eqx.is_array)),
+            ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
+            adversaries=adversaries,
+            step=jnp.zeros((), jnp.int32),
+        ),
     )
diff --git a/param_decomp/slow_eval.py b/param_decomp/slow_eval.py
index f3301bed6..7e3f6cc8f 100644
--- a/param_decomp/slow_eval.py
+++ b/param_decomp/slow_eval.py
@@ -806,16 +806,26 @@ def compute_hidden_acts_metrics(
 ) -> dict[str, float]:
     """Both hidden-acts recon eval metrics over the eval batches (`input_batches` holds
     the model's target-specific inputs — an LM's token ids), keyed by the torch
-    `<ClassName>[/<site>]` log keys. `state.components`/`state.ci_fn` are the restored
+    `<ClassName>[/<site>]` log keys. `state.decomposition.components`/`state.decomposition.ci_fn` are the restored
     trajectory; `base_key` seeds the stochastic variant's per-batch draws."""
     ci_key, stoch_key = random.split(base_key)
     ci_step = make_ci_hidden_acts_step(model, compiler_options)
     ci_reductions = accumulate_hidden_acts(
-        ci_step, model, state.components, state.ci_fn, input_batches, ci_key
+        ci_step,
+        model,
+        state.decomposition.components,
+        state.decomposition.ci_fn,
+        input_batches,
+        ci_key,
     )
     stoch_step = make_stochastic_hidden_acts_step(model, n_mask_samples, compiler_options)
     stoch_reductions = accumulate_hidden_acts(
-        stoch_step, model, state.components, state.ci_fn, input_batches, stoch_key
+        stoch_step,
+        model,
+        state.decomposition.components,
+        state.decomposition.ci_fn,
+        input_batches,
+        stoch_key,
     )
     return {
         **hidden_acts_log_entries("CIHiddenActsReconLoss", ci_reductions),
diff --git a/param_decomp/tests/stacked_parity/test_stacked_parity.py b/param_decomp/tests/stacked_parity/test_stacked_parity.py
index 4d05ec4d6..c68380caf 100644
--- a/param_decomp/tests/stacked_parity/test_stacked_parity.py
+++ b/param_decomp/tests/stacked_parity/test_stacked_parity.py
@@ -55,7 +55,7 @@
     mlp_family_site_cs,
 )
 from param_decomp.tests.test_llama8b import _tiny_cfg
-from param_decomp.train import TrainState, make_train_step
+from param_decomp.train import Decomposition, TrainingItem, TrainState, make_train_step
 from vendored_jax.llama import llama3_inv_freq
 
 FIXTURES = Path(__file__).resolve().parent / "stacked_fixtures.npz"
@@ -248,20 +248,22 @@ def test_train_trajectory_matches():
     )
     assert ppgd_cfg.coeff is not None
     state = TrainState(
-        components=vu, ci_fn=ci_fn,
-        components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
-        ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
-        adversaries={
-            ppgd_cfg.type: PersistentAdversary(
-                sources=sources,
-                opt_state=init_sources_adam_state(sources),
-                state_key=ppgd_cfg.type,
-                coeff=ppgd_cfg.coeff,
-                adam=ppgd_cfg.optimizer,
-                n_warmup=ppgd_cfg.n_warmup_steps,
-            )
-        },
-        step=jnp.zeros((), jnp.int32),
+        decomposition=Decomposition(components=vu, ci_fn=ci_fn),
+        training=TrainingItem(
+            components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
+            ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
+            adversaries={
+                ppgd_cfg.type: PersistentAdversary(
+                    sources=sources,
+                    opt_state=init_sources_adam_state(sources),
+                    state_key=ppgd_cfg.type,
+                    coeff=ppgd_cfg.coeff,
+                    adam=ppgd_cfg.optimizer,
+                    n_warmup=ppgd_cfg.n_warmup_steps,
+                )
+            },
+            step=jnp.zeros((), jnp.int32),
+        ),
     )  # fmt: skip
     loss_terms = build_loss_terms(
         (
@@ -299,11 +301,11 @@ def test_train_trajectory_matches():
                 f"step{step_idx} {fixture_key}: per-site {got!r} vs stacked {want!r}"
             )
 
-    assert isinstance(state.components, DecompVU)
+    assert isinstance(state.decomposition.components, DecompVU)
     for name in lm.site_names:
-        V, U = state.components.site(name)
+        V, U = state.decomposition.components.site(name)
         _assert_close(V, f[f"out::final_V::{name}"], f"final V {name}")
         _assert_close(U, f[f"out::final_U::{name}"], f"final U {name}")
-    final_sources = state.adversaries["PersistentPGDReconLoss"].sources
+    final_sources = state.training.adversaries["PersistentPGDReconLoss"].sources
     for name in lm.site_names:
         _assert_close(final_sources[name], f[f"out::final_src::{name}"], f"final src {name}")
diff --git a/param_decomp/tests/test_checkpoint.py b/param_decomp/tests/test_checkpoint.py
index ea4cddbf2..979cdb335 100644
--- a/param_decomp/tests/test_checkpoint.py
+++ b/param_decomp/tests/test_checkpoint.py
@@ -17,7 +17,9 @@
     sources_adam_ascend_project,
 )
 from param_decomp.checkpoint import (
+    init_from_parent,
     make_checkpoint_manager,
+    restore_decomposition_to_host,
     restore_latest,
     restore_step,
     save_state,
@@ -51,7 +53,7 @@
     init_sources_sharded,
 )
 from param_decomp.tests.test_llama8b import _tiny_cfg, _tiny_decomposed_lm
-from param_decomp.train import TrainState, make_train_step
+from param_decomp.train import Decomposition, TrainingItem, TrainState, make_train_step
 from vendored_jax.llama import LlamaConfig
 
 
@@ -112,11 +114,13 @@ def _build(seed: int):
     )
     ppgd_cfg = _ppgd_cfg(n_warmup=1)
     state = TrainState(
-        components=vu, ci_fn=ci_fn,
-        components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
-        ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
-        adversaries={ppgd_cfg.type: _adversary(src, ppgd_cfg)},
-        step=jnp.zeros((), jnp.int32),
+        decomposition=Decomposition(components=vu, ci_fn=ci_fn),
+        training=TrainingItem(
+            components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
+            ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
+            adversaries={ppgd_cfg.type: _adversary(src, ppgd_cfg)},
+            step=jnp.zeros((), jnp.int32),
+        ),
     )  # fmt: skip
     loss_terms = build_loss_terms(
         (
@@ -178,7 +182,7 @@ def test_persistent_adam_step_count_roundtrip_and_post_resume_bias_correction(tm
     for i in range(3):
         state, _ = step(lm, state, resid, jax.random.PRNGKey(i))
 
-    pre_save = state.adversaries[state_key].opt_state
+    pre_save = state.training.adversaries[state_key].opt_state
     n_ascents = int(pre_save.step_count)
     # Each train step runs n_warmup_steps (1) supplemental ascents + 1 final ascent.
     assert n_ascents == 3 * (1 + 1)
@@ -190,10 +194,10 @@ def test_persistent_adam_step_count_roundtrip_and_post_resume_bias_correction(tm
     restored = restore_latest(mgr, fresh)
     assert restored is not None
     loaded, _ = restored
-    loaded_adam = loaded.adversaries[state_key].opt_state
+    loaded_adam = loaded.training.adversaries[state_key].opt_state
 
     # (a) the step_count leaf survived the round-trip: present, fp32 scalar, value N.
-    assert state_key in loaded.adversaries
+    assert state_key in loaded.training.adversaries
     assert loaded_adam.step_count.dtype == jnp.float32
     assert loaded_adam.step_count.shape == ()
     assert float(loaded_adam.step_count) == float(n_ascents)
@@ -207,7 +211,7 @@ def test_persistent_adam_step_count_roundtrip_and_post_resume_bias_correction(tm
     adam_cfg = AdamPGDConfig(
         beta1=beta1, beta2=beta2, lr_schedule=ScheduleConfig(start_val=0.01, warmup_pct=0.025)
     )
-    loaded_sources = loaded.adversaries[state_key].sources
+    loaded_sources = loaded.training.adversaries[state_key].sources
     grads = {site: jnp.ones_like(v) for site, v in loaded_sources.items()}
     _, post_resume = sources_adam_ascend_project(
         loaded_sources, grads, loaded_adam, jnp.asarray(0.01), adam_cfg
@@ -230,6 +234,77 @@ def test_no_checkpoint_returns_none(tmp_path: Path):
     assert restore_latest(mgr, fresh) is None
 
 
+def test_saved_layout_is_two_items(tmp_path: Path):
+    """The on-disk item names are the cross-version contract every consumer keys on —
+    pin them."""
+    _, state, _, _ = _build(seed=1)
+    mgr = make_checkpoint_manager(tmp_path / "ckpts", keep_last=2)
+    save_state(mgr, 0, state)
+    step_dir = tmp_path / "ckpts" / "0"
+    assert (step_dir / "decomposition").is_dir()
+    assert (step_dir / "training").is_dir()
+    assert not (step_dir / "default").exists()
+
+
+def test_consumer_restores_decomposition_to_host(tmp_path: Path):
+    """The consumer path (`open_jax_run`): an eval_shape'd `Decomposition` abstract —
+    no optimizer/adversary knowledge — restores host-side, bit-equal to the saved
+    components + ci_fn."""
+    lm, state, step, resid = _build(seed=1)
+    for i in range(2):
+        state, _ = step(lm, state, resid, jax.random.PRNGKey(i))
+    mgr = make_checkpoint_manager(tmp_path / "ckpts", keep_last=2)
+    save_state(mgr, 2, state)
+
+    # A FRESH manager, as every real consumer opens: orbax pins an item's handler per
+    # manager instance, so the saving manager can't PyTreeRestore what it StandardSave'd.
+    consumer_mgr = make_checkpoint_manager(tmp_path / "ckpts", keep_last=2)
+    abstract = jax.eval_shape(lambda: state.decomposition)
+    restored = restore_decomposition_to_host(consumer_mgr, 2, abstract)
+    for a, b in zip(
+        jax.tree.leaves(state.decomposition),
+        jax.tree.leaves(restored),
+        strict=True,
+    ):
+        assert jnp.array_equal(jnp.asarray(a), jnp.asarray(b))
+
+
+def test_init_from_parent_restores_decomposition_only(tmp_path: Path):
+    """Fine-tune init (S33): `init_from_parent` restores ONLY the parent's
+    `decomposition` item (components + ci_fn) onto a fresh reference, keeping the fresh
+    reference's optimizer states / adversaries / `step=0` — never reading the parent's
+    `training` item, which may differ freely."""
+    lm, parent, step, resid = _build(seed=1)
+    for i in range(2):
+        parent, _ = step(lm, parent, resid, jax.random.PRNGKey(i))
+    mgr = make_checkpoint_manager(tmp_path / "ckpts", keep_last=2)
+    save_state(mgr, 2, parent)
+
+    # A fresh reference from a DIFFERENT seed: its components/ci_fn, optimizer state and
+    # adversaries all differ from the parent's, so the carry-over vs keep-fresh split is
+    # observable on every leaf.
+    _, fresh, _, _ = _build(seed=7)
+    finetuned = init_from_parent(tmp_path / "ckpts", parent_step=2, reference=fresh)
+
+    # decomposition carries over from the parent...
+    for a, b in zip(
+        jax.tree.leaves(finetuned.decomposition),
+        jax.tree.leaves(parent.decomposition),
+        strict=True,
+    ):
+        assert jnp.array_equal(jnp.asarray(a), jnp.asarray(b))
+    # ...while optimizer state, adversaries and step stay the fresh reference's.
+    for a, b in zip(
+        jax.tree.leaves(
+            (finetuned.training.components_opt_state, finetuned.training.ci_fn_opt_state)
+        ),
+        jax.tree.leaves((fresh.training.components_opt_state, fresh.training.ci_fn_opt_state)),
+        strict=True,
+    ):
+        assert jnp.array_equal(jnp.asarray(a), jnp.asarray(b))
+    assert int(finetuned.training.step) == 0
+
+
 def _build_sharded(seed: int, mesh: Mesh):
     """A `TrainState` placed exactly as the production trainer places it
     (`run_state.init_train_state`): C-sharded V/U + ci_fn, replicated sources, over the
@@ -258,11 +333,13 @@ def _build_sharded(seed: int, mesh: Mesh):
     )
     ppgd_cfg = _ppgd_cfg(n_warmup=1)
     state = TrainState(
-        components=vu, ci_fn=ci_fn,
-        components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
-        ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
-        adversaries={ppgd_cfg.type: _adversary(src, ppgd_cfg)},
-        step=jnp.asarray(7, jnp.int32),
+        decomposition=Decomposition(components=vu, ci_fn=ci_fn),
+        training=TrainingItem(
+            components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
+            ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
+            adversaries={ppgd_cfg.type: _adversary(src, ppgd_cfg)},
+            step=jnp.asarray(7, jnp.int32),
+        ),
     )  # fmt: skip
     return state
 
@@ -283,7 +360,7 @@ def test_sharded_roundtrip_bit_equal(tmp_path: Path):
     # The big V/U + ci_fn + sources leaves must be genuinely C-sharded over the mesh
     # (the multi-shard write path); only the small scalars (step) stay single-device.
     n_named = sum(isinstance(x.sharding, NamedSharding) for x in jax.tree.leaves(state))
-    assert n_named >= len(jax.tree.leaves(state.components)), n_named
+    assert n_named >= len(jax.tree.leaves(state.decomposition.components)), n_named
 
     mgr = make_checkpoint_manager(tmp_path / "ckpts", keep_last=2)
     save_state(mgr, 3, state)
diff --git a/param_decomp/tests/test_checkpoint_production_topology.py b/param_decomp/tests/test_checkpoint_production_topology.py
index 85b3ef3a3..aae8db4f6 100644
--- a/param_decomp/tests/test_checkpoint_production_topology.py
+++ b/param_decomp/tests/test_checkpoint_production_topology.py
@@ -58,7 +58,7 @@
     init_sources_sharded,
 )
 from param_decomp.tests.test_llama8b import _tiny_cfg, _tiny_decomposed_lm
-from param_decomp.train import TrainState, make_train_step
+from param_decomp.train import Decomposition, TrainingItem, TrainState, make_train_step
 
 # Needs >1 jax device (production topology); hangs at the default 1 device, so gated behind
 # --runmultidevice. Run via `make test-multidevice` (simulated CPU devices). See conftest.
@@ -134,11 +134,13 @@ def _build_sharded(seed: int):
             n_warmup=ppgd_cfg.n_warmup_steps,
         )
     state = TrainState(
-        components=vu, ci_fn=ci_fn,
-        components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
-        ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
-        adversaries=adversaries,
-        step=jnp.zeros((), jnp.int32),
+        decomposition=Decomposition(components=vu, ci_fn=ci_fn),
+        training=TrainingItem(
+            components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
+            ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
+            adversaries=adversaries,
+            step=jnp.zeros((), jnp.int32),
+        ),
     )  # fmt: skip
     # run.py:185 — normalize eager scalar stragglers (Adam `count`, `step`) onto the mesh
     # so the whole state is global-array placed, exactly as a production run / restore.
@@ -196,7 +198,7 @@ def test_sharded_roundtrip_persists_source_moments(tmp_path: Path):
     for i in range(2):
         state, _ = step(lm, state, resid, jax.random.PRNGKey(i))
     # The ascents must have advanced each term's Adam counter before we save.
-    _assert_moments_present(state.adversaries)
+    _assert_moments_present(state.training.adversaries)
 
     mgr = make_checkpoint_manager(tmp_path / "ckpts", keep_last=2)
     save_state(mgr, 2, state)
@@ -210,7 +212,7 @@ def test_sharded_roundtrip_persists_source_moments(tmp_path: Path):
     assert ckpt_step == 2
 
     # The persisted pytree itself carries the moments + step_count for every term.
-    _assert_moments_present(loaded.adversaries)
+    _assert_moments_present(loaded.training.adversaries)
 
     # Values come from disk, and each leaf is reconstructed onto the REFERENCE sharding.
     # Pull leaves to host before comparing: a sharded leaf and a single-device leaf can't
diff --git a/param_decomp/tests/test_finetune_resume.py b/param_decomp/tests/test_finetune_resume.py
index 772ab6b97..231384b05 100644
--- a/param_decomp/tests/test_finetune_resume.py
+++ b/param_decomp/tests/test_finetune_resume.py
@@ -29,7 +29,7 @@ def test_init_from_parent_loads_components_resets_schedule(tmp_path: Path):
     lm, parent_state, step, resid = _build(seed=1)
     for i in range(2):
         parent_state, _ = step(lm, parent_state, resid, jax.random.PRNGKey(i))
-    assert int(parent_state.step) == 2
+    assert int(parent_state.training.step) == 2
 
     parent_ckpt_dir = tmp_path / "parent" / "ckpts"
     mgr = make_checkpoint_manager(parent_ckpt_dir, keep_last=2)
@@ -42,30 +42,34 @@ def test_init_from_parent_loads_components_resets_schedule(tmp_path: Path):
 
     # components + ci_fn come from the parent.
     for a, b in zip(
-        jax.tree.leaves(finetuned.components), jax.tree.leaves(parent_state.components), strict=True
+        jax.tree.leaves(finetuned.decomposition.components),
+        jax.tree.leaves(parent_state.decomposition.components),
+        strict=True,
     ):
         assert jnp.array_equal(a, b)
     for a, b in zip(
-        jax.tree.leaves(finetuned.ci_fn), jax.tree.leaves(parent_state.ci_fn), strict=True
+        jax.tree.leaves(finetuned.decomposition.ci_fn),
+        jax.tree.leaves(parent_state.decomposition.ci_fn),
+        strict=True,
     ):
         assert jnp.array_equal(a, b)
 
     # step resets to 0 for the fresh schedule.
-    assert int(finetuned.step) == 0
-    assert finetuned.step.dtype == jnp.int32
+    assert int(finetuned.training.step) == 0
+    assert finetuned.training.step.dtype == jnp.int32
 
     # optimizer states + sources are the FRESH reference's (not the parent's). The fresh
     # sources are RNG-drawn from seed 7; the parent's from seed 1 — they must differ.
-    for state_key, fresh_adv in fresh.adversaries.items():
+    for state_key, fresh_adv in fresh.training.adversaries.items():
         for site, arr in fresh_adv.sources.items():
-            assert jnp.array_equal(finetuned.adversaries[state_key].sources[site], arr)
+            assert jnp.array_equal(finetuned.training.adversaries[state_key].sources[site], arr)
             assert not jnp.array_equal(
-                finetuned.adversaries[state_key].sources[site],
-                parent_state.adversaries[state_key].sources[site],
+                finetuned.training.adversaries[state_key].sources[site],
+                parent_state.training.adversaries[state_key].sources[site],
             )
     for a, b in zip(
-        jax.tree.leaves(finetuned.components_opt_state),
-        jax.tree.leaves(fresh.components_opt_state),
+        jax.tree.leaves(finetuned.training.components_opt_state),
+        jax.tree.leaves(fresh.training.components_opt_state),
         strict=True,
     ):
         assert jnp.array_equal(a, b)
diff --git a/param_decomp/tests/test_generic_model_io.py b/param_decomp/tests/test_generic_model_io.py
index 7640df7b3..af8129d3c 100644
--- a/param_decomp/tests/test_generic_model_io.py
+++ b/param_decomp/tests/test_generic_model_io.py
@@ -41,7 +41,7 @@
 from param_decomp.lm import DecomposedModel, run_stochastic_masked_output
 from param_decomp.recon import build_loss_terms
 from param_decomp.schedule import ScheduleConfig
-from param_decomp.train import TrainState, make_train_step
+from param_decomp.train import Decomposition, TrainingItem, TrainState, make_train_step
 
 B, T, D, C = 2, 3, 8, 5
 K_COORDS, M_AUX = 4, 2
@@ -221,12 +221,13 @@ def _initial_state(lm: DecomposedModel, components: DecompVU, ci_arch: Chunkwise
     opt_ci = optax.adamw(1e-2, weight_decay=0.0)
     ci_fn = build_ci_fn(ci_arch, lm.sites, random.PRNGKey(11))
     state = TrainState(
-        components=components,
-        ci_fn=ci_fn,
-        components_opt_state=opt_vu.init(eqx.filter(components, eqx.is_array)),
-        ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
-        adversaries={},
-        step=jnp.zeros((), jnp.int32),
+        decomposition=Decomposition(components=components, ci_fn=ci_fn),
+        training=TrainingItem(
+            components_opt_state=opt_vu.init(eqx.filter(components, eqx.is_array)),
+            ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
+            adversaries={},
+            step=jnp.zeros((), jnp.int32),
+        ),
     )
     return state, opt_vu, opt_ci
 
@@ -271,13 +272,15 @@ def test_train_step_runs_through_generic_target():
         mesh=None,
     )
 
-    V_before = jax.device_get(state.components.site(SITE)[0])  # host copy survives step donation
+    V_before = jax.device_get(
+        state.decomposition.components.site(SITE)[0]
+    )  # host copy survives step donation
     run_key = random.PRNGKey(3)
     for step_idx in range(2):
         state, metrics = step_fn(lm, state, inputs, random.fold_in(run_key, step_idx))
         assert jnp.isfinite(metrics["total"]), (step_idx, metrics["total"])
         assert "loss/StochasticReconLoss" in metrics
 
-    assert not jnp.allclose(state.components.site(SITE)[0], V_before), (
+    assert not jnp.allclose(state.decomposition.components.site(SITE)[0], V_before), (
         "V did not move — step is a no-op"
     )
diff --git a/param_decomp/tests/test_llama8b.py b/param_decomp/tests/test_llama8b.py
index a22f6ed80..3e9954972 100644
--- a/param_decomp/tests/test_llama8b.py
+++ b/param_decomp/tests/test_llama8b.py
@@ -48,7 +48,13 @@
     parse_site_name,
     site_name,
 )
-from param_decomp.train import TrainState, make_faith_warmup_step, make_train_step
+from param_decomp.train import (
+    Decomposition,
+    TrainingItem,
+    TrainState,
+    make_faith_warmup_step,
+    make_train_step,
+)
 from vendored_jax.llama import LlamaConfig, llama3_inv_freq
 
 
@@ -357,20 +363,22 @@ def test_step_trains_and_has_vpd_signature(site_cs: tuple[SiteC, ...]):
     )
     assert ppgd_cfg.coeff is not None
     state = TrainState(
-        components=vu, ci_fn=ci_fn,
-        components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
-        ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
-        adversaries={
-            ppgd_cfg.type: PersistentAdversary(
-                sources=src,
-                opt_state=init_sources_adam_state(src),
-                state_key=ppgd_cfg.type,
-                coeff=ppgd_cfg.coeff,
-                adam=ppgd_cfg.optimizer,
-                n_warmup=ppgd_cfg.n_warmup_steps,
-            )
-        },
-        step=jnp.zeros((), jnp.int32),
+        decomposition=Decomposition(components=vu, ci_fn=ci_fn),
+        training=TrainingItem(
+            components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
+            ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
+            adversaries={
+                ppgd_cfg.type: PersistentAdversary(
+                    sources=src,
+                    opt_state=init_sources_adam_state(src),
+                    state_key=ppgd_cfg.type,
+                    coeff=ppgd_cfg.coeff,
+                    adam=ppgd_cfg.optimizer,
+                    n_warmup=ppgd_cfg.n_warmup_steps,
+                )
+            },
+            step=jnp.zeros((), jnp.int32),
+        ),
     )  # fmt: skip
     loss_terms = build_loss_terms(
         (
@@ -405,9 +413,9 @@ def test_step_trains_and_has_vpd_signature(site_cs: tuple[SiteC, ...]):
         losses.append({k: float(v) for k, v in m.items()})
 
     assert all(jnp.isfinite(jnp.array(list(m.values()))).all() for m in losses)
-    assert int(state.step) == n_steps
+    assert int(state.training.step) == n_steps
     # SPEC S13: n_warmup + 1 source-Adam updates per training step, moments persist.
-    ppgd_adv = state.adversaries["PersistentPGDReconLoss"]
+    ppgd_adv = state.training.adversaries["PersistentPGDReconLoss"]
     assert float(ppgd_adv.opt_state.step_count) == n_steps * (n_warmup + 1)
     # SPEC S15: sources stay projected to [0,1].
     for v in ppgd_adv.sources.values():
@@ -415,11 +423,11 @@ def test_step_trains_and_has_vpd_signature(site_cs: tuple[SiteC, ...]):
     # SPEC S9: p annealed below its 2.0 start by step 4 of 100.
     assert losses[-1]["p_imp"] < 2.0
     # fp32 masters preserved through updates (SPEC N1).
-    assert isinstance(state.components, DecompVU)
-    for V, U in state.components.vu.values():
+    assert isinstance(state.decomposition.components, DecompVU)
+    for V, U in state.decomposition.components.vu.values():
         assert V.dtype == jnp.float32 and U.dtype == jnp.float32
-    assert isinstance(state.ci_fn, ChunkwiseTransformerCIFn)
-    assert state.ci_fn.chunks.in_proj_w.dtype == jnp.float32
+    assert isinstance(state.decomposition.ci_fn, ChunkwiseTransformerCIFn)
+    assert state.decomposition.ci_fn.chunks.in_proj_w.dtype == jnp.float32
 
 
 def test_faith_warmup_decreases_faith():
@@ -479,11 +487,13 @@ def make_state() -> TrainState:
         vu = init_decomp_vu(sites, jax.random.PRNGKey(1))
         ci_fn = _build_chunkwise_ci_fn(lm, jax.random.PRNGKey(2), n_blocks=1)
         return TrainState(
-            components=vu, ci_fn=ci_fn,
-            components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
-            ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
-            adversaries={},
-            step=jnp.zeros((), jnp.int32),
+            decomposition=Decomposition(components=vu, ci_fn=ci_fn),
+            training=TrainingItem(
+                components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
+                ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
+                adversaries={},
+                step=jnp.zeros((), jnp.int32),
+            ),
         )  # fmt: skip
 
     def run_step(n_ascent_steps: int) -> tuple[TrainState, dict[str, jax.Array]]:
@@ -531,8 +541,8 @@ def run_step(n_ascent_steps: int) -> tuple[TrainState, dict[str, jax.Array]]:
             ]
         )
     ).all()
-    assert state.adversaries == {}, "fresh adversary carries no persistent sources"
-    assert int(state.step) == 1
+    assert state.training.adversaries == {}, "fresh adversary carries no persistent sources"
+    assert int(state.training.step) == 1
 
     _, metrics_unascended = run_step(n_ascent_steps=0)
     assert float(metrics["loss/PGDReconLoss"]) >= float(metrics_unascended["loss/PGDReconLoss"]), (
diff --git a/param_decomp/tests/test_llama_simple_mlp.py b/param_decomp/tests/test_llama_simple_mlp.py
index 567e524ea..6aa01105a 100644
--- a/param_decomp/tests/test_llama_simple_mlp.py
+++ b/param_decomp/tests/test_llama_simple_mlp.py
@@ -49,7 +49,13 @@
     site_name,
     site_specs,
 )
-from param_decomp.train import TrainState, make_faith_warmup_step, make_train_step
+from param_decomp.train import (
+    Decomposition,
+    TrainingItem,
+    TrainState,
+    make_faith_warmup_step,
+    make_train_step,
+)
 
 
 def _tiny_cfg() -> LlamaSimpleMLPConfig:
@@ -375,20 +381,22 @@ def test_step_trains_and_has_vpd_signature():
     )
     assert ppgd_cfg.coeff is not None
     state = TrainState(
-        components=vu, ci_fn=ci_fn,
-        components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
-        ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
-        adversaries={
-            ppgd_cfg.type: PersistentAdversary(
-                sources=src,
-                opt_state=init_sources_adam_state(src),
-                state_key=ppgd_cfg.type,
-                coeff=ppgd_cfg.coeff,
-                adam=ppgd_cfg.optimizer,
-                n_warmup=ppgd_cfg.n_warmup_steps,
-            )
-        },
-        step=jnp.zeros((), jnp.int32),
+        decomposition=Decomposition(components=vu, ci_fn=ci_fn),
+        training=TrainingItem(
+            components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
+            ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
+            adversaries={
+                ppgd_cfg.type: PersistentAdversary(
+                    sources=src,
+                    opt_state=init_sources_adam_state(src),
+                    state_key=ppgd_cfg.type,
+                    coeff=ppgd_cfg.coeff,
+                    adam=ppgd_cfg.optimizer,
+                    n_warmup=ppgd_cfg.n_warmup_steps,
+                )
+            },
+            step=jnp.zeros((), jnp.int32),
+        ),
     )  # fmt: skip
     loss_terms = build_loss_terms(
         (
@@ -423,9 +431,9 @@ def test_step_trains_and_has_vpd_signature():
         losses.append({k: float(v) for k, v in m.items()})
 
     assert all(jnp.isfinite(jnp.array(list(m.values()))).all() for m in losses)
-    assert int(state.step) == n_steps
+    assert int(state.training.step) == n_steps
     # SPEC S13: n_warmup + 1 source-Adam updates per training step, moments persist.
-    ppgd_adv = state.adversaries["PersistentPGDReconLoss"]
+    ppgd_adv = state.training.adversaries["PersistentPGDReconLoss"]
     assert float(ppgd_adv.opt_state.step_count) == n_steps * (n_warmup + 1)
     # SPEC S15: sources stay projected to [0,1].
     for v in ppgd_adv.sources.values():
@@ -433,11 +441,11 @@ def test_step_trains_and_has_vpd_signature():
     # SPEC S9: p annealed below its 2.0 start by step 4 of 100.
     assert losses[-1]["p_imp"] < 2.0
     # fp32 masters preserved through updates (SPEC N1).
-    assert isinstance(state.components, DecompVU)
-    for V, U in state.components.vu.values():
+    assert isinstance(state.decomposition.components, DecompVU)
+    for V, U in state.decomposition.components.vu.values():
         assert V.dtype == jnp.float32 and U.dtype == jnp.float32
-    assert isinstance(state.ci_fn, ChunkwiseTransformerCIFn)
-    assert state.ci_fn.chunks.in_proj_w.dtype == jnp.float32
+    assert isinstance(state.decomposition.ci_fn, ChunkwiseTransformerCIFn)
+    assert state.decomposition.ci_fn.chunks.in_proj_w.dtype == jnp.float32
 
 
 def test_faith_warmup_decreases_faith():
diff --git a/param_decomp/tests/test_no_bake_invariant.py b/param_decomp/tests/test_no_bake_invariant.py
index 92f809081..dd5302681 100644
--- a/param_decomp/tests/test_no_bake_invariant.py
+++ b/param_decomp/tests/test_no_bake_invariant.py
@@ -20,7 +20,7 @@
 from param_decomp.recon import build_loss_terms
 from param_decomp.schedule import ScheduleConfig
 from param_decomp.tests.test_generic_model_io import SyntheticDecomposedModel
-from param_decomp.train import TrainState, make_train_step
+from param_decomp.train import Decomposition, TrainingItem, TrainState, make_train_step
 
 SITE = "block.0.proj"
 B, T, D, C = 2, 3, 32, 5
@@ -63,12 +63,13 @@ def _build_step_and_args():
     opt_vu = optax.adamw(1e-2, weight_decay=0.0)
     opt_ci = optax.adamw(1e-2, weight_decay=0.0)
     state = TrainState(
-        components=components,
-        ci_fn=ci_fn,
-        components_opt_state=opt_vu.init(eqx.filter(components, eqx.is_array)),
-        ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
-        adversaries={},
-        step=jnp.zeros((), jnp.int32),
+        decomposition=Decomposition(components=components, ci_fn=ci_fn),
+        training=TrainingItem(
+            components_opt_state=opt_vu.init(eqx.filter(components, eqx.is_array)),
+            ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
+            adversaries={},
+            step=jnp.zeros((), jnp.int32),
+        ),
     )
     loss_terms = build_loss_terms(
         (
diff --git a/param_decomp/tools/strip_checkpoint.py b/param_decomp/tools/strip_checkpoint.py
new file mode 100644
index 000000000..d1f43e718
--- /dev/null
+++ b/param_decomp/tools/strip_checkpoint.py
@@ -0,0 +1,163 @@
+"""Strip a finished run's pre-split orbax checkpoint to its decomposition item.
+
+A pre-split trainer checkpoint is one orbax item (`default`) dominated by
+resume-only state: the persistent-PGD adversary (sources + their Adam moments)
+and the V/U + CI-fn optimizer moments are ~90% of the bytes; the trained
+decomposition (`components` + `ci_fn`) that consumers actually restore is the
+small rest. Stripping rewrites `ckpts/<step>` into the split two-item layout
+(SPEC S22) with ONLY the `decomposition` item — this is the pre-split-run
+migration the split's clean format break requires, minus the training item.
+Consumers (`restore_decomposition*`) load a stripped checkpoint natively; the
+run can no longer warm-resume training — only strip finished or abandoned runs.
+
+Dry-run by default: prints the kept/dropped byte split from checkpoint metadata.
+`--apply` stages the stripped checkpoint next to `ckpts/`, verifies its tree
+against the original's metadata, then swaps it in. (If interrupted between the
+swap's delete and rename, the staged copy survives under `ckpts_strip_staging/`.)
+
+    python -m param_decomp.tools.strip_checkpoint <run_dir> [--step N] [--apply]
+"""
+
+import argparse
+import math
+import shutil
+import subprocess
+from collections.abc import Mapping
+from pathlib import Path
+from typing import Any, cast
+
+import jax
+import numpy as np
+import orbax.checkpoint as ocp
+
+DECOMPOSITION_KEYS = ("components", "ci_fn")
+DROPPABLE_KEYS = frozenset(
+    {
+        "step",
+        "adversaries",
+        "sources",
+        "sources_opt_state",
+        "components_opt_state",
+        "ci_fn_opt_state",
+    }
+)
+
+
+def open_manager(ckpts_dir: Path) -> ocp.CheckpointManager:
+    return ocp.CheckpointManager(
+        ckpts_dir.resolve(),
+        options=ocp.CheckpointManagerOptions(enable_async_checkpointing=False),
+    )
+
+
+def leaf_bytes(tree: Any) -> int:
+    leaves = [leaf for leaf in jax.tree.leaves(tree) if hasattr(leaf, "shape")]
+    assert leaves, "metadata subtree has no array leaves"
+    return sum(math.prod(leaf.shape) * np.dtype(leaf.dtype).itemsize for leaf in leaves)
+
+
+def tree_shapes(tree: Any) -> Any:
+    return jax.tree.map(lambda leaf: (tuple(leaf.shape), str(np.dtype(leaf.dtype))), tree)
+
+
+def du_bytes(path: Path) -> int:
+    du = subprocess.run(["du", "-sb", str(path)], capture_output=True, text=True, check=True)
+    return int(du.stdout.split()[0])
+
+
+def read_tree_metadata(item_dir: Path) -> Mapping[str, Any]:
+    metadata = ocp.PyTreeCheckpointer().metadata(item_dir).item_metadata
+    assert metadata is not None, f"no pytree metadata under {item_dir}"
+    return cast(Mapping[str, Any], metadata)
+
+
+def restore_decomposition_subtrees(
+    manager: ocp.CheckpointManager, step: int, metadata: Mapping[str, Any]
+) -> dict[str, Any]:
+    """Partial-restore only the decomposition subtrees of a pre-split (single
+    `default`-item) checkpoint, as host numpy (no shardings)."""
+    item = {
+        key: jax.tree.map(lambda m: jax.ShapeDtypeStruct(m.shape, m.dtype), metadata[key])
+        for key in DECOMPOSITION_KEYS
+    }
+    restore_args = jax.tree.map(lambda _: ocp.RestoreArgs(restore_type=np.ndarray), item)
+    restored = manager.restore(
+        step,
+        args=ocp.args.PyTreeRestore(item=item, restore_args=restore_args, partial_restore=True),
+    )
+    return cast(dict[str, Any], restored)
+
+
+def main() -> None:
+    assert __doc__ is not None
+    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
+    parser.add_argument("run_dir", type=Path)
+    parser.add_argument("--step", type=int, help="checkpoint step (default: latest)")
+    parser.add_argument(
+        "--apply", action="store_true", help="actually strip (default: dry-run report)"
+    )
+    args = parser.parse_args()
+
+    ckpts_dir = args.run_dir / "ckpts"
+    assert ckpts_dir.is_dir(), f"no ckpts dir under {args.run_dir}"
+    manager = open_manager(ckpts_dir)
+    step = args.step if args.step is not None else manager.latest_step()
+    assert step is not None, f"no checkpoints under {ckpts_dir}"
+    assert step in manager.all_steps(), f"step {step} not in {sorted(manager.all_steps())}"
+    assert (ckpts_dir / str(step) / "default").is_dir(), (
+        f"{ckpts_dir / str(step)} is not a pre-split single-item checkpoint "
+        "(already split/stripped?)"
+    )
+
+    metadata = read_tree_metadata(ckpts_dir / str(step) / "default")
+    present: set[str] = set(metadata.keys())
+    keep: set[str] = set(DECOMPOSITION_KEYS)
+    assert keep <= present, f"checkpoint lacks {keep - present}"
+    dropped_keys = present - keep
+    assert dropped_keys <= DROPPABLE_KEYS, (
+        f"unknown training-state keys: {dropped_keys - DROPPABLE_KEYS}"
+    )
+
+    kept = {key: leaf_bytes(metadata[key]) for key in DECOMPOSITION_KEYS}
+    dropped = {key: leaf_bytes(metadata[key]) for key in sorted(dropped_keys) if key != "step"}
+    on_disk = du_bytes(ckpts_dir / str(step))
+    total_logical = sum(kept.values()) + sum(dropped.values())
+    print(f"{args.run_dir.name} ckpts/{step}: {on_disk / 2**30:.1f} GiB on disk")
+    for label, group in (("keep", kept), ("drop", dropped)):
+        for key, size in sorted(group.items(), key=lambda kv: -kv[1]):
+            print(f"  {label}  {size / 2**30:9.2f} GiB  {key}")
+    kept_frac = sum(kept.values()) / total_logical
+    print(
+        f"stripped size ~{on_disk * kept_frac / 2**30:.1f} GiB "
+        f"({kept_frac:.1%} of logical bytes kept)"
+    )
+    if not args.apply:
+        print("dry run — pass --apply to strip")
+        return
+
+    decomposition = restore_decomposition_subtrees(manager, step, metadata)
+    staging_dir = args.run_dir / "ckpts_strip_staging"
+    assert not staging_dir.exists(), (
+        f"leftover staging dir from an interrupted strip: {staging_dir}"
+    )
+    staging_manager = open_manager(staging_dir)
+    staging_manager.save(
+        step, args=ocp.args.Composite(decomposition=ocp.args.StandardSave(decomposition))
+    )
+    staging_manager.wait_until_finished()
+
+    staged_metadata = read_tree_metadata(staging_dir / str(step) / "decomposition")
+    assert set(staged_metadata.keys()) == keep, staged_metadata.keys()
+    for key in DECOMPOSITION_KEYS:
+        assert tree_shapes(staged_metadata[key]) == tree_shapes(metadata[key]), (
+            f"staged {key} tree does not match the original"
+        )
+
+    shutil.rmtree(ckpts_dir / str(step))
+    (staging_dir / str(step)).rename(ckpts_dir / str(step))
+    shutil.rmtree(staging_dir)
+    print(f"stripped: {on_disk / 2**30:.1f} -> {du_bytes(ckpts_dir / str(step)) / 2**30:.1f} GiB")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/param_decomp/train.py b/param_decomp/train.py
index 6694d2157..6408f257d 100644
--- a/param_decomp/train.py
+++ b/param_decomp/train.py
@@ -63,9 +63,21 @@ def cast_floating(tree: Any, dtype: Any) -> Any:
 
 @jax.tree_util.register_dataclass
 @dataclass(frozen=True)
-class TrainState:
+class Decomposition:
+    """The trained PRODUCT: V/U components + the CI fn (fp32 masters). Checkpointed as
+    its own orbax item so consumers (harvest/autointerp/clustering/app) restore it with
+    zero knowledge of the training process (optimizer states, adversaries, step)."""
+
     components: DecompVU  # the universal trainable V/U pytree, fp32 masters
     ci_fn: CIFn  # fp32 masters
+
+
+@jax.tree_util.register_dataclass
+@dataclass(frozen=True)
+class TrainingItem:
+    """The trainer-only trajectory tail: both optimizer states, the persistent adversaries,
+    the step counter. Checkpointed as its own orbax item — no consumer restores it."""
+
     components_opt_state: optax.OptState
     ci_fn_opt_state: optax.OptState
     adversaries: dict[str, PersistentAdversary]
@@ -75,6 +87,17 @@ class TrainState:
     step: Array
 
 
+@jax.tree_util.register_dataclass
+@dataclass(frozen=True)
+class TrainState:
+    """The full training pytree, composed of the two checkpoint items so there is ONE
+    representation: `decomposition` is the trained product, `training` the trajectory tail.
+    Save/restore maps directly onto these two fields — no regrouping."""
+
+    decomposition: Decomposition
+    training: TrainingItem
+
+
 def _grad_norm_metrics(components_grad: DecompVU, ci_fn_grad: Any) -> dict[str, Array]:
     """Pre-clip gradient L2 norms, matching the torch `component_grad_norms` families:
     per-leaf `grad_norms/components<path>` / `grad_norms/ci_fns<path>` (paths are this
@@ -258,14 +281,16 @@ def step(
         batch: Any,
         key: PRNGKeyArray,
     ) -> tuple[TrainState, dict[str, Array]]:
-        step_f32 = state.step.astype(jnp.float32)
+        decomposition = state.decomposition
+        training = state.training
+        step_f32 = training.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)
+            taps = model.read_activations(batch, decomposition.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.
@@ -274,7 +299,7 @@ def step(
         # ── adversary ascents: params + CI detached (SPEC §4.5) ──
         prepared, recon_vjp = jax.vjp(
             lambda c: model.prepare_compute_weights(cast_floating(c, COMPUTE_DT)),
-            state.components,
+            decomposition.components,
         )
         prepared_detached = jax.lax.stop_gradient(prepared)
         prepared_ascend = replicate_for_ascend(prepared_detached)
@@ -286,7 +311,7 @@ def step(
         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,
+                decomposition.ci_fn,
             )
         ci_lower_detached = jax.lax.stop_gradient(ci).lower
 
@@ -304,7 +329,7 @@ def warmup_scoring_loss(sources: dict[str, Array]) -> Array:
         with jax.named_scope("pd_pgd_warmup_ascend"):
             warmed_advs = {
                 state_key: adv.warmup_ascend(warmup_scoring_loss, step_f32, total_steps)
-                for state_key, adv in state.adversaries.items()
+                for state_key, adv in training.adversaries.items()
             }
 
         # Fresh-PGD entries: ONE routing draw per entry per step, shared by all
@@ -463,7 +488,7 @@ def pre_built_fwd(
                 eqx.filter_value_and_grad(loss_fn, has_aux=True)(
                     (
                         prepared,
-                        state.components,
+                        decomposition.components,
                         ci,
                         {k: a.sources for k, a in warmed_advs.items()},
                     )
@@ -491,22 +516,23 @@ def pre_built_fwd(
 
         components_updates, new_components_opt_state = components_optimizer.update(
             components_grad,
-            state.components_opt_state,
-            eqx.filter(state.components, eqx.is_array),
+            training.components_opt_state,
+            eqx.filter(decomposition.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)
+            ci_fn_grad, training.ci_fn_opt_state, eqx.filter(decomposition.ci_fn, eqx.is_array)
         )
-        new_components = eqx.apply_updates(state.components, components_updates)
-        new_ci_fn = eqx.apply_updates(state.ci_fn, ci_fn_updates)
+        new_components = eqx.apply_updates(decomposition.components, components_updates)
+        new_ci_fn = eqx.apply_updates(decomposition.ci_fn, ci_fn_updates)
 
         new_state = TrainState(
-            components=new_components,
-            ci_fn=new_ci_fn,
-            components_opt_state=new_components_opt_state,
-            ci_fn_opt_state=new_ci_fn_opt_state,
-            adversaries=new_adversaries,
-            step=state.step + 1,
+            decomposition=Decomposition(components=new_components, ci_fn=new_ci_fn),
+            training=TrainingItem(
+                components_opt_state=new_components_opt_state,
+                ci_fn_opt_state=new_ci_fn_opt_state,
+                adversaries=new_adversaries,
+                step=training.step + 1,
+            ),
         )
         metrics = {
             "total": total_loss,
@@ -518,7 +544,7 @@ def pre_built_fwd(
             **grad_norm_metrics,
         }
         source_lrs = {
-            k: adv.source_lr(step_f32, total_steps) for k, adv in state.adversaries.items()
+            k: adv.source_lr(step_f32, total_steps) for k, adv in training.adversaries.items()
         }
         if len(source_lrs) == 1:
             metrics["src_lr"] = next(iter(source_lrs.values()))
diff --git a/param_decomp_lab/experiments/lm/load_run.py b/param_decomp_lab/experiments/lm/load_run.py
index c064d8591..20920bf30 100644
--- a/param_decomp_lab/experiments/lm/load_run.py
+++ b/param_decomp_lab/experiments/lm/load_run.py
@@ -3,8 +3,9 @@
 
 This is the reusable "load a JAX run" pattern. It reads a run dir
 (`runs/<p-id>/{launch_config.yaml, ckpts/}`), rebuilds the frozen
-target + `DecomposedModel` from the pinned config, restores the orbax checkpoint onto a
-reference `TrainState`, and exposes the pure forward a consumer needs:
+target + `DecomposedModel` from the pinned config, restores the checkpoint's
+`decomposition` item (the trained V/U + ci_fn — optimizer/adversary state is training's
+business and is never touched), and exposes the pure forward a consumer needs:
 
     run = open_jax_run(run_dir)                 # latest checkpoint
     fwd = run.forward(token_ids)                # one frozen, forward-only pass
@@ -33,11 +34,11 @@
 from jaxtyping import Array, Float, Int
 
 from param_decomp.built_run import BuiltRun
-from param_decomp.checkpoint import make_checkpoint_manager, restore_latest, restore_step
+from param_decomp.checkpoint import make_checkpoint_manager, restore_decomposition_to_host
 from param_decomp.ci_fn import CIFn
 from param_decomp.components import DecompVU
 from param_decomp.lm import DecomposedModel
-from param_decomp.run_state import build_optimizers, init_train_state
+from param_decomp.run_state import init_decomposition
 from param_decomp.sharding import hsdp_mesh, place_via_shardings
 from param_decomp.targets import llama_simple_mlp
 from param_decomp.targets.llama8b import (
@@ -46,7 +47,7 @@
     load_decomposed_lm_from_hf,
 )
 from param_decomp.targets.llama8b_sharding import place_target
-from param_decomp.train import COMPUTE_DT, TrainState, cast_floating
+from param_decomp.train import COMPUTE_DT, Decomposition, cast_floating
 from param_decomp_lab.experiments.lm.config import (
     LlamaSimpleMLPTargetConfig,
     TargetConfig,
@@ -120,7 +121,7 @@ class LoadedJaxRun:
     lm: DecomposedModel
     config: BuiltRun
     vocab_size: int
-    _state: TrainState
+    _decomposition: Decomposition
     _forward: Callable[
         [DecomposedModel, DecompVU, CIFn, Int[Array, "B T"]],
         tuple[dict[str, Array], dict[str, Array], Array],
@@ -137,10 +138,10 @@ def layer_activation_sizes(self) -> list[tuple[str, int]]:
         return [(s.name, s.C) for s in self.lm.sites]
 
     def forward(self, token_ids: Int[Array, "B T"]) -> HarvestForward:
-        ci_fn = self._state.ci_fn
+        ci_fn = self._decomposition.ci_fn
         assert isinstance(ci_fn, CIFn), "harvest is the transformer-CI-fn (LM) path only"
         lower_leaky_ci, component_acts, output_probs = self._forward(
-            self.lm, self._state.components, ci_fn, token_ids
+            self.lm, self._decomposition.components, ci_fn, token_ids
         )
         return HarvestForward(
             lower_leaky_ci=lower_leaky_ci,
@@ -149,30 +150,40 @@ def forward(self, token_ids: Int[Array, "B T"]) -> HarvestForward:
         )
 
 
-def open_jax_run(run_dir: Path, step: int | None = None) -> LoadedJaxRun:
-    """Open the run at `run_dir`; restore checkpoint `step` (latest if None)."""
-    cfg = load_run_dir_config(run_dir)
-    mesh = hsdp_mesh()
-    lm, vocab_size = build_target(cfg, mesh)
+def _restore_decomposition(
+    cfg: BuiltRun, lm: DecomposedModel, mesh: jax.sharding.Mesh, run_dir: Path, step: int | None
+) -> tuple[Decomposition, int]:
+    """Restore ONLY the trained decomposition from the run's checkpoint.
 
-    opt_vu, opt_ci, _ = build_optimizers(cfg.pd)
-    init_key, src_key = jax.random.split(jax.random.PRNGKey(cfg.pd.seed))
-    reference = init_train_state(
-        cfg.pd, lm, cfg.ci_fn, cfg.data, opt_vu, opt_ci, init_key, src_key, mesh
-    )
+    Consumers never need the optimizer moments or persistent-PGD adversary sources
+    training also checkpoints — and on a single device those dominate: an 8B run's
+    sources + Adam state materialize ~60GB at init and again at restore staging, which
+    wedges/OOMs one GPU. `jax.eval_shape` over `init_decomposition` yields the saved
+    decomposition's structure with ZERO allocation (and zero knowledge of training's
+    optimizers); leaves restore as host numpy, then `device_put` onto the consumer's
+    single default device."""
+    init_key, _ = jax.random.split(jax.random.PRNGKey(cfg.pd.seed))
+    abstract = jax.eval_shape(lambda: init_decomposition(lm, cfg.ci_fn, init_key, mesh))
 
     assert cfg.cadence.keep_last_n_checkpoints is not None, cfg.cadence
     manager = make_checkpoint_manager(run_dir / "ckpts", cfg.cadence.keep_last_n_checkpoints)
-    if step is None:
-        restored = restore_latest(manager, reference)
-        assert restored is not None, f"no checkpoints under {run_dir / 'ckpts'}"
-        state, resolved_step = restored
-    else:
-        state, resolved_step = restore_step(manager, reference, step), step
-    assert isinstance(state.components, DecompVU)
+    resolved_step = manager.latest_step() if step is None else step
+    assert resolved_step is not None, f"no checkpoints under {run_dir / 'ckpts'}"
+    decomposition = jax.device_put(restore_decomposition_to_host(manager, resolved_step, abstract))
+    return decomposition, resolved_step
+
+
+def open_jax_run(run_dir: Path, step: int | None = None) -> LoadedJaxRun:
+    """Open the run at `run_dir`; restore checkpoint `step` (latest if None). Restores
+    only the trained decomposition (see `_restore_decomposition`)."""
+    cfg = load_run_dir_config(run_dir)
+    mesh = hsdp_mesh()
+    lm, vocab_size = build_target(cfg, mesh)
+    decomposition, resolved_step = _restore_decomposition(cfg, lm, mesh, run_dir, step)
+    assert isinstance(decomposition.components, DecompVU)
 
     site_names = lm.site_names
-    u_norms = _u_norms(state.components, site_names)
+    u_norms = _u_norms(decomposition.components, site_names)
 
     # `model` is the filter_jit ARG (frozen weights traced, not baked). It embeds the token
     # ids internally — the harvest forward feeds tokens straight in.
@@ -209,7 +220,7 @@ def forward(
         lm=lm,
         config=cfg,
         vocab_size=vocab_size,
-        _state=state,
+        _decomposition=decomposition,
         _forward=forward,
     )
 
diff --git a/param_decomp_lab/experiments/lm/run.py b/param_decomp_lab/experiments/lm/run.py
index 872c711b1..5875479c6 100644
--- a/param_decomp_lab/experiments/lm/run.py
+++ b/param_decomp_lab/experiments/lm/run.py
@@ -281,7 +281,9 @@ def eval_fn(state: TrainState, now_step: int) -> "LogRecord":
             eval_batches.append(eval_tokens)
             # fold values >= pd.steps never collide with the train step keys
             eval_key = random.fold_in(run_key, pd.steps + eval_pass_index * eval.n_steps + j)
-            eval_metrics = eval_step_fn(lm, state.components, state.ci_fn, eval_tokens, eval_key)
+            eval_metrics = eval_step_fn(
+                lm, state.decomposition.components, state.decomposition.ci_fn, eval_tokens, eval_key
+            )
             for k, v in eval_metrics.items():
                 metric_sums[k] = metric_sums.get(k, jnp.zeros(())) + v
         eval_record: dict[str, LogValue] = {
@@ -292,7 +294,12 @@ def eval_fn(state: TrainState, now_step: int) -> "LogRecord":
             # is summed over distributions, divided by their count.
             attn_key = random.fold_in(run_key, 2 * pd.steps + eval_pass_index)
             reductions = accumulate_attn_patterns(
-                attn_step, lm, state.components, state.ci_fn, eval_batches, attn_key
+                attn_step,
+                lm,
+                state.decomposition.components,
+                state.decomposition.ci_fn,
+                eval_batches,
+                attn_key,
             )
             eval_record |= {
                 f"eval/loss/{k}": v
@@ -307,7 +314,11 @@ def eval_fn(state: TrainState, now_step: int) -> "LogRecord":
             # hidden-acts scalars ride the live `_step` axis through `eval_record`; the
             # figures' pure-host render + wandb.log happen OFF the loop on rank 0.
             site_reductions = accumulate_site_reductions(
-                slow_eval_step, lm, state.ci_fn, eval_batches, eval.slow_n_batches_accum
+                slow_eval_step,
+                lm,
+                state.decomposition.ci_fn,
+                eval_batches,
+                eval.slow_n_batches_accum,
             )
             hidden_acts_key = random.fold_in(run_key, 3 * pd.steps + eval_pass_index)
             hidden_acts = compute_hidden_acts_metrics(
@@ -321,7 +332,7 @@ def eval_fn(state: TrainState, now_step: int) -> "LogRecord":
             position_ci: dict[str, PositionCI] | None = None
             if position_ci_step is not None:
                 position_ci = accumulate_position_ci(
-                    position_ci_step, lm, state.ci_fn, eval_batches
+                    position_ci_step, lm, state.decomposition.ci_fn, eval_batches
                 )
                 identity_ci_errors = compute_identity_ci_errors(
                     perm_spec, position_ci, IDENTITY_CI_ERROR_TOLERANCE
@@ -335,7 +346,7 @@ def eval_fn(state: TrainState, now_step: int) -> "LogRecord":
             if perm_spec.want_uv_plots:
                 components = {
                     name: (np.asarray(V), np.asarray(U))
-                    for name, (V, U) in state.components.vu.items()
+                    for name, (V, U) in state.decomposition.components.vu.items()
                 }
             slow_renderer.submit(site_reductions, perm_spec, position_ci, components, now_step)
         if is_main and built.run.wandb is not None:
diff --git a/param_decomp_lab/experiments/resid_mlp/run.py b/param_decomp_lab/experiments/resid_mlp/run.py
index 526296395..2162d60e0 100644
--- a/param_decomp_lab/experiments/resid_mlp/run.py
+++ b/param_decomp_lab/experiments/resid_mlp/run.py
@@ -155,10 +155,10 @@ def single_feature_ci(
     uv_spec = toy_uv_eval.toy_uv_spec(lm, raw_cfg)
 
     def eval_fn(state: TrainState, now_step: int) -> dict[str, float]:
-        ci_lower, ci_upper = single_feature_ci(lm, state.ci_fn)
+        ci_lower, ci_upper = single_feature_ci(lm, state.decomposition.ci_fn)
         toy_uv_eval.log_uv_figure(
             uv_spec,
-            state.components.vu,
+            state.decomposition.components.vu,
             ci_upper,
             now_step,
             wandb_active=built.run.wandb is not None,
diff --git a/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py b/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py
index 3545eff3c..71d88d454 100644
--- a/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py
+++ b/param_decomp_lab/experiments/resid_mlp/test_resid_mlp.py
@@ -32,7 +32,13 @@
 from param_decomp.lm import DecomposedModel
 from param_decomp.recon import build_loss_terms
 from param_decomp.schedule import ScheduleConfig
-from param_decomp.train import TrainState, make_faith_warmup_step, make_train_step
+from param_decomp.train import (
+    Decomposition,
+    TrainingItem,
+    TrainState,
+    make_faith_warmup_step,
+    make_train_step,
+)
 from param_decomp_lab.experiments.resid_mlp.model import (
     ResidMLPConfig,
     ResidMLPTarget,
@@ -247,10 +253,12 @@ def _make_state_and_step(
     opt_vu = optax.chain(optax.clip_by_global_norm(0.01), optax.adamw(1e-3, weight_decay=0.0))
     opt_ci = optax.adamw(1e-3, weight_decay=0.0)
     state = TrainState(
-        components=vu, ci_fn=ci_fn,
-        components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
-        ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
-        adversaries={}, step=jnp.zeros((), jnp.int32),
+        decomposition=Decomposition(components=vu, ci_fn=ci_fn),
+        training=TrainingItem(
+            components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
+            ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
+            adversaries={}, step=jnp.zeros((), jnp.int32),
+        ),
     )  # fmt: skip
     loss_terms = build_loss_terms(_loss_metrics(), lm.site_names)
     step = make_train_step(
@@ -274,10 +282,10 @@ def test_step_trains_positionless_no_persistent_sources():
         state, m = step(lm, state, x @ target.W_E, jax.random.PRNGKey(100 + i))
         losses.append({k: float(v) for k, v in m.items()})
     assert all(jnp.isfinite(jnp.array(list(m.values()))).all() for m in losses)
-    assert int(state.step) == 6
-    assert state.adversaries == {}  # no persistent sources for the stochastic configs
-    assert isinstance(state.components, DecompVU)
-    for V, U in state.components.vu.values():
+    assert int(state.training.step) == 6
+    assert state.training.adversaries == {}  # no persistent sources for the stochastic configs
+    assert isinstance(state.decomposition.components, DecompVU)
+    for V, U in state.decomposition.components.vu.values():
         assert V.dtype == jnp.float32 and U.dtype == jnp.float32
 
 
@@ -353,10 +361,12 @@ def _faith_warmed_state(
     opt_vu = optax.chain(optax.clip_by_global_norm(0.01), optax.adamw(1e-3, weight_decay=0.0))
     opt_ci = optax.adamw(1e-3, weight_decay=0.0)
     state = TrainState(
-        components=vu, ci_fn=ci_fn,
-        components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
-        ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
-        adversaries={}, step=jnp.zeros((), jnp.int32),
+        decomposition=Decomposition(components=vu, ci_fn=ci_fn),
+        training=TrainingItem(
+            components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
+            ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
+            adversaries={}, step=jnp.zeros((), jnp.int32),
+        ),
     )  # fmt: skip
     loss_terms = build_loss_terms(_recovery_loss_metrics(), lm.site_names)
     step = make_train_step(
@@ -403,7 +413,7 @@ def test_end_to_end_pretrain_decompose_recovers_identity():
         totals.append(float(m["total"]))
     assert totals[-1] < totals[0], (totals[0], totals[-1])
 
-    ci_lower = single_feature_ci(lm, state.ci_fn, n_features=5)
+    ci_lower = single_feature_ci(lm, state.decomposition.ci_fn, n_features=5)
     err = identity_ci_error(ci_lower["layers.0.mlp_in"], tolerance=0.2)
     assert err == 0, (
         f"mlp_in did not recover identity (err={err}):\n{jnp.round(ci_lower['layers.0.mlp_in'], 2)}"
diff --git a/param_decomp_lab/experiments/tms/run.py b/param_decomp_lab/experiments/tms/run.py
index 56e60a80e..1f43e50c4 100644
--- a/param_decomp_lab/experiments/tms/run.py
+++ b/param_decomp_lab/experiments/tms/run.py
@@ -138,10 +138,10 @@ def single_feature_ci(
     uv_spec = toy_uv_eval.toy_uv_spec(lm, raw_cfg)
 
     def eval_fn(state: TrainState, now_step: int) -> dict[str, float]:
-        ci_lower, ci_upper = single_feature_ci(lm, state.ci_fn)
+        ci_lower, ci_upper = single_feature_ci(lm, state.decomposition.ci_fn)
         toy_uv_eval.log_uv_figure(
             uv_spec,
-            state.components.vu,
+            state.decomposition.components.vu,
             ci_upper,
             now_step,
             wandb_active=built.run.wandb is not None,
diff --git a/param_decomp_lab/experiments/tms/test_tms.py b/param_decomp_lab/experiments/tms/test_tms.py
index 89932c322..fceffd990 100644
--- a/param_decomp_lab/experiments/tms/test_tms.py
+++ b/param_decomp_lab/experiments/tms/test_tms.py
@@ -26,7 +26,13 @@
 from param_decomp.lm import DecomposedModel
 from param_decomp.recon import build_loss_terms
 from param_decomp.schedule import ScheduleConfig
-from param_decomp.train import TrainState, make_faith_warmup_step, make_train_step
+from param_decomp.train import (
+    Decomposition,
+    TrainingItem,
+    TrainState,
+    make_faith_warmup_step,
+    make_train_step,
+)
 from param_decomp_lab.experiments.tms.model import (
     HiddenLayerInit,
     TMSConfig,
@@ -197,10 +203,12 @@ def _make_state_and_step(
     opt_vu = optax.chain(optax.clip_by_global_norm(0.01), optax.adamw(1e-3, weight_decay=0.0))
     opt_ci = optax.adamw(1e-3, weight_decay=0.0)
     state = TrainState(
-        components=vu, ci_fn=ci_fn,
-        components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
-        ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
-        adversaries={}, step=jnp.zeros((), jnp.int32),
+        decomposition=Decomposition(components=vu, ci_fn=ci_fn),
+        training=TrainingItem(
+            components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
+            ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
+            adversaries={}, step=jnp.zeros((), jnp.int32),
+        ),
     )  # fmt: skip
     loss_terms = build_loss_terms(_loss_metrics(), lm.site_names)
     step = make_train_step(
@@ -224,12 +232,12 @@ def test_step_trains_positionless_no_persistent_sources():
         state, m = step(lm, state, x, jax.random.PRNGKey(100 + i))
         losses.append({k: float(v) for k, v in m.items()})
     assert all(jnp.isfinite(jnp.array(list(m.values()))).all() for m in losses)
-    assert int(state.step) == 6
+    assert int(state.training.step) == 6
     # no persistent sources for the TMS stochastic configs
-    assert state.adversaries == {}
+    assert state.training.adversaries == {}
     # fp32 masters preserved
-    assert isinstance(state.components, DecompVU)
-    for V, U in state.components.vu.values():
+    assert isinstance(state.decomposition.components, DecompVU)
+    for V, U in state.decomposition.components.vu.values():
         assert V.dtype == jnp.float32 and U.dtype == jnp.float32
 
 
@@ -297,10 +305,12 @@ def _faith_warmed_state(
     opt_vu = optax.chain(optax.clip_by_global_norm(0.01), optax.adamw(1e-3, weight_decay=0.0))
     opt_ci = optax.adamw(1e-3, weight_decay=0.0)
     state = TrainState(
-        components=vu, ci_fn=ci_fn,
-        components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
-        ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
-        adversaries={}, step=jnp.zeros((), jnp.int32),
+        decomposition=Decomposition(components=vu, ci_fn=ci_fn),
+        training=TrainingItem(
+            components_opt_state=opt_vu.init(eqx.filter(vu, eqx.is_array)),
+            ci_fn_opt_state=opt_ci.init(eqx.filter(ci_fn, eqx.is_array)),
+            adversaries={}, step=jnp.zeros((), jnp.int32),
+        ),
     )  # fmt: skip
     loss_terms = build_loss_terms(_recovery_loss_metrics(), lm.site_names)
     step = make_train_step(
@@ -345,7 +355,7 @@ def test_end_to_end_pretrain_decompose_recovers_identity():
         totals.append(float(m["total"]))
     assert totals[-1] < totals[0], (totals[0], totals[-1])
 
-    ci_lower = single_feature_ci(lm, state.ci_fn, n_features=5)
+    ci_lower = single_feature_ci(lm, state.decomposition.ci_fn, n_features=5)
     err1 = identity_ci_error(ci_lower["linear1"], tolerance=0.2)
     err2 = identity_ci_error(ci_lower["linear2"], tolerance=0.2)
     assert err1 == 0, (