diff --git a/bitsandbytes/optim/adagrad.py b/bitsandbytes/optim/adagrad.py index b871f2bf4..2566d1e47 100644 --- a/bitsandbytes/optim/adagrad.py +++ b/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, diff --git a/bitsandbytes/optim/lamb.py b/bitsandbytes/optim/lamb.py index 6dcfd383f..a48dcf565 100644 --- a/bitsandbytes/optim/lamb.py +++ b/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,6 +97,7 @@ 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`): @@ -104,8 +105,15 @@ def __init__( 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, ) diff --git a/tests/test_optim.py b/tests/test_optim.py index d9c24bdb5..7c64cfcec 100644 --- a/tests/test_optim.py +++ b/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)