train: flush the trailing partial grad-accum window - #8
Merged
Conversation
Floor division dropped any trailing partial accumulation window: its microbatches ran backward() but the grads were never stepped. Worst case, a run with fewer microbatches than grad_accum (tiny smoke runs) trained ZERO steps and returned []. Now total_steps uses ceil and the partial window is flushed after the loop, keeping the uniform 1/grad_accum grad scaling (the short window takes a proportionally smaller step, matching common practice). NB: changes training numerics slightly vs the banked lr3e-5_ep1 / lr1e-4_ep2 runs whenever n_micro % grad_accum != 0 — one extra (small) optimizer step and a one-step shift in the lr schedule. Tests: partial-window flush now expects the extra step; new test pins the shorter-than-one-window case (1 microbatch, grad_accum=4 -> 1 step, adapter actually moves). Co-Authored-By: Claude Fable 5 <[email protected]>
dtch1997
added a commit
that referenced
this pull request
Jul 17, 2026
Full pipeline re-run on main (post PRs #6/#7/#8) — all six reproduction criteria pass. Fresh summary/ia_validation/figures + REPRO.md committed. Verdict: REPRODUCED Co-authored-by: dtch1997 <[email protected]> Co-authored-by: Claude Opus 4.8 <[email protected]>
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
1 participant
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.
Fixes the tail-drop quirk flagged in #7.
The bug
total_steps = n_micro // grad_accum(floor) meant any trailing partial accumulation window ranbackward()on its microbatches but never got an optimizer step — those examples silently contributed nothing. The worst case: a run with fewer microbatches thangrad_accum(easy to hit in smoke configs) trained zero steps and returned an empty loss list.The fix
total_stepsusesceil, and after the epoch loop a non-empty partial window is flushed: clip, step, scheduler step, loss appended. Gradient scaling stays the uniform1/grad_accum, so the short window takes a proportionally smaller step rather than an up-weighted one — matching common practice for hand-rolled accumulation loops.Reproducibility note
This intentionally changes training numerics relative to the banked
lr3e-5_ep1/lr1e-4_ep2results whenevern_micro % grad_accum != 0: one extra small optimizer step and a one-step shift in the lr schedule. For the full config (2000 rows, batch 8, accum 4 → 250 microbatches) that's a 2-microbatch flush after 62 full steps — negligible, but a fresh run is no longer bit-identical to the banked ones.Tests
test_partial_accum_window_flushes: 3 microbatches / accum 2 → 2 optimizer steps (was 1 + dropped tail).test_run_shorter_than_one_window_still_trains: 1 microbatch / accum 4 → 1 step, and asserts the adapter's B matrices actually moved (this returned[]and trained nothing before the fix).17 tests, all green locally; CI covers both ends of the peft range.
🤖 Generated with Claude Code