fix(eagle3): flush trailing partial grad-accum window each epoch#2257
Merged
HuiyingLi merged 2 commits intoMay 19, 2026
Merged
Conversation
khazic
requested review from
HuiyingLi,
ZhiyuLi-Nvidia,
adil-a,
akoumpa,
athitten,
hemildesai,
pthombre and
zyzhou5
as code owners
Contributor
|
/ok to test 62f21a7 |
4 tasks
When ``num_batches_per_epoch`` is not a multiple of ``grad_accumulation_steps``, the EAGLE-3 recipe's training loop silently dropped the trailing micro-batches: their gradients reached ``backward()`` but never an ``optimizer.step()`` -- the next epoch's ``zero_grad`` wiped them. Up to ``grad_accumulation_steps - 1`` micro-batches per epoch were wasted, and the LR scheduler (sized via floor division) was off-by-N at the end of training. Two coordinated fixes in ``recipes/llm/train_eagle3.py``: 1. Compute optimizer steps with ceil division (new ``_optim_steps_per_epoch`` helper) so the LR scheduler covers the trailing flush and ``progress`` does not saturate at ``min_lr_ratio`` prematurely. 2. After the inner per-batch loop exits, if there is a partially-filled accumulation window, rescale its gradients by ``grad_accumulation_steps / pending_micro_batches`` (each micro-batch had divided its loss by the full accumulation count anticipating a full window) and run one final ``clip_grad_norm_`` / ``optimizer.step`` / ``lr_scheduler.step`` / ``zero_grad`` so the magnitude is comparable to a normal step. Tests: - ``tests/unit_tests/recipes/llm/test_train_eagle3_grad_accum.py`` exercises ``_optim_steps_per_epoch`` on divisible / non-divisible / degenerate inputs. Signed-off-by: khazic <khazzz1c@gmail.com>
khazic
force-pushed
the
fix/eagle3-flush-trailing-grad-accum
branch
from
62f21a7 to
70d2932
Compare
Contributor
|
/ok to test 70d2932 |
…nternals Signed-off-by: khazic <khazzz1c@gmail.com>
Contributor
|
/ok to test 9b5e95d |
svcnvidia-nemo-ci
added
the
waiting-on-customer
HuiyingLi
approved these changes
May 19, 2026
svcnvidia-nemo-ci
removed
the
waiting-on-customer
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
3 participants
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
When
num_batches_per_epochis not a multiple ofgrad_accumulation_steps, the EAGLE-3 recipe's training loop silently dropped the trailing micro-batches: their gradients reachedbackward()but never anoptimizer.step()— the next epoch'szero_gradwiped them. Up tograd_accumulation_steps - 1micro-batches per epoch were wasted, and the LR scheduler (sized via floor division) was off-by-N over a full run.Two coordinated fixes in
recipes/llm/train_eagle3.py:_optim_steps_per_epochhelper so the LR scheduler covers the trailing flush andprogressdoes not saturate atmin_lr_ratioprematurely.grad_accumulation_steps / pending_micro_batches(every micro-batch had divided its loss by the full accumulation count, expecting a full window) and run one finalclip_grad_norm_/optimizer.step/lr_scheduler.step/zero_gradso the trailing step's update magnitude matches every other step.Test plan
tests/unit_tests/recipes/llm/test_train_eagle3_grad_accum.pyexercising_optim_steps_per_epochon divisible / non-divisible / degenerate inputs.pytest tests/unit_tests/recipes/llm/test_train_eagle3_grad_accum.py tests/unit_tests/speculative/— 32 passed, 2 skipped (FA2 CUDA path).