[Fix] Match foreach and per-parameter rounding in SignSGD and Tiger - #520
Merged
Merged
Conversation
…ger` `_step_foreach` uses `_foreach_lerp_` while `_step_per_param` uses `mul_().add_()`. The two are equal in exact arithmetic but round differently, and a `sign()` update turns that into a full step when the momentum buffer lands near zero. Use `lerp_` in the per-parameter path so both agree.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #520 +/- ##
=========================================
Coverage 100.00% 100.00%
=========================================
Files 128 128
Lines 12444 12444
=========================================
Hits 12444 12444
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
kozistr
approved these changes
Aug 24, 2026
kozistr
left a comment
kozistr
left a comment
Owner
There was a problem hiding this comment.
looks good to me!
afair, _foreach_lerp_ landed in Pytorch 2.0, and MPS support was added in 2.4 (but, we've already used here). so maybe, it's time to bump the minimum pytorch version from 1.10 to 2.0.
anyway, thanks!
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
2 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.
Problem (Why?)
Close #519
Partly addresses #519, the
SignSGD/Tigerhalf._step_foreachuses_foreach_lerp_while_step_per_paramusesmul_().add_(). Same result in exactarithmetic, different rounding, and
sign()turns that into a full step when the buffer lands near zero.Solution (What/How?)
lerp_in_step_per_paramso both paths round the same wayOne line in each. This keeps
_foreach_lerp_as a single kernel. Moving theforeachpath to_foreach_mul_+_foreach_add_instead would work equally well if you would rather matchtorch.optim.SGD.SGDWis left alone here, since fixing it changes results for anyone using it with momentum.Other changes (bug fixes, small refactors)
N/A
Notes
Added
test_sign_based_foreach_parityfor both optimizers. It fails onmainand passes with this change.Checklist
just formatbefore commitjust checkshows no errors)just testruff checkis clean; I did not runpyright. Suite: 2494 passed before, 2496 after, same twopre-existing failures.