diff --git a/CLAUDE.md b/CLAUDE.md index fae0420fb..f75d65b5d 100644 --- a/CLAUDE.md +++ b/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`. diff --git a/param_decomp/donation.py b/param_decomp/donation.py new file mode 100644 index 000000000..a6001e6f1 --- /dev/null +++ b/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, + ) diff --git a/param_decomp/run.py b/param_decomp/run.py index 17ebca0eb..e2274e39a 100644 --- a/param_decomp/run.py +++ b/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 diff --git a/param_decomp/tests/test_llama_simple_mlp.py b/param_decomp/tests/test_llama_simple_mlp.py index 567e524ea..db9c8cb6a 100644 --- a/param_decomp/tests/test_llama_simple_mlp.py +++ b/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) diff --git a/param_decomp/tests/test_pretrain.py b/param_decomp/tests/test_pretrain.py index 623d20564..44bef4d3d 100644 --- a/param_decomp/tests/test_pretrain.py +++ b/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"]) diff --git a/param_decomp/train.py b/param_decomp/train.py index 6694d2157..e15362f97 100644 --- a/param_decomp/train.py +++ b/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) diff --git a/pretrain/train.py b/pretrain/train.py index e53d5cbab..db208a8a8 100644 --- a/pretrain/train.py +++ b/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