Make checkpointing better#1647

Draft

finbarrtimbers wants to merge 11 commits into

mainfrom

finbarr/check-script-in

Commits on Apr 29, 2026

Extract checkpoint-state save into maybe_save_checkpoint_state and log size/bandwidth Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
finbarrtimbers
committed
Cut redundant work in checkpoint save: per-rank driver state, ref_policy rewrites, RNG dup Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
finbarrtimbers
committed
Drop legacy checkpoint fallbacks for driver_state.pt and per-device CUDA RNG Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
finbarrtimbers
committed
Replace _ref_policy_dirty with should_save_ref_policy predicate; restore driver_state guard Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
finbarrtimbers
committed
Assert checkpoint_state_freq divides ref_policy_update_freq so saves land on update steps Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
finbarrtimbers
committed
Restore WHY comments around ref_policy raw torch.save and deepspeed mpu/save dance Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
finbarrtimbers
committed
Add CHANGELOG entry for #1647 Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
finbarrtimbers
committed