Fix double optimizer step in DefaultTrainingLoop with static capture - #1912
Fix double optimizer step in DefaultTrainingLoop with static capture#1912wilyan09007 wants to merge 1 commit into
Conversation
StaticCaptureTraining runs the backward pass and steps the optimizer through its grad scaler, but DefaultTrainingLoop skipped only loss.backward() and called optimizer.step() unconditionally. The gradients are still populated at that point, so every minibatch got a second update: the effective learning rate doubled, stateful optimizers advanced twice, and the update bypassed the grad scaler's inf/nan check while the minibatch count and scheduler still recorded one step. Move optimizer.step() into the eager branch so the capture wrapper is the sole owner of the update. Signed-off-by: William <wilyan090@gmail.com>
CODEOWNERS review mapCurrent for commit ⏳ @dallasfoster — 2 file(s)
⏳ @laserkelvin — 2 file(s)
No CODEOWNER
Comment |
Greptile SummaryThis PR prevents
Important Files Changed
Reviews (1): Last reviewed commit: "Fix double optimizer step in DefaultTrai..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| DefaultTrainingLoop( | ||
| enable_static_capture=enable_static_capture, | ||
| use_progress_bars=False, | ||
| use_amp=False, |
There was a problem hiding this comment.
Passing use_amp=False with a single CPU minibatch exercises only the fallback wrapper path, not the grad-scaler ownership or CUDA graph record/replay behavior central to this fix. The test can therefore pass without detecting missing or duplicate optimizer updates in those production capture modes.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
There was a problem hiding this comment.
On CPU there is no graph or enabled scaler for this test to skip. _StaticCapture.__init__ takes the CPU branch, which sets cuda_graphs_enabled = False and builds the scaler with _init_amp_scaler(False) no matter what use_amp is; use_gradscaler is only read on the CUDA branch. So turning use_amp back on would not reach graph record/replay or an enabled scaler, it would only add CPU bfloat16 autocast, and that is what I turned off so the two paths can be compared with an exact torch.equal on the parameter deltas.
Step ownership is covered. A disabled GradScaler.step(optim) still calls through to optimizer.step, so call_count == 1 on the capture path is counting the wrapper's own update, which is the behavior the fix depends on.
Enabled-scaler inf/nan skipping and graph replay are GPU-only and belong to capture.py, where test/utils/test_capture.py already parametrizes over cuda:0. What changed here is a branch in the training loop that does not vary by device or capture mode.
There was a problem hiding this comment.
You're right. The bug is a missing conditional in the training loop — optimizer.step() was called unconditionally instead of only in the eager branch — and that branch condition doesn't vary by device or capture mode. Testing it on CPU is sufficient and cleaner: it isolates exactly the invariant being fixed (step count == 1) without CUDA graph or scaler complexity that lives in a different layer.
The disabled-scaler call-through point is the key one I missed: GradScaler.step with enabled=False delegates directly to optimizer.step, so the mock's call_count accurately reflects both the wrapper's own update and any spurious second call from the loop. That's precisely what the test needs to observe.
My original concern was misplaced — the AMP/graph behavior is orthogonal to this fix and is already covered in test/utils/test_capture.py. The test is correct as written.
Tip: You can customize Greptile's behavior for this repo with .greptile/rules.md and .greptile/config.json.
1 participant
PhysicsNeMo Pull Request
Description
StaticCaptureTrainingruns the backward pass and steps the optimizer through its grad scaler, but the loop skipped onlyloss.backward()and calledoptimizer.step()unconditionally. The gradients are still populated there, so withenable_static_capture=Trueevery minibatch got a second update: the effective learning rate doubled, stateful optimizers advanced twice, and the update bypassed the grad scaler's inf/nan check, while the minibatch count and scheduler still recorded one.This moves
optimizer.step()into the eager branch, leaving the capture wrapper the sole owner of the update.test_single_optimizer_step_per_minibatchruns one deterministic minibatch each way on CPU and asserts identical parameter deltas and a singleoptimizer.stepcall; it fails on main. Every other test in that file setsenable_static_capture=False, so this path had no coverage.Unrelated and left alone: the loop's
model.zero_grad(set_to_none=True)before each wrapped step is whatcapture.pywarns can invalidate a recorded CUDA graph.Closes #1911
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.