Skip to content

fix(optim): reject bias_correction and adam_w_mode in LAMB - #2088

Open
ConnorMoss02 wants to merge 1 commit into
bitsandbytes-foundation:mainfrom
ConnorMoss02:fix/lamb-reject-unsupported-args
Open

ConnorMoss02 wants to merge 1 commit into
bitsandbytes-foundation:mainfrom
ConnorMoss02:fix/lamb-reject-unsupported-args

Conversation

@ConnorMoss02

Copy link
Copy Markdown

Problem

LAMB, LAMB8bit and LAMB32bit accept bias_correction and adam_w_mode, but neither reaches the base optimizer: the constructors never pass them to super().__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:

class         bias_correction T vs F    adam_w_mode T vs F
LAMB                        0.000000              0.000000
LAMB8bit                    0.000000              0.000000
LAMB32bit                   0.000000              0.000000

control: LAMB max_unorm 1e-4 vs 10.0  ->  max|diff| = 0.499796

The control uses max_unorm, which does reach the optimizer, to show the harness detects a real difference.

Fix

Raise on bias_correction=False and adam_w_mode=False in all three classes, mirroring the amsgrad guard already in this file.

Notes

  • No public API changes. The guards only reject values that were already silently ignored, and no code in the repo constructs these classes with them.
  • The bias_correction1 / bias_correction2 names 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-files passes; 15 LAMB tests pass.

Relates to #1261. Same class as #2085, which covers nesterov and dampening in SGD and LARS; the two are independent and can be reviewed separately.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant