From 6c371e487233aca43c0e2d96e13381af124a3072 Mon Sep 17 00:00:00 2001 From: Connor Moss Date: Mon, 14 Sep 2026 21:57:02 -0400 Subject: [PATCH] fix(optim): reject bias_correction and adam_w_mode in LAMB 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. --- bitsandbytes/optim/lamb.py | 20 ++++++++++++++++++++ tests/test_optim.py | 18 ++++++++++++++++++ 2 files changed, 38 insertions(+) diff --git a/bitsandbytes/optim/lamb.py b/bitsandbytes/optim/lamb.py index 15af97d6d..b852d156f 100644 --- a/bitsandbytes/optim/lamb.py +++ b/bitsandbytes/optim/lamb.py @@ -50,6 +50,13 @@ def __init__( max_unorm (`float`, defaults to 1.0): The maximum gradient norm. """ + # Validate unsupported parameters + if not bias_correction: + raise ValueError("LAMB does not support bias_correction=False") + + if not adam_w_mode: + raise ValueError("LAMB does not support adam_w_mode=False") + super().__init__( "lamb", params, @@ -111,6 +118,12 @@ def __init__( clipping; it is honored by the 32-bit LAMB / LAMB32bit optimizers. """ # Validate unsupported parameters + if not bias_correction: + raise ValueError("LAMB8bit does not support bias_correction=False") + + if not adam_w_mode: + raise ValueError("LAMB8bit does not support adam_w_mode=False") + if amsgrad: raise ValueError("LAMB8bit does not support amsgrad=True") @@ -176,6 +189,13 @@ def __init__( max_unorm (`float`, defaults to 1.0): The maximum gradient norm. """ + # Validate unsupported parameters + if not bias_correction: + raise ValueError("LAMB32bit does not support bias_correction=False") + + if not adam_w_mode: + raise ValueError("LAMB32bit does not support adam_w_mode=False") + super().__init__( "lamb", params, diff --git a/tests/test_optim.py b/tests/test_optim.py index 29736311d..9f5fe1bb9 100644 --- a/tests/test_optim.py +++ b/tests/test_optim.py @@ -741,3 +741,21 @@ def test_adagrad8bit_rejects_non_8_optim_bits(): bnb.optim.Adagrad8bit(p, optim_bits=32) # default (optim_bits=8) still constructs bnb.optim.Adagrad8bit(p) + + +@pytest.mark.parametrize( + "optim_cls", + [bnb.optim.LAMB, bnb.optim.LAMB8bit, bnb.optim.LAMB32bit], + ids=id_formatter("opt"), +) +@pytest.mark.parametrize("kwarg", ["bias_correction", "adam_w_mode"], ids=id_formatter("kwarg")) +def test_lamb_rejects_unsupported_flags(optim_cls, kwarg): + # Both flags were accepted by the constructors but never reached the base + # optimizer, so a non-default value was silently ignored: the update applies + # bias correction and AdamW-style decay regardless. Reject instead; mirrors + # the LAMB8bit amsgrad guard (relates to #1261). + p = [torch.nn.Parameter(torch.randn(8, 8))] + with pytest.raises(ValueError): + optim_cls(p, lr=1e-3, **{kwarg: False}) + # defaults still construct + optim_cls(p, lr=1e-3)