Checkpoint per-query MAGIC backwards mid-walk so preemption loses one segment - #454
Open
luciaquirke wants to merge 1 commit into
Open
Checkpoint per-query MAGIC backwards mid-walk so preemption loses one segment#454luciaquirke wants to merge 1 commit into
luciaquirke wants to merge 1 commit into
Conversation
… segment
A single per-query backward can run for days, but unlike the aggregate
backward it never saved its BackwardState, so any interruption redid the
whole trajectory walk. Save the state (atomic tmp+rename, as before) about
once per rematerialization segment by default, or at the configured
backward_save_every cadence, and resume from it, under a per-query filename
(backward_q{i}_rank{r}.pt) so one query's stale state can never be resumed
by another.
Harden the shared resume path while here: an unreadable state file falls
back to a fresh backward with a warning instead of crashing, and ranks
verify they agree on the resume step before using saved state — a crash
mid-cadence can leave ranks holding different steps, which would deadlock
the replay collectives.
Co-Authored-By: Claude Fable 5 <[email protected]>
Claude-Session: https://claude.ai/code/session_016ZCeBy6TkkCW1BwC9peVWj
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 join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Related to #453 (same accumulators-are-checkpointable hardening theme as #449–#452).
A per-query MAGIC backward can run 19–72h, but unlike the aggregate backward it never saved its
BackwardState— an interruption redid the entire trajectory walk (~200 GPU-h lost to this in the last 48h). This makes the per-query backward resumable at checkpoint-segment granularity.Changes
compute_per_query_magic_scoresnow passesresume/save_everythrough toTrainer.backward, defaulting the cadence to about one sqrt save-mode rematerialization segment (isqrt(num_steps)backward steps) whenbackward_save_everyis 0, so an interrupted query loses at most ~one segment. An explicitbackward_save_everyis honored as-is.backward_q{i}_rank{r}.ptvia a newstate_prefixparameter onTrainer.backward) so a stale state file from one query can never be resumed by another. The files use the existing atomic tmp+rename save and are deleted on successful completion, as before.load_backward_statedeletes incomplete checkpoints while filtering.Testing
test_per_query_backward_resumes_mid_query: preempts query 0 after its second mid-backward save, asserts the state file exists, then resumes and checks the scores match an uninterrupted run exactly and the state file is cleaned up.test_per_query_backward_ignores_corrupt_state: a garbage state file produces the warning and correct scores.tests/test_per_query_magic.py,test_magic.py,test_save_mode_final.py,test_step_state.py,test_grad_clipping.py: 64 passed, 2 skipped (CUDA).The rank-consensus check only exercises under
dist; it is intentionally minimal (two all-reduces on a scalar). Not covered by a unit test since the suite here runs single-process.🤖 Generated with Claude Code
https://claude.ai/code/session_016ZCeBy6TkkCW1BwC9peVWj