feat(train): donate state into the faith-warmup and pretrain steps

CLAUDE.md

@@ -158,7 +158,7 @@ and returns JAX-native as the #10 torch->jax adapter.
   `DataConfig` / `EvalConfig` / … + the `TargetSites` protocol). The engine + numerics
   (`run.py` = `run_decomposition_training`, `lm.py` / `train.py` / `ci_fn.py` /
   `targets/llama8b.py` / `targets/llama_simple_mlp.py` / `adversary.py` / `recon.py` / `losses.py` /
-  `checkpoint.py` / `sharding.py` / `eval.py` / `slow_eval.py` / `arithmetic_eval.py` +
+  `checkpoint.py` / `sharding.py` / `donation.py` / `eval.py` / `slow_eval.py` / `arithmetic_eval.py` +
   `log.py`) plus `configs/`
   (the self-contained run yamls) and `tests/` (incl. the `tests/equivalence/` frozen
   torch↔JAX goldens). The torch oracle lives at git tag `torch-oracle`.

param_decomp/donation.py

@@ -0,0 +1,50 @@
+"""Loud verification that a donating jitted step actually reused its input buffers.
+
+Dispatch-time donation failure is SILENT: the `Some donated buffers were not usable`
+warning is lowering-time only (jax `mlir.py`), equinox's `donate="all"` suppresses even
+that, and on GPU a runtime-blocked donation falls back to a fresh allocation + copy with
+no signal at all — the step then peaks ~one full copy of the donated state above steady
+state (the resume OOM; lore 2026-07-03--resume-oom-is-buffer-donation-asymmetry).
+Aliasing is observable only at the buffer-pointer level, so callers snapshot the input
+pointers before the one step they distrust and check reuse after.
+"""
+
+import jax
+
+
+def buffer_bytes_by_ptr(tree: object) -> dict[int, int]:
+    """Device-buffer pointer -> shard nbytes over every addressable shard of `tree`'s
+    array leaves. Holds no array references (the per-shard views die with the
+    comprehension), so a snapshot cannot itself block donation."""
+    return {
+        shard.data.unsafe_buffer_pointer(): shard.data.nbytes
+        for leaf in jax.tree.leaves(tree)
+        if isinstance(leaf, jax.Array)
+        for shard in leaf.addressable_shards
+    }
+
+
+def reused_fraction(in_buffers: dict[int, int], out_tree: object) -> float:
+    """Bytes-weighted fraction of `in_buffers` reappearing under `out_tree`. Pointer SETS,
+    not positions: XLA may cross-match same-aval leaves (e.g. Adam m against v)."""
+    total = sum(in_buffers.values())
+    reused = sum(
+        nbytes
+        for ptr, nbytes in buffer_bytes_by_ptr(out_tree).items()
+        if in_buffers.get(ptr) == nbytes
+    )
+    return reused / total
+
+
+def warn_if_not_donated(in_buffers: dict[int, int], out_tree: object, what: str) -> None:
+    """Print `DONATION FAILED` when <95% of the snapshotted input bytes were reused."""
+    fraction = reused_fraction(in_buffers, out_tree)
+    if fraction < 0.95:
+        total = sum(in_buffers.values())
+        print(
+            f"[rank {jax.process_index()}] DONATION FAILED on {what}: only "
+            f"{fraction * total / 2**30:.2f}/{total / 2**30:.2f} GiB of input buffers "
+            "were reused; this step peaked ~one full copy of the donated state above "
+            "steady state (lore 2026-07-03--resume-oom-is-buffer-donation-asymmetry)",
+            flush=True,
+        )

param_decomp/run.py

@@ -58,6 +58,7 @@
 )
 from param_decomp.ci_fn import CIFnArch
 from param_decomp.configs import Cadence, PDConfig, ProfileConfig, flatten_typed_lists
+from param_decomp.donation import buffer_bytes_by_ptr, warn_if_not_donated
 from param_decomp.lm import DecomposedModel
 from param_decomp.recon import build_loss_terms
 from param_decomp.run_state import build_optimizers, init_train_state
@@ -407,10 +408,18 @@ def _init_or_restore_state(
         warmed_components = state.components
         t0 = time.time()
         faith_warmup_loss = None
+        # The warmup step donates components + opt state; dispatch-time donation failure
+        # is silent (a resident extra copy right before the run's peak phase), so check
+        # buffer reuse loudly on the first iteration.
+        warmup_in_buffers = buffer_bytes_by_ptr((warmed_components, faith_warmup_opt_state))
         for _ in range(pd.faithfulness_warmup_steps):
             warmed_components, faith_warmup_opt_state, faith_warmup_loss = faith_warmup_step(
                 lm, warmed_components, faith_warmup_opt_state
             )
+            if warmup_in_buffers is not None:
+                out_trees = (warmed_components, faith_warmup_opt_state)
+                warn_if_not_donated(warmup_in_buffers, out_trees, "the first faith-warmup step")
+                warmup_in_buffers = None
             if _sigterm_consensus():
                 # No valid checkpoint exists yet (the step-0 save happens only after warmup
                 # completes, and resume skips warmup whenever a checkpoint is present — a

param_decomp/tests/test_llama_simple_mlp.py

@@ -34,6 +34,7 @@
     SCScope,
     UniformKSubsetRoutingConfig,
 )
+from param_decomp.donation import buffer_bytes_by_ptr, reused_fraction
 from param_decomp.lm import DecomposedModel
 from param_decomp.recon import build_loss_terms
 from param_decomp.schedule import ScheduleConfig
@@ -457,6 +458,25 @@ def test_faith_warmup_decreases_faith():
     assert float(loss) < first_loss * 0.9, (first_loss, float(loss))
 
 
+def test_faith_warmup_step_donates():
+    """The warmup step reuses the components + opt-state buffers (donation, checked at
+    the pointer level) and leaves the model — the non-donated first arg — readable."""
+    cfg = _tiny_cfg()
+    sites = site_specs(cfg, canonical_site_cs(_MIXED_SITE_CS))
+    lm = _tiny_decomposed_model(cfg, sites, jax.random.PRNGKey(0))
+    vu = init_decomp_vu(sites, jax.random.PRNGKey(1))
+    opt = optax.adamw(1e-2, weight_decay=0.0)
+    wstep = make_faith_warmup_step(opt)
+    ostate = opt.init(eqx.filter(vu, eqx.is_array))
+    vu, ostate, _ = wstep(lm, vu, ostate)  # settle: inputs below are jit outputs
+    in_buffers = buffer_bytes_by_ptr((vu, ostate))
+    new_vu, new_ostate, _ = wstep(lm, vu, ostate)
+    jax.block_until_ready((new_vu, new_ostate))
+    fraction = reused_fraction(in_buffers, (new_vu, new_ostate))
+    assert fraction >= 0.95, f"only {fraction:.2%} of warmup state bytes reused"
+    jax.block_until_ready(eqx.filter(lm, eqx.is_array))  # model survives donation
+
+
 def test_decomp_vu_shapes_fp32():
     cfg = _tiny_cfg()
     sites = site_specs(cfg, _MIXED_SITE_CS)

param_decomp/tests/test_pretrain.py

@@ -4,6 +4,7 @@
 import tempfile
 from pathlib import Path
 
+import equinox as eqx
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -12,6 +13,7 @@
 import pytest
 
 import param_decomp.targets.llama_simple_mlp as lsm
+from param_decomp.donation import buffer_bytes_by_ptr, reused_fraction
 from pretrain.cache import torch_model_config_dict, write_pretrain_cache
 from pretrain.config import PretrainConfig, PretrainDataConfig
 from pretrain.models import (
@@ -21,7 +23,15 @@
     init_model,
     model_logits,
 )
-from pretrain.train import train
+from pretrain.train import (
+    TrainState,
+    _make_checkpoint_manager,
+    _restore_latest,
+    _save,
+    make_optimizer,
+    make_train_step,
+    train,
+)
 
 
 def _tiny_mlp_cfg() -> LlamaSimpleMLPConfig:
@@ -149,5 +159,64 @@ def test_training_smoke_loss_decreases():
         assert target.lm_head.shape == (mc.vocab_size, mc.n_embd)
 
 
+def _tiny_train_setup():
+    mc = _tiny_mlp_cfg()
+    cfg = PretrainConfig(
+        model=mc,
+        data=PretrainDataConfig(dir=Path("/tmp"), tokenizer_name="x"),
+        global_batch=4,
+        num_iterations=4,
+        learning_rate=1e-3,
+        warmup_iters=0,
+        learning_rate_decay_frac=0.1,
+        weight_decay=0.0,
+        grad_clip=1.0,
+        dtype="float32",
+        run_name="t",
+    )
+    model = init_model(mc, jax.random.PRNGKey(0))
+    optimizer = make_optimizer(cfg, model)
+    state = TrainState(
+        model=model,
+        opt_state=optimizer.init(eqx.filter(model, eqx.is_array)),
+        step=jnp.zeros((), jnp.int32),
+    )
+    return mc, make_train_step(cfg, optimizer), state
+
+
+def _tokens(mc: LlamaSimpleMLPConfig, seed: int) -> jax.Array:
+    return jax.random.randint(jax.random.PRNGKey(seed), (4, mc.block_size + 1), 0, mc.vocab_size)
+
+
+def test_pretrain_step_donates_state():
+    """The train step reuses the model + opt-state buffers (donation, pointer-checked)."""
+    mc, step_fn, state = _tiny_train_setup()
+    state, _ = step_fn(state, _tokens(mc, 1))  # settle: state is now jit outputs
+    in_buffers = buffer_bytes_by_ptr(state)
+    new_state, _ = step_fn(state, _tokens(mc, 2))
+    jax.block_until_ready(new_state)
+    fraction = reused_fraction(in_buffers, new_state)
+    assert fraction >= 0.95, f"only {fraction:.2%} of state bytes reused"
+
+
+def test_restored_pretrain_state_is_donatable(tmp_path: Path):
+    """Orbax round-trip preserves donatability: `_restore_latest` re-materialises the
+    restored tree as jit outputs, so the first resumed step's donation aliases instead of
+    silently copying (jax#18617)."""
+    mc, step_fn, state = _tiny_train_setup()
+    state, _ = step_fn(state, _tokens(mc, 1))
+    mgr = _make_checkpoint_manager(tmp_path / "ckpts", keep_last=1)
+    _save(mgr, 1, state)
+    restored = _restore_latest(mgr, state)
+    assert restored is not None
+    restored_state, restored_step = restored
+    assert restored_step == 1
+    in_buffers = buffer_bytes_by_ptr(restored_state)
+    new_state, _ = step_fn(restored_state, _tokens(mc, 2))
+    jax.block_until_ready(new_state)
+    fraction = reused_fraction(in_buffers, new_state)
+    assert fraction >= 0.95, f"only {fraction:.2%} of restored state bytes reused"
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v"])

param_decomp/train.py

@@ -537,7 +537,12 @@ def make_faith_warmup_step(
     compiler_options: dict[str, bool | int | str] | None = None,
 ) -> Callable[[DecomposedModel, DecompVU, optax.OptState], tuple[DecompVU, optax.OptState, Array]]:
     """`model` is the jit ARG (frozen weights traced, not baked) — `weight_deltas` reads its
-    per-site W slices, so closing over the model would bake them into the HLO."""
+    per-site W slices, so closing over the model would bake them into the HLO.
+
+    Donating (all but the model, matching the train step): the warmup loop rebinds
+    components + opt state every iteration right before the run's peak-memory phase, so
+    the step must not hold both copies live. Callers may not reuse the passed-in
+    components/opt state after a call."""
 
     def warmup_step(
         model: DecomposedModel, components: DecompVU, opt_state: optax.OptState
@@ -549,4 +554,4 @@ def loss_fn(components_: DecompVU) -> Array:
         updates, opt_state = opt.update(grad, opt_state, eqx.filter(components, eqx.is_array))
         return eqx.apply_updates(components, updates), opt_state, loss
 
-    return filter_jit(warmup_step, compiler_options=compiler_options)
+    return filter_jit(warmup_step, donate="all-except-first", compiler_options=compiler_options)

pretrain/train.py

@@ -38,6 +38,7 @@
 from orbax.checkpoint.type_handlers import ArrayHandler, register_type_handler
 
 from param_decomp.data import BatchSchedule, ShardServer, scan_shards
+from param_decomp.donation import buffer_bytes_by_ptr, warn_if_not_donated
 from param_decomp.sharding import hsdp_mesh, init_distributed
 from pretrain.cache import (
     cache_dir_for,
@@ -119,7 +120,9 @@ def make_train_step(cfg: PretrainConfig, optimizer: optax.GradientTransformation
     compute_dtype = jnp.bfloat16 if cfg.dtype == "bfloat16" else jnp.float32
     block = cfg.block_size
 
-    @eqx.filter_jit
+    # Donate everything: the loop rebinds `state` every step (holding old+new state live
+    # would double the resident model+Adam memory) and `tokens` is fresh per step.
+    @eqx.filter_jit(donate="all")
     def step_fn(state: TrainState, tokens: Int[Array, "b tplus1"]) -> tuple[TrainState, Array]:
         def loss_fn(model: PretrainModel) -> Array:
             cast_model = _cast_arrays(model, compute_dtype)
@@ -188,12 +191,18 @@ def _save(mgr: ocp.CheckpointManager, step: int, state: TrainState) -> None:
 def _restore_latest(
     mgr: ocp.CheckpointManager, reference: TrainState
 ) -> tuple[TrainState, int] | None:
+    """Restore the newest checkpoint onto `reference`'s shapes/shardings, re-materialised
+    as jit outputs: orbax-restored arrays are not reliably donatable to the jitted train
+    step (jax#18617), and a dispatch-time donation failure silently copies — the first
+    resumed step then peaks one full `TrainState` above steady state (the resume OOM
+    mechanism, `param_decomp.checkpoint.restore_step`'s precedent)."""
     step = mgr.latest_step()
     if step is None:
         return None
     abstract = jax.tree.map(ocp.utils.to_shape_dtype_struct, reference)
     restored = mgr.restore(step, args=ocp.args.StandardRestore(abstract))
-    return cast(TrainState, restored), step
+    rematerialize = jax.jit(lambda t: t, out_shardings=jax.tree.map(lambda r: r.format, reference))
+    return cast(TrainState, rematerialize(restored)), step
 
 
 class MetricsSink:
@@ -287,11 +296,19 @@ def train(cfg: PretrainConfig) -> None:
     tokens_per_step = cfg.global_batch * cfg.block_size
     window_t0 = time.time()
 
+    # The restored state should donate like fresh state (rematerialised in
+    # `_restore_latest`); dispatch-time failure is silent, so check the first resumed step.
+    resumed_state_buffers = buffer_bytes_by_ptr(state) if start_step > 0 else None
+
     for step in range(start_step, cfg.num_iterations):
         tokens = _global_token_batch(server.local_batch(step), mesh, cfg.global_batch)
         state, loss = step_fn(state, tokens)
         now = step + 1
 
+        if resumed_state_buffers is not None:
+            warn_if_not_donated(resumed_state_buffers, state, "the first resumed step")
+            resumed_state_buffers = None
+
         if now % cfg.log_every == 0 or now == cfg.num_iterations:
             jax.block_until_ready(loss)
             dt = time.time() - window_t0