Repository navigation
Unified recipe: log per-step losses without waiting for the GPU - #2171
Open
peterdsharpe wants to merge 14 commits into
Open
peterdsharpe wants to merge 14 commits into
peterdsharpe wants to merge 14 commits into
Conversation
step_time_s is host time up to the per-step loss reduction, so it misses the device work still queued and the wait for other ranks. The reduction already syncs, so the clock right after it gives the true time per iteration at no cost: log it as iter_time_s (console Iter:, TensorBoard, JSONL train_step and val_step records). Epoch summary records gain epoch_seconds and, on GPU, peak_memory_gib (peak allocated memory, maximum over ranks). Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
The loader does host-side work for each batch inside next() (the lookahead sample's host-to-device copy and transforms). The loop called next() at the start of each step, when the GPU had drained, so the GPU sat idle for that work. Fetch batch i + 1 right after the backward is launched, before the optimizer step. Same samples in the same order. Under cfg.profile the loop now stops at the same step by its bound, so no extra batch is loaded at the stop. Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
Contributor
CODEOWNERS review mapCurrent for commit ⏳ @coreyjadams — 3 file(s)
⏳ @peterdsharpe — 3 file(s)
No CODEOWNER
Comment |
Contributor
|
[Medium impact] The PR appears safe to merge; no actionable new issue was established. SummaryThe recipe now starts each step's readback without waiting, logs completed records in order, and waits for remaining records at the epoch's end.
Reviews (2) · Last reviewed commit: "Recipe: log landed steps on an error thr..." · Reviewed by Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Reset the CUDA peak statistic at the start of _run_epoch instead of after the read, so the first epoch's peak_memory_gib no longer includes allocations made before it, such as loading a checkpoint. Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
peterdsharpe
marked this pull request as draft
October 26, 2026 16:32
Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
Drop the hand prefetch, iter_time_s, epoch_seconds and peak_memory_gib; the next commit replaces them with the fix for the per-step host sync. Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
Each training step read its loss back to the host (`.cpu()` in `_reduce_and_average`), which waited for the GPU just before the loop asked the datapipe for the next batch, so its prefetch work ran on an idle GPU. The per-step records now stay on the device; every `logging.log_every_n_steps` steps one fused all-reduce of the stacked records is copied to pinned host memory without blocking, and each record is logged once its copy has landed, at most one per step. The epoch's end logs the rest. The lr is read at its own step. Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
peterdsharpe
changed the title
… land Replace the per-period stacked readback with one non-blocking readback per step (_start_step_readback). The loop logs every earlier step whose copy has landed, in order, through the log_step closure, and the epoch's end waits for the rest. Each step does the same all-reduce as main, so the logged values equal main's for any number of ranks. Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
An error inside the step loop (OOM, Ctrl-C, preemption) skipped the epoch-end drain and lost every pending record. The loop now logs, in order and without waiting, the records whose copies have already reached the host, then re-raises. The in-loop and error paths share log_landed(). Pending records are a typed _PendingStep NamedTuple instead of a positional tuple, log_step reads them as landed host copies, the collective note is back at the readback call, and _reduce_and_average's docstring no longer describes a per-step caller. The CHANGELOG says that step_time_s now includes the host's wait for the GPU. Tests: a step that raises logs only the landed records; a 2-rank gloo test checks that the readback is the all-rank mean; the CUDA no-sync test restores the previous sync debug mode; the source-inspection test is removed, since the no-sync test covers it. Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
…nd wait for the real copy in the test's held events The error path is an ExitStack callback on the existing with line instead of a nine-line context manager: on any exit it logs, without waiting, the steps whose copies have landed, which after a normal epoch end is none. The tests' held readback events now wrap the real CUDA event and wait on it before a record is read. Before, they dropped it, so on CUDA a drained record could be read before its non_blocking copy had landed, and test_epoch_end_drains_records_still_in_flight[val-cuda] failed about once in five runs. Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
peterdsharpe
marked this pull request as ready for review
October 27, 2026 04:08
Conflicts in train.py, test_train_helpers.py and CHANGELOG.md. - train.py: the per-step readback takes main's metric_counts (NVIDIA#2244), so a step whose field has no finite target still logs NaN for it, and keeps main's materialized loss (NVIDIA#1983) and output_layout (NVIDIA#2254). - test_train_helpers.py: main's tests kept; the per-step tests pass a metric calculator stub (main's loop reads mask_nonfinite_targets) and also run with masked targets, against the blocking readback with metric_counts. - CHANGELOG.md: both entries kept. Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
…rd pass in the no-sync test The 2-rank readback test adds a metric whose target is finite on rank 0 only, so the per-step readback must average it over that rank alone, as _reduce_and_average does. The CUDA no-sync test now runs the real forward_pass, LossCalculator and MetricCalculator, weighted by area with masked targets, so a sync added to the loss or metric code fails it. Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
The CHANGELOG entry is two sentences. The README's Logging section says that each step's records are written once its values reach the host, usually about one step later, and what step_time_s now measures. The _run_epoch docstring names the train_step phase and the deferred logging. Signed-off-by: Peter Sharpe <peterdsharpe@gmail.com>
Collaborator
Author
|
/ok to test 8316d0c |
This branch has not been deployed
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.
Summary
The unified recipe's training loop copied each step's loss to the host, which made the host wait for the GPU at every step. That wait came just before the loop asked the datapipe for the next batch, so the datapipe's own prefetch work ran on an idle GPU. Each step now reads its values back without blocking, and the prefetch overlaps compute again. Transolver steps take 8.7% less time and GeoTransolver steps 1.2% less, with bitwise identical losses.
What changed
In
examples/cfd/external_aerodynamics/unified_external_aero_recipe/src/train.py:_run_epochno longer calls_reduce_and_averageper step. That call's.cpu()was the per-step sync, as noted in an earlier review._start_step_readbackruns the same fused all-reduce as main's per-step path, then anon_blockingcopy into pinned host memory, and records a CUDA event._start_step_readbackalso takes each field's count of finite targets (metric_counts, from Unified external aero recipe: weight loss and metrics by cell area or volume #2244), as_reduce_and_averagedoes. A step whose field has no finite target still logs NaN for that field.log_stepclosure holds main's per-step logging block, with the same output. Its threeadd_scalarcalls are folded into one loop, and two comments are shorter. Each step's lr is read at its own step._reduce_and_averagenow serves only the epoch summary.The recipe README's Logging section now says that per-step records arrive after their step, and what
step_time_smeasures.Results
These runs predate main's #1983, #2254 and #2244. They use main's
geotransolver_volumeandtransolver_volumeconfigs on DrivAerML volume (200k points, compiled, bf16). Each model ran on 2 GB300 GPUs for 50 steps per run, and each version ran twice in turn:mainmainrun read its samples from a cold file cache, so its row uses the secondmainrun.mainbitwise. A 4-GPU Transolver run also matchesmainbitwise.Limits
step_time_snow includes the host's wait for the GPU inside the step, which main's left out. So it reads higher than on main even though steps are faster. On Transolver it went from 94.3 to 102.1 ms, while wall time per step went from 117.4 to 107.2 ms. Use epoch wall time per step to compare speed.Tests
New tests in
tests/test_train_helpers.pycheck that every step is logged once, in order, with the blocking readback's values bitwise, with and without masked targets. They cover the records logged at the epoch's end, the profile stop and the per-step lr. A CUDA test runs the recipe's real forward pass, loss and metrics undertorch.cuda.set_sync_debug_mode("error"), and counts explicit event, stream and device waits. Another test makes a step raise mid-epoch and checks that only the records already landed are logged, in order. A 2-process gloo test checks that each step's readback is the all-rank mean, also when only one rank has a finite target for a field. CI skipsexamples/, so I ran the recipe tests and pre-commit locally, and they pass.Checklist
Dependencies
None.
Review Process
All PRs are reviewed by the PhysicsNeMo team before merging.
Depending on which files are changed, GitHub may automatically assign a maintainer for review.
We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.
AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.