fix(optim): reject bias_correction and adam_w_mode in LAMB - #2088
Open
ConnorMoss02 wants to merge 1 commit into
Open
ConnorMoss02 wants to merge 1 commit into
ConnorMoss02 wants to merge 1 commit into
Conversation
All three LAMB classes accepted `bias_correction` and `adam_w_mode` but neither reached the base optimizer, so a non-default value was silently ignored: the update applies bias correction and AdamW-style decay regardless. Reject non-default values instead, mirroring the existing LAMB8bit amsgrad guard in the same file.
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.
Problem
LAMB,LAMB8bitandLAMB32bitacceptbias_correctionandadam_w_mode, but neither reaches the base optimizer: the constructors never pass them tosuper().__init__(). The update applies bias correction and AdamW-style decay regardless, so a non-default value is silently ignored.Reproduced on CPU, 5 steps, identical seeds:
The control uses
max_unorm, which does reach the optimizer, to show the harness detects a real difference.Fix
Raise on
bias_correction=Falseandadam_w_mode=Falsein all three classes, mirroring theamsgradguard already in this file.Notes
bias_correction1/bias_correction2names in the Triton kernels are local variables computing the correction unconditionally; they do not read this flag.Tests (
tests/test_optim.py)test_lamb_rejects_unsupported_flags, parametrized over the three classes and both flags. Defaults still construct. All six cases fail with the source reverted.pre-commit run --all-filespasses; 15 LAMB tests pass.Relates to #1261. Same class as #2085, which covers
nesterovanddampeningin SGD and LARS; the two are independent and can be reviewed separately.