fix(optim): reject bias_correction and adam_w_mode in LAMB

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,

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)