fix(optim): honor user-provided max_unorm in LAMB (+ 8-bit arg guards)

bitsandbytes/optim/adagrad.py

@@ -95,6 +95,7 @@ def __init__(
                 The epsilon value prevents division by zero in the optimizer.
             optim_bits (`int`, defaults to 8):
                 The number of bits of the optimizer state.
+                Note: This parameter is not used in Adagrad8bit as it always uses 8-bit optimization.
             args (`object`, defaults to `None`):
                 An object with additional arguments.
             min_8bit_size (`int`, defaults to 4096):
@@ -110,6 +111,10 @@ def __init__(
             raise ValueError("Initial accumulator value != 0.0 not supported!")
         if lr_decay != 0.0:
             raise ValueError("Lr Decay != 0.0 not supported!")
+        if optim_bits != 8:
+            # We allow the default value of 8 to maintain compatibility with the function signature,
+            # but any other value is invalid since Adagrad8bit always uses 8-bit optimization
+            raise ValueError("Adagrad8bit only supports optim_bits=8 (default value for compatibility)")
         super().__init__(
             "adagrad",
             params,

bitsandbytes/optim/lamb.py

@@ -60,7 +60,7 @@ def __init__(
             optim_bits,
             args,
             min_8bit_size,
-            max_unorm=1.0,
+            max_unorm=max_unorm,
         )
 
 
@@ -97,15 +97,23 @@ def __init__(
                 The weight decay value for the optimizer.
             amsgrad (`bool`, defaults to `False`):
                 Whether to use the [AMSGrad](https://hf.co/papers/1904.09237) variant of Adam that uses the maximum of past squared gradients instead.
+                Note: This parameter is not supported in LAMB8bit and must be False.
             adam_w_mode (`bool`, defaults to `True`):
                 Whether to use the AdamW variant.
             args (`object`, defaults to `None`):
                 An object with additional arguments.
             min_8bit_size (`int`, defaults to 4096):
                 The minimum number of elements of the parameter tensors for 8-bit optimization.
             max_unorm (`float`, defaults to 1.0):
-                The maximum gradient norm.
+                The maximum update norm for trust-ratio clipping.
+                Note: The 8-bit blockwise update does not apply update-norm clipping, so this
+                value is stored on the optimizer but has no effect on LAMB8bit. It is honored by
+                the 32-bit LAMB / LAMB32bit optimizers.
         """
+        # Validate unsupported parameters
+        if amsgrad:
+            raise ValueError("LAMB8bit does not support amsgrad=True")
+
         super().__init__(
             "lamb",
             params,
@@ -116,7 +124,7 @@ def __init__(
             8,
             args,
             min_8bit_size,
-            max_unorm=1.0,
+            max_unorm=max_unorm,
         )
 
 
@@ -172,5 +180,5 @@ def __init__(
             32,
             args,
             min_8bit_size,
-            max_unorm=1.0,
+            max_unorm=max_unorm,
         )

tests/test_optim.py

@@ -681,3 +681,53 @@ def test_ademamix_state_dict_no_nan(optim_name, optim_factory, device):
 
     for p_a, p_b in zip(model.parameters(), model2.parameters()):
         torch.testing.assert_close(p_a, p_b)
+
+
+@pytest.mark.parametrize(
+    "optim_cls", [bnb.optim.LAMB, bnb.optim.LAMB8bit, bnb.optim.LAMB32bit], ids=id_formatter("opt")
+)
+@pytest.mark.parametrize("max_unorm", [0.0, 0.5, 2.0], ids=id_formatter("max_unorm"))
+def test_lamb_max_unorm_threaded_to_config(optim_cls, max_unorm):
+    # The LAMB constructors accepted a `max_unorm` argument but passed a hardcoded 1.0
+    # to the base optimizer, so the user value never reached the optimizer config.
+    # (For LAMB8bit the value is stored but the 8-bit blockwise kernel does not apply
+    # update-norm clipping; see test_lamb_max_unorm_changes_update for the 32-bit path.)
+    p = [torch.nn.Parameter(torch.randn(8, 8))]
+    opt = optim_cls(p, max_unorm=max_unorm)
+    assert opt.args.max_unorm == max_unorm
+
+
+def test_lamb_max_unorm_changes_update():
+    # Behavioral regression: on the 32-bit LAMB path, `max_unorm` drives trust-ratio
+    # clipping of the update. Before the fix the argument was ignored (always 1.0), so a
+    # tight and a loose value produced identical updates. Runs on CPU, which implements
+    # the 32-bit LAMB kernel.
+    def one_step(max_unorm):
+        torch.manual_seed(0)
+        p = torch.nn.Parameter(torch.randn(128, 128))
+        opt = bnb.optim.LAMB([p], lr=1e-1, max_unorm=max_unorm)
+        p.grad = torch.randn(128, 128) * 5.0
+        opt.step()
+        return p.detach().clone()
+
+    tight = one_step(1e-4)  # aggressive clipping
+    loose = one_step(10.0)  # effectively no clipping
+    assert not torch.allclose(tight, loose)
+
+
+def test_lamb8bit_rejects_amsgrad():
+    # amsgrad is unused by the base optimizer; mirror the Adam8bit/AdamW8bit guard (relates to #1261).
+    p = [torch.nn.Parameter(torch.randn(8, 8))]
+    with pytest.raises(ValueError):
+        bnb.optim.LAMB8bit(p, amsgrad=True)
+    # default (amsgrad=False) still constructs
+    bnb.optim.LAMB8bit(p)
+
+
+def test_adagrad8bit_rejects_non_8_optim_bits():
+    # optim_bits is ignored (Adagrad8bit always uses 8-bit); guard invalid values (relates to #1261).
+    p = [torch.nn.Parameter(torch.randn(8, 8))]
+    with pytest.raises(ValueError):
+        bnb.optim.Adagrad8bit(p, optim_bits=32)
+    # default (optim_bits=8) still constructs
+    bnb.optim.Adagrad8bit(p)