From 030bc99dcd50b0ad6373b89e447a1e107d5fd646 Mon Sep 17 00:00:00 2001 From: oli Date: Mon, 6 Jul 2026 10:01:45 +0000 Subject: [PATCH 1/5] feat(checkpoint): split checkpoints into decomposition | training items (S22/S33 amended) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each step now stores two orbax items: decomposition (V/U components + ci_fn — the trained product) and training (both optimizer states, persistent adversaries, step — the process). Consumers (open_jax_run, fine-tune init S33) restore ONLY decomposition via restore_decomposition[_to_host] against an eval_shape'd run_state.init_decomposition reference — zero knowledge of training's optimizer/adversary structure, and none of the ~60GB opt-moment/bsc-source materialization that wedged single-GPU harvest restores (supersedes the interim partial-restore 1ca895860, which still needed eval_shape(init_train_state) incl. build_optimizers). Runs saved before the split (single default item) keep writing that layout from their immutable workspaces; every restore path dual-reads both layouts (_stored_as_split, grounded in the stored dir layout: /default vs /decomposition). Both items save sharded — the no-on-loop-full-gather property is unchanged. SPEC S22/S33 amended (pending Oli sign-off). New tests: two-item layout pin, consumer host-restore, legacy dual-read across restore_latest / restore_decomposition_to_host / init_from_parent. Validated vs real stored data: p-594db290 (8B L18, 60GB-class ckpt) opens on CPU in 22s. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01EwzrwmBiMpkvDovpGgVysn --- CLAUDE.md | 2 +- param_decomp/CLAUDE.md | 9 +- param_decomp/SPEC.md | 4 +- param_decomp/checkpoint.py | 186 +++++++++++++++--- .../pile_ppgd_bsc_RESUMESMOKE_dp8.yaml | 152 ++++++++++++++ param_decomp/run_state.py | 29 ++- param_decomp/tests/test_checkpoint.py | 84 +++++++- param_decomp/train.py | 11 ++ param_decomp_lab/experiments/lm/load_run.py | 65 +++--- 9 files changed, 475 insertions(+), 67 deletions(-) create mode 100644 param_decomp/configs/pile_ppgd_bsc_RESUMESMOKE_dp8.yaml diff --git a/CLAUDE.md b/CLAUDE.md index f852c6063..27ad9f253 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// 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//... # orbax sharded checkpoints (JAX trainer) + ckpts//{decomposition,training}/ # orbax sharded checkpoints (JAX trainer): the trained product vs the trainer-only tail; pre-split runs have a single default/ item instead 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..50bd97937 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); pre-split runs' single `default` item stays dual-readable (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..0c7c3f6bc 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`). 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). Runs launched before the split store one `default` item carrying the whole `TrainState` and keep writing that layout from their immutable workspaces; every restore path dual-reads both layouts (`checkpoint._stored_as_split`). **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 (`/` + a combined `` = Σmse/Σn over all sites). `CIHiddenActsReconLoss` is the deterministic `lower_leaky` CI mask, no delta, one forward (tight torch parity); `StochasticHiddenActsReconLoss` draws `n_mask_samples` stochastic CI masks (`mask = ci + (1−ci)·s`) WITH weight deltas (the delta component is always built in JAX runs) — its draws are NOT seed-aligned to torch, so exact bitwise parity is impossible there (expected). Masked + clean run in COMPUTE_DT (bf16, matching the trained model, mirroring `load_run.py`); the MSE reduction is fp32. **AMENDED 2026-06-16** (Oli-approved): superseded the prior "keep-on-bridge / seam refused" decision — these now have a native JAX eval path; the `pd-offline-eval` torch bridge remains available for cross-framework parity checks. | -| S33 | Fine-tune init (`ExperimentConfig.resume_provenance`, LM-only): a fresh run whose own `ckpts/` is EMPTY and whose `resume_provenance is not None` initializes from a PARENT checkpoint — it loads the parent's `ckpts/` onto the fresh reference `TrainState` and keeps ONLY the trained `components` (V/U) + `ci_fn`; the optimizer states, persistent sources, and `step` are the FRESH reference's (`step = 0`). Rationale: a fine-tune runs a NEW LR/p-anneal schedule computed over the new `cfg.steps` from 0, so carrying stale Adam momentum / a stale adversary would mis-scale the restart; faith warmup is also skipped (the parent's V/U is already faithful). The run records lineage via `resume_provenance` in `config.yaml` + `wandb.config`. The parent's decomposition STRUCTURE (sites names + C, ci-fn arch) must match the new config's — asserted from the parent's pinned `config.yaml` (`run.py::assert_finetune_structural_compat`) before the orbax restore; only LR / coeffs / eps / seq / batch / steps may change. This is distinct from same-config requeue-resume (S22): on a subsequent requeue the run's own `ckpts/` is non-empty, so `restore_latest` from its own dir wins and provenance is ignored. | +| 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/` `decomposition` item (the trained `components` (V/U) + `ci_fn`; dual-read for pre-split `default`-item runs, 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..3f3d65624 100644 --- a/param_decomp/checkpoint.py +++ b/param_decomp/checkpoint.py @@ -1,30 +1,45 @@ """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` — `TrainingItem` (both optimizer states, the persistent adversaries, the + step counter), the trainer-only trajectory tail. Only trainer resume touches it. + +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. + +Runs saved before the split store one item named `default` carrying the whole +`TrainState`; those runs keep training from immutable workspaces and keep writing that +layout, so every restore path here dual-reads both layouts (`_stored_as_split`). 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). """ import dataclasses +from collections.abc import Callable +from dataclasses import dataclass from pathlib import Path -from typing import cast +from typing import Any, cast import jax +import numpy as np +import optax import orbax.checkpoint as ocp +from jaxtyping import Array from orbax.checkpoint.type_handlers import ArrayHandler, register_type_handler -from param_decomp.train import TrainState +from param_decomp.adversary import PersistentAdversary +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 @@ -33,6 +48,41 @@ register_type_handler(jax.Array, ArrayHandler(use_replica_parallel=False), override=True) +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class TrainingItem: + """The `training` checkpoint item: the trajectory state that is training's business + only — no consumer restores it.""" + + components_opt_state: optax.OptState + ci_fn_opt_state: optax.OptState + adversaries: dict[str, PersistentAdversary] + step: Array + + +def _split(state: TrainState) -> tuple[Decomposition, TrainingItem]: + return ( + Decomposition(components=state.components, ci_fn=state.ci_fn), + TrainingItem( + components_opt_state=state.components_opt_state, + ci_fn_opt_state=state.ci_fn_opt_state, + adversaries=state.adversaries, + step=state.step, + ), + ) + + +def _join(decomposition: Decomposition, training: TrainingItem) -> TrainState: + return TrainState( + components=decomposition.components, + ci_fn=decomposition.ci_fn, + components_opt_state=training.components_opt_state, + ci_fn_opt_state=training.ci_fn_opt_state, + adversaries=training.adversaries, + step=training.step, + ) + + def make_checkpoint_manager(ckpt_dir: Path, keep_last: int) -> ocp.CheckpointManager: return ocp.CheckpointManager( ckpt_dir.resolve(), @@ -43,16 +93,48 @@ def make_checkpoint_manager(ckpt_dir: Path, keep_last: int) -> ocp.CheckpointMan ) +def _stored_as_split(mgr: ocp.CheckpointManager, step: int) -> bool: + """Whether checkpoint `step` is stored as the two-item layout (`decomposition/` + + `training/`) or as the single `default/` item of runs launched before the split + (those runs keep writing it from their immutable workspaces).""" + step_dir = Path(mgr.directory) / str(step) + assert step_dir.is_dir(), f"no checkpoint dir {step_dir}" + split = (step_dir / "decomposition").is_dir() + legacy = (step_dir / "default").is_dir() + assert split != legacy, f"unrecognized checkpoint layout under {step_dir}" + return split + + def save_state(mgr: ocp.CheckpointManager, step: int, state: TrainState) -> None: - mgr.save(step, args=ocp.args.StandardSave(state)) + decomposition, training = _split(state) + mgr.save( + step, + args=ocp.args.Composite( + decomposition=ocp.args.StandardSave(decomposition), + training=ocp.args.StandardSave(training), + ), + ) mgr.wait_until_finished() def restore_step(mgr: ocp.CheckpointManager, reference: TrainState, step: int) -> TrainState: """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)) + if _stored_as_split(mgr, step): + abstract_decomp, abstract_training = jax.tree.map( + ocp.utils.to_shape_dtype_struct, _split(reference) + ) + composite = mgr.restore( + step, + args=ocp.args.Composite( + decomposition=ocp.args.StandardRestore(abstract_decomp), + training=ocp.args.StandardRestore(abstract_training), + ), + ) + restored = _join(composite["decomposition"], composite["training"]) + else: + abstract = jax.tree.map(ocp.utils.to_shape_dtype_struct, reference) + restored = mgr.restore(step, args=ocp.args.StandardRestore(abstract)) # 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,26 +155,82 @@ def restore_latest( return restore_step(mgr, reference, step), step +def _legacy_partial_restore( + mgr: ocp.CheckpointManager, + step: int, + abstract: Decomposition, + restore_args_for_leaf: Callable[[Any], ocp.RestoreArgs], +) -> Decomposition: + """Pull ONLY `components` + `ci_fn` out of a pre-split `default` item: + `partial_restore` skips every saved key absent from the item, so the + optimizer/adversary subtrees are never read off disk or materialized.""" + item = {"components": abstract.components, "ci_fn": abstract.ci_fn} + restore_args = jax.tree.map(restore_args_for_leaf, item) + restored = mgr.restore( + step, + args=ocp.args.PyTreeRestore(item=item, restore_args=restore_args, partial_restore=True), + ) + return Decomposition(components=restored["components"], ci_fn=restored["ci_fn"]) + + +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).""" + if _stored_as_split(mgr, step): + composite = mgr.restore( + step, args=ocp.args.Composite(decomposition=ocp.args.StandardRestore(abstract)) + ) + return cast(Decomposition, composite["decomposition"]) + return _legacy_partial_restore( + mgr, + step, + abstract, + lambda leaf: ocp.ArrayRestoreArgs(sharding=leaf.sharding, dtype=leaf.dtype), + ) + + +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 + if _stored_as_split(mgr, step): + 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"]) + return _legacy_partial_restore(mgr, step, abstract, host_args_for_leaf) + + 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. + abstract = jax.tree.map(ocp.utils.to_shape_dtype_struct, _split(reference)[0]) + parent = restore_decomposition(parent_mgr, parent_step, abstract) return dataclasses.replace( reference, components=parent.components, diff --git a/param_decomp/configs/pile_ppgd_bsc_RESUMESMOKE_dp8.yaml b/param_decomp/configs/pile_ppgd_bsc_RESUMESMOKE_dp8.yaml new file mode 100644 index 000000000..8cdd03cc7 --- /dev/null +++ b/param_decomp/configs/pile_ppgd_bsc_RESUMESMOKE_dp8.yaml @@ -0,0 +1,152 @@ +# Save->resume smoke for the checkpoint decomposition|training split (bridge task +# checkpoint-split-decomposition-item). pile_ppgd_bsc.yaml (persistent bsc adversary -- +# the training-item subtree that motivated the split) shrunk to dp=8 / 40 steps / +# save_every=10 / warmup=5 / evals+wandb off. Procedure: pd-lm launch -> scancel after +# the step-20 save -> pd-lm --run_id resubmit -> expect 'resumed from checkpoint step +# 20' and two-item step dirs (decomposition/ + training/); rm the run dir after. +cadence: + keep_last_n_checkpoints: 2 + save_every: 10 + train_log_every: 5 +data: + buffer_size: 1000 + column_name: input_ids + data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet + dataset_name: parquet + eval_split: train + is_tokenized: true + max_seq_len: 512 + revision: null + shuffle_each_epoch: true + streaming: false + tokenizer_name: EleutherAI/gpt-neox-20b + train_split: train +eval: + batch_size: 128 + every: 100000 + metrics: + - n_batches_accum: 1 + type: CIHistograms + - ci_alive_threshold: 0.0 + type: ComponentActivationDensity + - ci_alive_threshold: 0.0 + groups: + layer_0: + - h.0.* + layer_1: + - h.1.* + layer_2: + - h.2.* + layer_3: + - h.3.* + total: + - '*' + type: CI_L0 + - rounding_threshold: 0.0 + type: CEandKLLosses + - type: CIMeanPerComponent + - coeff: null + type: StochasticHiddenActsReconLoss + - type: CIHiddenActsReconLoss + - coeff: null + init: random + mask_scope: c + n_steps: 20 + name: PGDReconLoss_20step + step_size: 0.1 + type: PGDReconLoss + n_steps: 1 + slow_every: 100000 + slow_on_first_step: true +pd: + batch_size: 64 + ci_config: + blocks_per_chunk: 1 + d_model: 2048 + mlp_hidden: 8192 + n_blocks: 8 + n_heads: 16 + type: chunkwise_transformer + ci_fn_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: null + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + components_optimizer: + betas: + - 0.9 + - 0.999 + grad_clip_norm: 0.01 + lr_schedule: + final_val_frac: 0.1 + fn_type: cosine + start_val: 5.0e-05 + warmup_pct: 0.0 + weight_decay: 0.0 + decomposition_targets: + - C: 3072 + module_pattern: h.*.mlp.c_fc + - C: 3584 + module_pattern: h.*.mlp.down_proj + - C: 512 + module_pattern: h.*.attn.q_proj + - C: 512 + module_pattern: h.*.attn.k_proj + - C: 1024 + module_pattern: h.*.attn.v_proj + - C: 1024 + module_pattern: h.*.attn.o_proj + faithfulness_warmup_lr: 0.001 + faithfulness_warmup_steps: 5 + faithfulness_warmup_weight_decay: 0.0 + loss_metrics: + - coeff: 0.0002 + eps: 1.0e-12 + frequency: + coeff: 0.0001 + reference_token_count: 65536 + pnorm: + final_val_frac: 0.2 + fn_type: linear + start_val: 2.0 + type: ImportanceMinimalityLoss + - coeff: 0.5 + routing: + type: uniform_k_subset + type: StochasticReconSubsetLoss + - coeff: 0.5 + n_warmup_steps: 2 + optimizer: + beta1: 0.5 + beta2: 0.99 + eps: 1.0e-08 + lr_schedule: + final_val_frac: 1.0 + fn_type: constant + start_val: 0.01 + warmup_pct: 0.025 + type: adam + scope: + type: bsc + type: PersistentPGDReconLoss + - coeff: 10000000.0 + type: FaithfulnessLoss + seed: 0 + steps: 40 +run_name: jax-pile4l-ckpt-split-resume-smoke +runtime: + dp: 8 + remat_recon_forwards: false +target: + output_extract: 0 + spec: + kind: pretrained + model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP + run_path: goodfire/spd/runs/t-9d2b8f02 + weights_dtype: bfloat16 diff --git a/param_decomp/run_state.py b/param_decomp/run_state.py index a20bfc436..f3ace34a7 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, 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 = { diff --git a/param_decomp/tests/test_checkpoint.py b/param_decomp/tests/test_checkpoint.py index ea4cddbf2..a57a5ef01 100644 --- a/param_decomp/tests/test_checkpoint.py +++ b/param_decomp/tests/test_checkpoint.py @@ -8,6 +8,7 @@ import jax import jax.numpy as jnp import optax +import orbax.checkpoint as ocp from jax.sharding import Mesh from param_decomp.adversary import ( @@ -17,7 +18,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 +54,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, TrainState, make_train_step from vendored_jax.llama import LlamaConfig @@ -230,6 +233,85 @@ 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 and the + dual-read dispatch key 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: Decomposition(components=state.components, ci_fn=state.ci_fn)) + restored = restore_decomposition_to_host(consumer_mgr, 2, abstract) + for a, b in zip( + jax.tree.leaves((state.components, state.ci_fn)), + jax.tree.leaves((restored.components, restored.ci_fn)), + strict=True, + ): + assert jnp.array_equal(jnp.asarray(a), jnp.asarray(b)) + + +def _save_legacy_single_item(mgr: ocp.CheckpointManager, step: int, state: TrainState) -> None: + """Write the pre-split layout (one `default` item carrying the whole `TrainState`) — + what every run launched before the split keeps writing from its immutable workspace.""" + mgr.save(step, args=ocp.args.StandardSave(state)) + mgr.wait_until_finished() + + +def test_legacy_single_item_dual_read(tmp_path: Path): + """Pre-split checkpoints must stay readable by every restore path: full trainer + restore (`restore_latest`), consumer partial restore + (`restore_decomposition_to_host`), and fine-tune init (`init_from_parent`).""" + 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_legacy_single_item(mgr, 2, state) + assert (tmp_path / "ckpts" / "2" / "default").is_dir() + + _, fresh, _, _ = _build(seed=7) + restored = restore_latest(mgr, fresh) + assert restored is not None + loaded, ckpt_step = restored + assert ckpt_step == 2 + for a, b in zip(jax.tree.leaves(state), jax.tree.leaves(loaded), strict=True): + assert jnp.array_equal(jnp.asarray(a), jnp.asarray(b)) + + consumer_mgr = make_checkpoint_manager(tmp_path / "ckpts", keep_last=2) + abstract = jax.eval_shape(lambda: Decomposition(components=state.components, ci_fn=state.ci_fn)) + decomp = restore_decomposition_to_host(consumer_mgr, 2, abstract) + for a, b in zip( + jax.tree.leaves((state.components, state.ci_fn)), + jax.tree.leaves((decomp.components, decomp.ci_fn)), + strict=True, + ): + assert jnp.array_equal(jnp.asarray(a), jnp.asarray(b)) + + finetuned = init_from_parent(tmp_path / "ckpts", parent_step=2, reference=fresh) + for a, b in zip( + jax.tree.leaves(finetuned.components), jax.tree.leaves(state.components), strict=True + ): + assert jnp.array_equal(jnp.asarray(a), jnp.asarray(b)) + assert int(finetuned.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 diff --git a/param_decomp/train.py b/param_decomp/train.py index 6694d2157..601713b5a 100644 --- a/param_decomp/train.py +++ b/param_decomp/train.py @@ -61,6 +61,17 @@ def cast_floating(tree: Any, dtype: Any) -> Any: return jax.tree.map(lambda a: a.astype(dtype) if eqx.is_inexact_array(a) else a, tree) +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +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 + ci_fn: CIFn + + @jax.tree_util.register_dataclass @dataclass(frozen=True) class TrainState: 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//{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, ) From d73096e4a6952a122778231bab752bfb5a41707e Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Wed, 8 Jul 2026 03:02:22 +0000 Subject: [PATCH 2/5] chore(checkpoint): drop manual RESUMESMOKE test config from PR The pile_ppgd_bsc_RESUMESMOKE_dp8.yaml was a manual e2e testing artefact used to validate the two-item checkpoint format on 8 GPUs, not code intended to be committed. Removing per review. Crew-Address: task/checkpoint-split-decomposition-item --- .../pile_ppgd_bsc_RESUMESMOKE_dp8.yaml | 152 ------------------ 1 file changed, 152 deletions(-) delete mode 100644 param_decomp/configs/pile_ppgd_bsc_RESUMESMOKE_dp8.yaml diff --git a/param_decomp/configs/pile_ppgd_bsc_RESUMESMOKE_dp8.yaml b/param_decomp/configs/pile_ppgd_bsc_RESUMESMOKE_dp8.yaml deleted file mode 100644 index 8cdd03cc7..000000000 --- a/param_decomp/configs/pile_ppgd_bsc_RESUMESMOKE_dp8.yaml +++ /dev/null @@ -1,152 +0,0 @@ -# Save->resume smoke for the checkpoint decomposition|training split (bridge task -# checkpoint-split-decomposition-item). pile_ppgd_bsc.yaml (persistent bsc adversary -- -# the training-item subtree that motivated the split) shrunk to dp=8 / 40 steps / -# save_every=10 / warmup=5 / evals+wandb off. Procedure: pd-lm launch -> scancel after -# the step-20 save -> pd-lm --run_id resubmit -> expect 'resumed from checkpoint step -# 20' and two-item step dirs (decomposition/ + training/); rm the run dir after. -cadence: - keep_last_n_checkpoints: 2 - save_every: 10 - train_log_every: 5 -data: - buffer_size: 1000 - column_name: input_ids - data_files: /mnt/data/artifacts/mechanisms/param-decomp/datasets/pile_neox_tok_512/*.parquet - dataset_name: parquet - eval_split: train - is_tokenized: true - max_seq_len: 512 - revision: null - shuffle_each_epoch: true - streaming: false - tokenizer_name: EleutherAI/gpt-neox-20b - train_split: train -eval: - batch_size: 128 - every: 100000 - metrics: - - n_batches_accum: 1 - type: CIHistograms - - ci_alive_threshold: 0.0 - type: ComponentActivationDensity - - ci_alive_threshold: 0.0 - groups: - layer_0: - - h.0.* - layer_1: - - h.1.* - layer_2: - - h.2.* - layer_3: - - h.3.* - total: - - '*' - type: CI_L0 - - rounding_threshold: 0.0 - type: CEandKLLosses - - type: CIMeanPerComponent - - coeff: null - type: StochasticHiddenActsReconLoss - - type: CIHiddenActsReconLoss - - coeff: null - init: random - mask_scope: c - n_steps: 20 - name: PGDReconLoss_20step - step_size: 0.1 - type: PGDReconLoss - n_steps: 1 - slow_every: 100000 - slow_on_first_step: true -pd: - batch_size: 64 - ci_config: - blocks_per_chunk: 1 - d_model: 2048 - mlp_hidden: 8192 - n_blocks: 8 - n_heads: 16 - type: chunkwise_transformer - ci_fn_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: null - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - components_optimizer: - betas: - - 0.9 - - 0.999 - grad_clip_norm: 0.01 - lr_schedule: - final_val_frac: 0.1 - fn_type: cosine - start_val: 5.0e-05 - warmup_pct: 0.0 - weight_decay: 0.0 - decomposition_targets: - - C: 3072 - module_pattern: h.*.mlp.c_fc - - C: 3584 - module_pattern: h.*.mlp.down_proj - - C: 512 - module_pattern: h.*.attn.q_proj - - C: 512 - module_pattern: h.*.attn.k_proj - - C: 1024 - module_pattern: h.*.attn.v_proj - - C: 1024 - module_pattern: h.*.attn.o_proj - faithfulness_warmup_lr: 0.001 - faithfulness_warmup_steps: 5 - faithfulness_warmup_weight_decay: 0.0 - loss_metrics: - - coeff: 0.0002 - eps: 1.0e-12 - frequency: - coeff: 0.0001 - reference_token_count: 65536 - pnorm: - final_val_frac: 0.2 - fn_type: linear - start_val: 2.0 - type: ImportanceMinimalityLoss - - coeff: 0.5 - routing: - type: uniform_k_subset - type: StochasticReconSubsetLoss - - coeff: 0.5 - n_warmup_steps: 2 - optimizer: - beta1: 0.5 - beta2: 0.99 - eps: 1.0e-08 - lr_schedule: - final_val_frac: 1.0 - fn_type: constant - start_val: 0.01 - warmup_pct: 0.025 - type: adam - scope: - type: bsc - type: PersistentPGDReconLoss - - coeff: 10000000.0 - type: FaithfulnessLoss - seed: 0 - steps: 40 -run_name: jax-pile4l-ckpt-split-resume-smoke -runtime: - dp: 8 - remat_recon_forwards: false -target: - output_extract: 0 - spec: - kind: pretrained - model_class: param_decomp_lab.experiments.lm.pretrain.models.llama_simple_mlp.LlamaSimpleMLP - run_path: goodfire/spd/runs/t-9d2b8f02 - weights_dtype: bfloat16 From 62b100deb507e104fb22fae15ac03e50377cb4af Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Wed, 8 Jul 2026 03:16:26 +0000 Subject: [PATCH 3/5] refactor(checkpoint): drop legacy single-item dual-read (clean format break) Per review: no in-code compat for pre-split `default`-item checkpoints. Pre-split runs are migrated by an ad-hoc script when needed, rather than bloating every restore path with a fallback. - Remove `_stored_as_split` layout dispatch and `_legacy_partial_restore`. - Collapse restore_step / restore_decomposition / restore_decomposition_to_host to the two-item (`decomposition`+`training`) path only. - SPEC S22/S33, CLAUDE.md: state the clean break; drop dual-read wording. - Replace the legacy dual-read test with a two-item init_from_parent test (keeps S33 fine-tune-init coverage: decomposition carries over, opt state/adversaries/step stay the fresh reference's). Suite green at 1 and 4 simulated devices (param_decomp/tests/test_checkpoint.py, 7 passed). Crew-Address: task/checkpoint-split-decomposition-item --- CLAUDE.md | 2 +- param_decomp/CLAUDE.md | 4 +- param_decomp/SPEC.md | 4 +- param_decomp/checkpoint.py | 94 +++++++-------------------- param_decomp/tests/test_checkpoint.py | 53 ++++++--------- 5 files changed, 49 insertions(+), 108 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 27ad9f253..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// 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//{decomposition,training}/ # orbax sharded checkpoints (JAX trainer): the trained product vs the trainer-only tail; pre-split runs have a single default/ item instead + ckpts//{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 50bd97937..22385d343 100644 --- a/param_decomp/CLAUDE.md +++ b/param_decomp/CLAUDE.md @@ -212,8 +212,8 @@ 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), TWO items per step — `decomposition` (V/U + ci_fn, the product every consumer restores alone) and `training` (opt states + adversaries + step, -trainer-only); pre-split runs' single `default` item stays dual-readable (SPEC S22); -SIGTERM → save → SLURM requeue → resume from latest. +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. diff --git a/param_decomp/SPEC.md b/param_decomp/SPEC.md index 0c7c3f6bc..33ca69a29 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). **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`). 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). Runs launched before the split store one `default` item carrying the whole `TrainState` and keep writing that layout from their immutable workspaces; every restore path dual-reads both layouts (`checkpoint._stored_as_split`). **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`). 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 (`/` + a combined `` = Σmse/Σn over all sites). `CIHiddenActsReconLoss` is the deterministic `lower_leaky` CI mask, no delta, one forward (tight torch parity); `StochasticHiddenActsReconLoss` draws `n_mask_samples` stochastic CI masks (`mask = ci + (1−ci)·s`) WITH weight deltas (the delta component is always built in JAX runs) — its draws are NOT seed-aligned to torch, so exact bitwise parity is impossible there (expected). Masked + clean run in COMPUTE_DT (bf16, matching the trained model, mirroring `load_run.py`); the MSE reduction is fp32. **AMENDED 2026-06-16** (Oli-approved): superseded the prior "keep-on-bridge / seam refused" decision — these now have a native JAX eval path; the `pd-offline-eval` torch bridge remains available for cross-framework parity checks. | -| S33 | Fine-tune init (`ExperimentConfig.resume_provenance`, LM-only): a fresh run whose own `ckpts/` is EMPTY and whose `resume_provenance is not None` initializes from a PARENT checkpoint — it restores the parent's `ckpts/` `decomposition` item (the trained `components` (V/U) + `ci_fn`; dual-read for pre-split `default`-item runs, 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. | +| 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/` `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 3f3d65624..91daafe0a 100644 --- a/param_decomp/checkpoint.py +++ b/param_decomp/checkpoint.py @@ -12,10 +12,6 @@ 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. -Runs saved before the split store one item named `default` carrying the whole -`TrainState`; those runs keep training from immutable workspaces and keep writing that -layout, so every restore path here dual-reads both layouts (`_stored_as_split`). - Synchronous saves (no async): a SIGTERM-triggered save must be on disk before the process exits for SLURM requeue-resume. @@ -26,10 +22,9 @@ """ import dataclasses -from collections.abc import Callable from dataclasses import dataclass from pathlib import Path -from typing import Any, cast +from typing import cast import jax import numpy as np @@ -93,18 +88,6 @@ def make_checkpoint_manager(ckpt_dir: Path, keep_last: int) -> ocp.CheckpointMan ) -def _stored_as_split(mgr: ocp.CheckpointManager, step: int) -> bool: - """Whether checkpoint `step` is stored as the two-item layout (`decomposition/` + - `training/`) or as the single `default/` item of runs launched before the split - (those runs keep writing it from their immutable workspaces).""" - step_dir = Path(mgr.directory) / str(step) - assert step_dir.is_dir(), f"no checkpoint dir {step_dir}" - split = (step_dir / "decomposition").is_dir() - legacy = (step_dir / "default").is_dir() - assert split != legacy, f"unrecognized checkpoint layout under {step_dir}" - return split - - def save_state(mgr: ocp.CheckpointManager, step: int, state: TrainState) -> None: decomposition, training = _split(state) mgr.save( @@ -120,21 +103,17 @@ def save_state(mgr: ocp.CheckpointManager, step: int, state: TrainState) -> None def restore_step(mgr: ocp.CheckpointManager, reference: TrainState, step: int) -> TrainState: """Restore checkpoint `step` onto `reference`'s shapes/dtypes/shardings (a freshly-initialised, correctly-placed `TrainState`).""" - if _stored_as_split(mgr, step): - abstract_decomp, abstract_training = jax.tree.map( - ocp.utils.to_shape_dtype_struct, _split(reference) - ) - composite = mgr.restore( - step, - args=ocp.args.Composite( - decomposition=ocp.args.StandardRestore(abstract_decomp), - training=ocp.args.StandardRestore(abstract_training), - ), - ) - restored = _join(composite["decomposition"], composite["training"]) - else: - abstract = jax.tree.map(ocp.utils.to_shape_dtype_struct, reference) - restored = mgr.restore(step, args=ocp.args.StandardRestore(abstract)) + abstract_decomp, abstract_training = jax.tree.map( + ocp.utils.to_shape_dtype_struct, _split(reference) + ) + composite = mgr.restore( + step, + args=ocp.args.Composite( + decomposition=ocp.args.StandardRestore(abstract_decomp), + training=ocp.args.StandardRestore(abstract_training), + ), + ) + restored = _join(composite["decomposition"], 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 @@ -155,40 +134,15 @@ def restore_latest( return restore_step(mgr, reference, step), step -def _legacy_partial_restore( - mgr: ocp.CheckpointManager, - step: int, - abstract: Decomposition, - restore_args_for_leaf: Callable[[Any], ocp.RestoreArgs], -) -> Decomposition: - """Pull ONLY `components` + `ci_fn` out of a pre-split `default` item: - `partial_restore` skips every saved key absent from the item, so the - optimizer/adversary subtrees are never read off disk or materialized.""" - item = {"components": abstract.components, "ci_fn": abstract.ci_fn} - restore_args = jax.tree.map(restore_args_for_leaf, item) - restored = mgr.restore( - step, - args=ocp.args.PyTreeRestore(item=item, restore_args=restore_args, partial_restore=True), - ) - return Decomposition(components=restored["components"], ci_fn=restored["ci_fn"]) - - 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).""" - if _stored_as_split(mgr, step): - composite = mgr.restore( - step, args=ocp.args.Composite(decomposition=ocp.args.StandardRestore(abstract)) - ) - return cast(Decomposition, composite["decomposition"]) - return _legacy_partial_restore( - mgr, - step, - abstract, - lambda leaf: ocp.ArrayRestoreArgs(sharding=leaf.sharding, dtype=leaf.dtype), + composite = mgr.restore( + step, args=ocp.args.Composite(decomposition=ocp.args.StandardRestore(abstract)) ) + return cast(Decomposition, composite["decomposition"]) def restore_decomposition_to_host( @@ -199,16 +153,14 @@ def restore_decomposition_to_host( `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 - if _stored_as_split(mgr, step): - 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"]) - return _legacy_partial_restore(mgr, step, abstract, host_args_for_leaf) + 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: diff --git a/param_decomp/tests/test_checkpoint.py b/param_decomp/tests/test_checkpoint.py index a57a5ef01..647dce939 100644 --- a/param_decomp/tests/test_checkpoint.py +++ b/param_decomp/tests/test_checkpoint.py @@ -8,7 +8,6 @@ import jax import jax.numpy as jnp import optax -import orbax.checkpoint as ocp from jax.sharding import Mesh from param_decomp.adversary import ( @@ -234,8 +233,8 @@ def test_no_checkpoint_returns_none(tmp_path: Path): def test_saved_layout_is_two_items(tmp_path: Path): - """The on-disk item names are the cross-version contract every consumer and the - dual-read dispatch key on — pin them.""" + """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) @@ -268,45 +267,35 @@ def test_consumer_restores_decomposition_to_host(tmp_path: Path): assert jnp.array_equal(jnp.asarray(a), jnp.asarray(b)) -def _save_legacy_single_item(mgr: ocp.CheckpointManager, step: int, state: TrainState) -> None: - """Write the pre-split layout (one `default` item carrying the whole `TrainState`) — - what every run launched before the split keeps writing from its immutable workspace.""" - mgr.save(step, args=ocp.args.StandardSave(state)) - mgr.wait_until_finished() - - -def test_legacy_single_item_dual_read(tmp_path: Path): - """Pre-split checkpoints must stay readable by every restore path: full trainer - restore (`restore_latest`), consumer partial restore - (`restore_decomposition_to_host`), and fine-tune init (`init_from_parent`).""" - lm, state, step, resid = _build(seed=1) +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): - state, _ = step(lm, state, resid, jax.random.PRNGKey(i)) + parent, _ = step(lm, parent, resid, jax.random.PRNGKey(i)) mgr = make_checkpoint_manager(tmp_path / "ckpts", keep_last=2) - _save_legacy_single_item(mgr, 2, state) - assert (tmp_path / "ckpts" / "2" / "default").is_dir() + 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) - restored = restore_latest(mgr, fresh) - assert restored is not None - loaded, ckpt_step = restored - assert ckpt_step == 2 - for a, b in zip(jax.tree.leaves(state), jax.tree.leaves(loaded), strict=True): - assert jnp.array_equal(jnp.asarray(a), jnp.asarray(b)) + finetuned = init_from_parent(tmp_path / "ckpts", parent_step=2, reference=fresh) - consumer_mgr = make_checkpoint_manager(tmp_path / "ckpts", keep_last=2) - abstract = jax.eval_shape(lambda: Decomposition(components=state.components, ci_fn=state.ci_fn)) - decomp = restore_decomposition_to_host(consumer_mgr, 2, abstract) + # decomposition carries over from the parent... for a, b in zip( - jax.tree.leaves((state.components, state.ci_fn)), - jax.tree.leaves((decomp.components, decomp.ci_fn)), + jax.tree.leaves((finetuned.components, finetuned.ci_fn)), + jax.tree.leaves((parent.components, parent.ci_fn)), strict=True, ): assert jnp.array_equal(jnp.asarray(a), jnp.asarray(b)) - - finetuned = init_from_parent(tmp_path / "ckpts", parent_step=2, reference=fresh) + # ...while optimizer state, adversaries and step stay the fresh reference's. for a, b in zip( - jax.tree.leaves(finetuned.components), jax.tree.leaves(state.components), strict=True + jax.tree.leaves((finetuned.components_opt_state, finetuned.ci_fn_opt_state)), + jax.tree.leaves((fresh.components_opt_state, fresh.ci_fn_opt_state)), + strict=True, ): assert jnp.array_equal(jnp.asarray(a), jnp.asarray(b)) assert int(finetuned.step) == 0 From e4d1b9f65a1af1f743dad523cd75f85c5b6caa19 Mon Sep 17 00:00:00 2001 From: "PD User (shared)" Date: Wed, 8 Jul 2026 03:59:27 +0000 Subject: [PATCH 4/5] =?UTF-8?q?refactor(train):=20TrainState=20composes=20?= =?UTF-8?q?{decomposition,=20training}=20=E2=80=94=20one=20representation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Per review: rather than a flat TrainState that the checkpoint code regroups into two items (the _split/_join helpers), TrainState now IS the composition of the two checkpoint items: TrainState{decomposition: Decomposition, training: TrainingItem} so save/restore map straight onto its own fields — _split/_join deleted, no regrouping anywhere. TrainingItem moved from checkpoint.py to train.py (it's now a TrainState field). SPEC S22 records the one-representation invariant. Churn: every state. access becomes state.decomposition. / state.training. — the hot train step, run.py resume/warmup, run_state init, slow_eval hidden-acts, the lab run/eval paths, and all constructing tests. Pure regrouping, no numeric change. Validated (1 + 4 sim devices): test_checkpoint, test_checkpoint_production_topology, test_generic_model_io, test_no_bake_invariant, test_llama8b, test_llama_simple_mlp, test_finetune_resume, stacked_parity, tms, resid_mlp — all green (the only failure, test_pretrained_target_converts_with_wildcards, is the pre-existing tip failure). basedpyright clean; ruff clean. Crew-Address: task/checkpoint-split-decomposition-item --- param_decomp/SPEC.md | 2 +- param_decomp/checkpoint.py | 69 ++++--------------- param_decomp/experiments/invariance_check.py | 32 +++++---- param_decomp/run.py | 12 ++-- param_decomp/run_state.py | 15 ++-- param_decomp/slow_eval.py | 16 ++++- .../stacked_parity/test_stacked_parity.py | 38 +++++----- param_decomp/tests/test_checkpoint.py | 54 ++++++++------- .../test_checkpoint_production_topology.py | 18 ++--- param_decomp/tests/test_finetune_resume.py | 26 ++++--- param_decomp/tests/test_generic_model_io.py | 21 +++--- param_decomp/tests/test_llama8b.py | 66 ++++++++++-------- param_decomp/tests/test_llama_simple_mlp.py | 50 ++++++++------ param_decomp/tests/test_no_bake_invariant.py | 15 ++-- param_decomp/train.py | 61 +++++++++------- param_decomp_lab/experiments/lm/run.py | 21 ++++-- param_decomp_lab/experiments/resid_mlp/run.py | 4 +- .../experiments/resid_mlp/test_resid_mlp.py | 38 ++++++---- param_decomp_lab/experiments/tms/run.py | 4 +- param_decomp_lab/experiments/tms/test_tms.py | 38 ++++++---- 20 files changed, 327 insertions(+), 273 deletions(-) diff --git a/param_decomp/SPEC.md b/param_decomp/SPEC.md index 33ca69a29..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). **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`). 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). | +| 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. | diff --git a/param_decomp/checkpoint.py b/param_decomp/checkpoint.py index 91daafe0a..d6990c2f5 100644 --- a/param_decomp/checkpoint.py +++ b/param_decomp/checkpoint.py @@ -5,8 +5,11 @@ - `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` — `TrainingItem` (both optimizer states, the persistent adversaries, the - step counter), the trainer-only trajectory tail. Only trainer resume touches it. +- `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 @@ -22,18 +25,14 @@ """ import dataclasses -from dataclasses import dataclass from pathlib import Path from typing import cast import jax import numpy as np -import optax import orbax.checkpoint as ocp -from jaxtyping import Array from orbax.checkpoint.type_handlers import ArrayHandler, register_type_handler -from param_decomp.adversary import PersistentAdversary from param_decomp.train import Decomposition, TrainState # Replica-parallel writes (multiple hosts cooperatively writing a REPLICATED array) @@ -43,41 +42,6 @@ register_type_handler(jax.Array, ArrayHandler(use_replica_parallel=False), override=True) -@jax.tree_util.register_dataclass -@dataclass(frozen=True) -class TrainingItem: - """The `training` checkpoint item: the trajectory state that is training's business - only — no consumer restores it.""" - - components_opt_state: optax.OptState - ci_fn_opt_state: optax.OptState - adversaries: dict[str, PersistentAdversary] - step: Array - - -def _split(state: TrainState) -> tuple[Decomposition, TrainingItem]: - return ( - Decomposition(components=state.components, ci_fn=state.ci_fn), - TrainingItem( - components_opt_state=state.components_opt_state, - ci_fn_opt_state=state.ci_fn_opt_state, - adversaries=state.adversaries, - step=state.step, - ), - ) - - -def _join(decomposition: Decomposition, training: TrainingItem) -> TrainState: - return TrainState( - components=decomposition.components, - ci_fn=decomposition.ci_fn, - components_opt_state=training.components_opt_state, - ci_fn_opt_state=training.ci_fn_opt_state, - adversaries=training.adversaries, - step=training.step, - ) - - def make_checkpoint_manager(ckpt_dir: Path, keep_last: int) -> ocp.CheckpointManager: return ocp.CheckpointManager( ckpt_dir.resolve(), @@ -89,12 +53,11 @@ def make_checkpoint_manager(ckpt_dir: Path, keep_last: int) -> ocp.CheckpointMan def save_state(mgr: ocp.CheckpointManager, step: int, state: TrainState) -> None: - decomposition, training = _split(state) mgr.save( step, args=ocp.args.Composite( - decomposition=ocp.args.StandardSave(decomposition), - training=ocp.args.StandardSave(training), + decomposition=ocp.args.StandardSave(state.decomposition), + training=ocp.args.StandardSave(state.training), ), ) mgr.wait_until_finished() @@ -103,17 +66,15 @@ def save_state(mgr: ocp.CheckpointManager, step: int, state: TrainState) -> None def restore_step(mgr: ocp.CheckpointManager, reference: TrainState, step: int) -> TrainState: """Restore checkpoint `step` onto `reference`'s shapes/dtypes/shardings (a freshly-initialised, correctly-placed `TrainState`).""" - abstract_decomp, abstract_training = jax.tree.map( - ocp.utils.to_shape_dtype_struct, _split(reference) - ) + abstract = jax.tree.map(ocp.utils.to_shape_dtype_struct, reference) composite = mgr.restore( step, args=ocp.args.Composite( - decomposition=ocp.args.StandardRestore(abstract_decomp), - training=ocp.args.StandardRestore(abstract_training), + decomposition=ocp.args.StandardRestore(abstract.decomposition), + training=ocp.args.StandardRestore(abstract.training), ), ) - restored = _join(composite["decomposition"], composite["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 @@ -181,10 +142,6 @@ def init_from_parent(parent_ckpt_dir: Path, parent_step: int, reference: TrainSt assert parent_step in parent_mgr.all_steps(), ( f"parent step {parent_step} not in {parent_ckpt_dir} (have {sorted(parent_mgr.all_steps())})" ) - abstract = jax.tree.map(ocp.utils.to_shape_dtype_struct, _split(reference)[0]) + abstract = jax.tree.map(ocp.utils.to_shape_dtype_struct, reference.decomposition) parent = restore_decomposition(parent_mgr, parent_step, abstract) - return dataclasses.replace( - reference, - components=parent.components, - ci_fn=parent.ci_fn, - ) + 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 f3ace34a7..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 Decomposition, TrainState +from param_decomp.train import Decomposition, TrainingItem, TrainState def optax_schedule(config: ScheduleConfig, total_steps: int) -> Callable[[ArrayLike], Array]: @@ -165,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 - `[/]` log keys. `state.components`/`state.ci_fn` are the restored + `[/]` 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 647dce939..979cdb335 100644 --- a/param_decomp/tests/test_checkpoint.py +++ b/param_decomp/tests/test_checkpoint.py @@ -53,7 +53,7 @@ init_sources_sharded, ) from param_decomp.tests.test_llama8b import _tiny_cfg, _tiny_decomposed_lm -from param_decomp.train import Decomposition, TrainState, make_train_step +from param_decomp.train import Decomposition, TrainingItem, TrainState, make_train_step from vendored_jax.llama import LlamaConfig @@ -114,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( ( @@ -180,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) @@ -192,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) @@ -209,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 @@ -257,11 +259,11 @@ def test_consumer_restores_decomposition_to_host(tmp_path: Path): # 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: Decomposition(components=state.components, ci_fn=state.ci_fn)) + 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.components, state.ci_fn)), - jax.tree.leaves((restored.components, restored.ci_fn)), + jax.tree.leaves(state.decomposition), + jax.tree.leaves(restored), strict=True, ): assert jnp.array_equal(jnp.asarray(a), jnp.asarray(b)) @@ -286,19 +288,21 @@ def test_init_from_parent_restores_decomposition_only(tmp_path: Path): # decomposition carries over from the parent... for a, b in zip( - jax.tree.leaves((finetuned.components, finetuned.ci_fn)), - jax.tree.leaves((parent.components, parent.ci_fn)), + 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.components_opt_state, finetuned.ci_fn_opt_state)), - jax.tree.leaves((fresh.components_opt_state, fresh.ci_fn_opt_state)), + 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.step) == 0 + assert int(finetuned.training.step) == 0 def _build_sharded(seed: int, mesh: Mesh): @@ -329,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 @@ -354,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/train.py b/param_decomp/train.py index 601713b5a..6408f257d 100644 --- a/param_decomp/train.py +++ b/param_decomp/train.py @@ -68,15 +68,16 @@ class Decomposition: 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 - ci_fn: CIFn + components: DecompVU # the universal trainable V/U pytree, fp32 masters + ci_fn: CIFn # fp32 masters @jax.tree_util.register_dataclass @dataclass(frozen=True) -class TrainState: - components: DecompVU # the universal trainable V/U pytree, fp32 masters - ci_fn: CIFn # fp32 masters +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] @@ -86,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` / `grad_norms/ci_fns` (paths are this @@ -269,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. @@ -285,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) @@ -297,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 @@ -315,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 @@ -474,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()}, ) @@ -502,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, @@ -529,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/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, ( From 4dc887afde592119f83854818da07e508a21264b Mon Sep 17 00:00:00 2001 From: oli Date: Wed, 8 Jul 2026 19:16:24 +0000 Subject: [PATCH 5/5] =?UTF-8?q?feat(tools):=20strip=5Fcheckpoint=20?= =?UTF-8?q?=E2=80=94=20convert=20a=20pre-split=20ckpt=20to=20its=20decompo?= =?UTF-8?q?sition=20item?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The pre-split-run migration the S22 clean format break requires, doubling as the storage lever: a pre-split checkpoint is ~90% resume-only state (persistent adversary + all optimizer moments); the tool partial-restores components+ci_fn as host numpy and rewrites ckpts/ as the two-item layout with ONLY the decomposition item, staged + tree-verified before an in-place swap. Dry-run by default, --apply to strip. The run can no longer warm-resume — for finished or abandoned runs only. Validated on a scratch copy of p-e118ece9/ckpts/355000 (19.2 -> 2.1 GiB, 7s): restore_decomposition_to_host over the stripped checkpoint is byte-equal to the original across all 163 decomposition leaves. Handles both current (adversaries) and old jax_runs-era (sources/sources_opt_state) key layouts; unknown keys fail. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01RkKMDdM1PxaL66RmEJv8QN --- param_decomp/tools/strip_checkpoint.py | 163 +++++++++++++++++++++++++ 1 file changed, 163 insertions(+) create mode 100644 param_decomp/tools/strip_checkpoint.py 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/` 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 [--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()