Skip to content

Unified aero recipe: coordinate step failures and restore scheduler/scaler state - #1972

Draft
peterdsharpe wants to merge 4 commits into
NVIDIA:mainfrom
peterdsharpe:pr/recipe-train-robustness
Draft

peterdsharpe wants to merge 4 commits into
NVIDIA:mainfrom
peterdsharpe:pr/recipe-train-robustness

Conversation

@peterdsharpe

@peterdsharpe peterdsharpe commented Sep 8, 2026

Copy link
Copy Markdown
Collaborator

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_threshold rejects 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_restored and identifies missing scaler state. trajectory_exactness is unverified: RNG and stochastic data state are not restored, so this change does not promise exact stochastic continuation.

Validation

  • 36 focused training-helper tests passed, including real CPU/Gloo DDP epoch-loop tests for second-load failures on either rank in both training and validation.
  • Real save/load tests exercise _finish_epoch and reconciliation, verifying scheduler continuity and deterministic continuation while demonstrating that dropout continuation remains unverified.
  • All changed-file pre-commit hooks passed.
  • Independently reviewed by a second agent. Merged current main; no full GPU recipe run was performed.

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.
@copy-pr-bot

copy-pr-bot Bot commented Sep 8, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@github-actions

github-actions Bot commented Sep 8, 2026

Copy link
Copy Markdown
Contributor

CODEOWNERS review map

Current for commit 20679ac05ee4. An approval covers every file listed for that owner; one owner is sufficient for shared files.

@coreyjadams — 3 file(s)
  • examples/cfd/external_aerodynamics/unified_external_aero_recipe/conf/base.yaml
  • examples/cfd/external_aerodynamics/unified_external_aero_recipe/src/train.py
  • examples/cfd/external_aerodynamics/unified_external_aero_recipe/tests/test_train_helpers.py
@peterdsharpe — 3 file(s)
  • examples/cfd/external_aerodynamics/unified_external_aero_recipe/conf/base.yaml
  • examples/cfd/external_aerodynamics/unified_external_aero_recipe/src/train.py
  • examples/cfd/external_aerodynamics/unified_external_aero_recipe/tests/test_train_helpers.py

No CODEOWNER

  • CHANGELOG.md

Comment /codeowners-info to refresh.

@peterdsharpe

Copy link
Copy Markdown
Collaborator Author

/ok to test 20679ac

@peterdsharpe peterdsharpe changed the title Unified aero recipe: fail all ranks together on a step failure and resume exactly Unified aero recipe: coordinate step failures and restore scheduler/scaler state Sep 9, 2026
@peterdsharpe

Copy link
Copy Markdown
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
peterdsharpe force-pushed the pr/recipe-train-robustness branch from 7148485 to 4c53a31 Compare September 9, 2026 13:40
@peterdsharpe

Copy link
Copy Markdown
Collaborator Author

/ok to test 4c53a31

@peterdsharpe

Copy link
Copy Markdown
Collaborator Author

/ok to test e624496

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant