fix(checkpoint): make orbax-restored TrainState donatable — solve the resume OOM (S22)

[ocg-goodfire]

Description

Three commits:

1. restore_step re-materialises the restored tree as jit outputs (non-donating identity jit pinned to reference.format via out_shardings), replacing the device_put(format) coercion. Orbax-restored arrays (per-shard device_put + make_array_from_single_device_arrays) are not reliably donatable to the jitted train step (jax#18617, never fixed upstream). When dispatch-time donation fails it is silent (the Some donated buffers were not usable warning is lowering-time only — jax/_src/interpreters/mlir.py), and the first resumed step peaks ~one full TrainState above any fresh step: the full32L resume OOM (lore 2026-07-03--resume-oom-is-buffer-donation-asymmetry). Jit outputs are indistinguishable from fresh-init state; the transient 2×-state cost equals the peak restore already paid while reference and the restored tree coexisted. Non-donating deliberately: a plain copy cannot hard-fail at dispatch (GPU errors loudly on donation of externally-referenced buffers), and the reuse donation would buy is off the step's peak. Covers the fine-tune path too (init_from_parent → restore_step).
   Plus a bytes-weighted buffer-reuse guard after the first resumed step (run.py): if <95% of state bytes were reused, each affected rank prints DONATION FAILED … — the silent failure mode is silent no more.

2. invariance_check.py fed float residuals where the target now takes token ids (same fallout class as fix(hidden-acts-eval): derive waist leading from CI output, not token inputs (S31) #919) — validation-stack item 3 was unrunnable. Now feeds randint tokens.

3. dp8 save→resume smoke config documenting the scancel-after-save → pd-lm --run_id resubmit procedure.

Related Issue

Bridge task solve-resumption-oom; upstream jax#18617.

Motivation and Context

Large-state donating runs can train fresh but die on any preemption/requeue: full32L resumes OOM'd at the default mem fraction, and for p-8e2380e2 (attn C=4096) the mem-fraction band-aid is unusable in both directions (0.97 starves NCCL clique init; even max fraction cannot fit fresh-peak + one state). The restored state must donate.

How Has This Been Tested?

• Repro harness (task workspace repro/repro_donation.py): orbax StandardSave/Restore + donating jitted step, donation observed directly via per-shard unsafe_buffer_pointer aliasing + memory_stats. On CPU, raw restored state flakily fails donation on a random subset of leaves — exclusively the sharded ones; the identity-jit re-materialisation fixes it deterministically. (GPU 1-dev and 8-proc×1-GPU multi-controller runs did not reproduce the raw failure at small scale — 48/48 rank×mode combos aliased — so CPU + the at-scale diag logs are the only places the failure is visible; the fix makes the restored arrays constructively equivalent to fresh-init ones rather than relying on that.)
• test_checkpoint.py::test_restored_state_is_donatable pins the contract (pointer aliasing through a donating jit); ran 3× stable at 4 sim devices.
• Full param_decomp/tests/ green at default device count AND --xla_force_host_platform_device_count=4; make type clean; invariance_check.py passes at 4 sim devices (worst rel 1.02e-04).
• End-to-end dp8 trainer smoke (LlamaSimpleMLP pile, scavenge jobs 170727/170729): fresh run → scancel → SIGTERM save at step 70 → pd-lm --run_id resubmit → resumed from checkpoint step 70, first resumed step trains, no DONATION FAILED line from any rank, losses continue the trajectory, per-rank peak identical (18.22 GB). Run + workspaces deleted after.
• NOT yet validated at 128-proc scale (the original OOM artifact p-05101e15 was GC'd; p-8e2380e2 is live). The guard turns any residual at-scale failure loud on the next real resume.

Does this PR introduce a breaking change?

No. restore_step now requires the reference on a single device set (one jit over the tree) — production references already are (_ensure_global); the sharded test fixture was aligned to that contract.

🤖 Generated with Claude Code