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

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

The LAMB, LAMB8bit and LAMB32bit 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. On the 32-bit LAMB path
(LAMB with the default optim_bits=32, and LAMB32bit) max_unorm drives
the trust-ratio clipping of the update, so this silently ignored the
setting -- e.g. a tighter max_unorm had no effect and max_unorm could
not be changed from 1.0. Thread the argument through in all three
classes.

Note: the 8-bit blockwise update kernel does not apply update-norm
clipping, so max_unorm has no numerical effect on LAMB8bit; the value
is now stored consistently and this limitation is documented in the
docstring.

Also complete the 8-bit optimizer signature/doc cleanup started for
Adam8bit/AdamW8bit (relates to #1261), mirroring their guards on
parameters the base optimizer ignores:
- LAMB8bit: raise on amsgrad=True (unused) and note it in the docstring.
- Adagrad8bit: raise on optim_bits != 8 (always 8-bit) and note it.

Tests (tests/test_optim.py):
- test_lamb_max_unorm_changes_update: behavioral check on the 32-bit
  path -- a tight vs loose max_unorm now yields different updates
  (fails on main where the value was ignored).
- test_lamb_max_unorm_threaded_to_config: the value reaches
  optimizer.args for all three LAMB classes.
- guards for LAMB8bit amsgrad and Adagrad8bit optim_bits.

Relates to #1261

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

* Reject non-default max_unorm in LAMB8bit (review feedback)

The 8-bit blockwise update does not implement update-norm clipping, so
threading max_unorm through LAMB8bit left a silent no-op: a user passing
e.g. max_unorm=0.1 would reasonably assume it takes effect. Reject any
non-default value instead, consistent with the amsgrad guard (the
default 1.0 is allowed for signature compatibility).

Config-threading test now covers the 32-bit classes only; add a guard
test for LAMB8bit max_unorm.

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Author: egeozkoc

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,29 @@ 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: This parameter is not supported in LAMB8bit and must be left at the
+                default 1.0. The 8-bit blockwise update does not implement update-norm
+                clipping; it is honored by the 32-bit LAMB / LAMB32bit optimizers.
         """
+        # Validate unsupported parameters
+        if amsgrad:
+            raise ValueError("LAMB8bit does not support amsgrad=True")
+
+        if max_unorm != 1.0:
+            # We allow the default value of 1.0 to maintain compatibility with the function
+            # signature, but the 8-bit blockwise update does not implement update-norm
+            # clipping, so any other value would be silently ignored.
+            raise ValueError("LAMB8bit only supports max_unorm=1.0 (default value for compatibility)")
+
         super().__init__(
             "lamb",
             params,
@@ -116,7 +130,7 @@ def __init__(
             8,
             args,
             min_8bit_size,
-            max_unorm=1.0,
+            max_unorm=max_unorm,
         )
 
 
@@ -172,5 +186,5 @@ def __init__(
             32,
             args,
             min_8bit_size,
-            max_unorm=1.0,
+            max_unorm=max_unorm,
         )

tests/test_optim.py

@@ -681,3 +681,63 @@ 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.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.
+    # (LAMB8bit rejects non-default values instead, since the 8-bit blockwise kernel
+    # does not apply update-norm clipping; see test_lamb8bit_rejects_non_default_max_unorm.)
+    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_lamb8bit_rejects_non_default_max_unorm():
+    # The 8-bit blockwise update does not implement update-norm clipping, so a
+    # non-default max_unorm would be silently ignored; reject it instead.
+    p = [torch.nn.Parameter(torch.randn(8, 8))]
+    with pytest.raises(ValueError):
+        bnb.optim.LAMB8bit(p, max_unorm=0.5)
+    with pytest.raises(ValueError):
+        bnb.optim.LAMB8bit(p, max_unorm=0.0)
+    # default (max_unorm=1.0) still constructs
+    bnb.optim.LAMB8bit(p, max_unorm=1.0)
+
+
+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)