This repository was archived by the owner on Jul 17, 2026. It is now read-only.
feat(train): donate state into the faith-warmup and pretrain steps#971
Open
claude-spd1 wants to merge 1 commit into
Open
feat(train): donate state into the faith-warmup and pretrain steps#971claude-spd1 wants to merge 1 commit into
claude-spd1 wants to merge 1 commit into
Conversation
The faith-warmup step and the pretrain train step both rebind multi-GiB state (components/model + Adam moments) every iteration with no buffer donation, so old+new copies stayed co-resident — the warmup right before the run's peak-memory phase. Donate matching the train-step convention (S22-adjacent; no numerics change): warmup keeps the model arg (all-except-first), pretrain donates all. Donation failure at dispatch is silent (equinox suppresses even the lowering-time warning), so: pretrain's restore re-materialises the orbax tree as jit outputs (jax#18617, `checkpoint.restore_step`'s mechanism), and a new shared `param_decomp/donation.py` provides bytes-weighted buffer-pointer-reuse checks, wired one-shot into the first warmup iteration and the first resumed pretrain step (prints DONATION FAILED). Pointer-aliasing tests pin donation for both steps and for the restored pretrain state. Bridge task donate-warmup-pretrain-steps (from the 2026-07-10 compile audit); guard precedent PR #932. Crew-Address: task/donate-warmup-pretrain-steps
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Two jitted steps rebound multi-GiB state every iteration with no buffer donation, keeping old+new copies co-resident:
param_decomp/train.py::make_faith_warmup_step) — components + warmup Adam state, right before the run's peak-memory phase. Nowdonate="all-except-first", exactly the train step's convention (the model stays a non-donated traced arg).pretrain/train.py::make_train_step) — full model + Adam state. Nowdonate="all"(the loop rebindsstate;tokensis fresh per step).Because dispatch-time donation failure is silent (the lowering-time warning is the only one that exists, and equinox suppresses even that under
donate="all"), the change ships with loud verification:param_decomp/donation.py: bytes-weighted buffer-pointer-reuse helpers (buffer_bytes_by_ptr/reused_fraction/warn_if_not_donated, the PR fix(checkpoint): make orbax-restored TrainState donatable — solve the resume OOM (S22) #932 guard shape). One-shotDONATION FAILEDchecks wired into the first faith-warmup iteration (run.py) and the first resumed pretrain step._restore_latestre-materialises the orbax-restored tree as jit outputs (identity jit pinned to the reference's formats) — orbax-restored arrays are not reliably donatable (jax#18617), so without this, adding donation would reintroduce the resume-OOM mechanism on the pretrain side. Same mechanism as fix(checkpoint): make orbax-restored TrainState donatable — solve the resume OOM (S22) #932'srestore_step; deliberately non-overlapping with that open PR's hunks (itsrun.pycopy of the guard can dedupe ontodonation.pywhen it merges).Stale-reference audit (the M2-class hazard): every consumer of the pre-step bindings at both call sites rebinds before the next donating call — run.py's warmup loop only
dataclasses.replaces the old state; the pretrain loop's eval/save/cache-write all read the post-step state synchronously. Test loops all rebind.No numerics change; SPEC untouched (donation is a memory-aliasing property only).
Related Issue
Bridge task
donate-warmup-pretrain-steps(promoted from the 2026-07-10 JAX compile audit, donation-gaps finding). Guard precedent: #932.Motivation and Context
Peak-memory headroom is an active concern (lp-* campaign). Without donation the warmup phase pays a full extra copy of components + Adam state exactly where headroom matters, and the pretrain trainer doubles its resident state for its whole run.
How Has This Been Tested?
Pointer-aliasing tests (the fix(checkpoint): make orbax-restored TrainState donatable — solve the resume OOM (S22) #932 verification pattern):
test_llama_simple_mlp.py::test_faith_warmup_step_donates,test_pretrain.py::test_pretrain_step_donates_state,test_pretrain.py::test_restored_pretrain_state_is_donatable— ≥95% bytes-weightedunsafe_buffer_pointerreuse through the donating step, model arg survives.Full
param_decomp/tests/green at default device count ANDXLA_FLAGS="--xla_force_host_platform_device_count=4";make test,make type,make checkall clean.1-GPU H100 memory probe (donation on vs off, separate processes, 6 rebinding iterations each;
compiled.memory_analysis()+device.memory_stats()):Steady residency drops by exactly one copy of the donated state in both;
alias == donated stateproves aliasing is baked into the executable, the peak/in_use drop proves the runtime honored it (and noDONATION FAILEDline fired).Does this PR introduce a breaking change?
No.
make_faith_warmup_stepcallers may no longer reuse the passed-in components/opt state after a call (all in-tree callers rebind). The pretrain_restore_latestnow returns jit-output arrays (constructively fresh), transient 2×state at restore — the same cost restore already paid.🤖 Generated with Claude Code
Crew-Address: task/donate-warmup-pretrain-steps