Unified aero recipe: coordinate step failures and restore scheduler/scaler state - #1972
Draft
peterdsharpe wants to merge 4 commits into
Draft
peterdsharpe wants to merge 4 commits into
peterdsharpe wants to merge 4 commits into
Conversation
A per-sample reader/transform exception or a forward-pass failure on one rank previously let that rank exit while its peers blocked in the next all_reduce until the NCCL timeout. Every step now all-reduces a failure flag before backward so the failing rank re-raises its real error and healthy ranks fail with a "peer rank failed" message instead of hanging. The loader is iterated manually so its own per-sample re-raise sits inside the same barrier as the forward pass. An optional training.divergence_loss_threshold (default null) treats a non-finite loss, or a finite loss above the threshold, as a step failure. Checkpointing: step the epoch-mode scheduler before saving so the resumed run reproduces the continuous one, persist and restore fp16 GradScaler state, always save the final epoch, and tag checkpoints with a small recipe metadata block. Checkpoints written before the tag are migrated on load by advancing the scheduler once; fp16 resumes without scaler state are reported as non-exact.
Contributor
CODEOWNERS review mapCurrent for commit ⏳ @coreyjadams — 3 file(s)
⏳ @peterdsharpe — 3 file(s)
No CODEOWNER
Comment |
Collaborator
Author
|
/ok to test 20679ac |
Collaborator
Author
|
/ok to test 7148485 |
…2-sept9 Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
peterdsharpe
force-pushed
the
pr/recipe-train-robustness
branch
from
September 9, 2026 13:40
7148485 to
4c53a31
Compare
Collaborator
Author
|
/ok to test 4c53a31 |
Collaborator
Author
|
/ok to test e624496 |
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
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
Coordinate rank-local data-loading and device-transfer failures before DDP forward, then forward/loss failures before backward. This prevents a failed loader from entering the error all-reduce while a healthy rank is inside DDP buffer synchronization or bucket rebuilding. Errors identify the phase and retain the failing rank's original exception. Failures inside model collectives, backward, or an unusable CUDA context still rely on the process-group timeout.
An optional
training.divergence_loss_thresholdrejects non-finite or excessive losses. Epoch schedulers advance before checkpointing, legacy pre-step scheduler state is migrated, fp16 GradScaler state is saved, and terminal epochs are checkpointed even off the periodic cadence.The checkpoint-resume record reports
scheduler_scaler_state_restoredand identifies missing scaler state.trajectory_exactnessisunverified: RNG and stochastic data state are not restored, so this change does not promise exact stochastic continuation.Validation
_finish_epochand reconciliation, verifying scheduler continuity and deterministic continuation while demonstrating that dropout continuation remains unverified.