diff --git a/bitsandbytes/optim/lars.py b/bitsandbytes/optim/lars.py index c2f5aa784..c83b247d8 100644 --- a/bitsandbytes/optim/lars.py +++ b/bitsandbytes/optim/lars.py @@ -33,11 +33,12 @@ def __init__( momentum (`float`, defaults to 0): The momentum value speeds up the optimizer by taking bigger steps. dampening (`float`, defaults to 0): - The dampening value reduces the momentum of the optimizer. + The dampening value reduces the momentum of the optimizer. Not supported; only + the default of 0 is accepted. weight_decay (`float`, defaults to 1e-2): The weight decay value for the optimizer. nesterov (`bool`, defaults to `False`): - Whether to use Nesterov momentum. + Whether to use Nesterov momentum. Not supported; only the default of `False` is accepted. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): @@ -47,6 +48,13 @@ def __init__( max_unorm (`float`, defaults to 0.02): The maximum gradient norm. """ + # Validate unsupported parameters + if nesterov: + raise ValueError("LARS does not support nesterov=True") + + if dampening != 0: + raise ValueError("LARS does not support dampening != 0") + if momentum == 0: raise NotImplementedError("LARS without momentum is not supported!") super().__init__( @@ -87,11 +95,12 @@ def __init__( momentum (`float`, defaults to 0): The momentum value speeds up the optimizer by taking bigger steps. dampening (`float`, defaults to 0): - The dampening value reduces the momentum of the optimizer. + The dampening value reduces the momentum of the optimizer. Not supported; only + the default of 0 is accepted. weight_decay (`float`, defaults to 1e-2): The weight decay value for the optimizer. nesterov (`bool`, defaults to `False`): - Whether to use Nesterov momentum. + Whether to use Nesterov momentum. Not supported; only the default of `False` is accepted. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): @@ -99,6 +108,13 @@ def __init__( max_unorm (`float`, defaults to 0.02): The maximum gradient norm. """ + # Validate unsupported parameters + if nesterov: + raise ValueError("LARS8bit does not support nesterov=True") + + if dampening != 0: + raise ValueError("LARS8bit does not support dampening != 0") + if momentum == 0: raise NotImplementedError("LARS without momentum is not supported!") super().__init__( @@ -139,11 +155,12 @@ def __init__( momentum (`float`, defaults to 0): The momentum value speeds up the optimizer by taking bigger steps. dampening (`float`, defaults to 0): - The dampening value reduces the momentum of the optimizer. + The dampening value reduces the momentum of the optimizer. Not supported; only + the default of 0 is accepted. weight_decay (`float`, defaults to 1e-2): The weight decay value for the optimizer. nesterov (`bool`, defaults to `False`): - Whether to use Nesterov momentum. + Whether to use Nesterov momentum. Not supported; only the default of `False` is accepted. args (`object`, defaults to `None`): An object with additional arguments. min_8bit_size (`int`, defaults to 4096): @@ -151,6 +168,13 @@ def __init__( max_unorm (`float`, defaults to 0.02): The maximum gradient norm. """ + # Validate unsupported parameters + if nesterov: + raise ValueError("LARS32bit does not support nesterov=True") + + if dampening != 0: + raise ValueError("LARS32bit does not support dampening != 0") + if momentum == 0: raise NotImplementedError("LARS without momentum is not supported!") super().__init__( diff --git a/bitsandbytes/optim/sgd.py b/bitsandbytes/optim/sgd.py index 75fc71474..dcc5ff68b 100644 --- a/bitsandbytes/optim/sgd.py +++ b/bitsandbytes/optim/sgd.py @@ -29,11 +29,12 @@ def __init__( momentum (`float`, defaults to 0): The momentum value speeds up the optimizer by taking bigger steps. dampening (`float`, defaults to 0): - The dampening value reduces the momentum of the optimizer. + The dampening value reduces the momentum of the optimizer. Not supported; only + the default of 0 is accepted. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. nesterov (`bool`, defaults to `False`): - Whether to use Nesterov momentum. + Whether to use Nesterov momentum. Not supported; only the default of `False` is accepted. optim_bits (`int`, defaults to 32): The number of bits of the optimizer state. args (`object`, defaults to `None`): @@ -41,6 +42,13 @@ def __init__( min_8bit_size (`int`, defaults to 4096): The minimum number of elements of the parameter tensors for 8-bit optimization. """ + # Validate unsupported parameters + if nesterov: + raise ValueError("SGD does not support nesterov=True") + + if dampening != 0: + raise ValueError("SGD does not support dampening != 0") + if momentum == 0: raise NotImplementedError("SGD without momentum is not supported!") super().__init__( @@ -79,16 +87,24 @@ def __init__( momentum (`float`, defaults to 0): The momentum value speeds up the optimizer by taking bigger steps. dampening (`float`, defaults to 0): - The dampening value reduces the momentum of the optimizer. + The dampening value reduces the momentum of the optimizer. Not supported; only + the default of 0 is accepted. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. nesterov (`bool`, defaults to `False`): - Whether to use Nesterov momentum. + Whether to use Nesterov momentum. Not supported; only the default of `False` is accepted. 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. """ + # Validate unsupported parameters + if nesterov: + raise ValueError("SGD8bit does not support nesterov=True") + + if dampening != 0: + raise ValueError("SGD8bit does not support dampening != 0") + if momentum == 0: raise NotImplementedError("SGD without momentum is not supported!") super().__init__( @@ -127,16 +143,24 @@ def __init__( momentum (`float`, defaults to 0): The momentum value speeds up the optimizer by taking bigger steps. dampening (`float`, defaults to 0): - The dampening value reduces the momentum of the optimizer. + The dampening value reduces the momentum of the optimizer. Not supported; only + the default of 0 is accepted. weight_decay (`float`, defaults to 0.0): The weight decay value for the optimizer. nesterov (`bool`, defaults to `False`): - Whether to use Nesterov momentum. + Whether to use Nesterov momentum. Not supported; only the default of `False` is accepted. 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. """ + # Validate unsupported parameters + if nesterov: + raise ValueError("SGD32bit does not support nesterov=True") + + if dampening != 0: + raise ValueError("SGD32bit does not support dampening != 0") + if momentum == 0: raise NotImplementedError("SGD without momentum is not supported!") super().__init__( diff --git a/docs/source/reference/optim/sgd.mdx b/docs/source/reference/optim/sgd.mdx index a0d09d1e8..55814145c 100644 --- a/docs/source/reference/optim/sgd.mdx +++ b/docs/source/reference/optim/sgd.mdx @@ -2,7 +2,7 @@ Stochastic gradient descent (SGD) is a basic gradient descent optimizer to minimize loss given a set of model parameters and updates the parameters in the opposite direction of the gradient. The update is performed on a randomly sampled mini-batch of data from the dataset. -bitsandbytes also supports momentum and Nesterov momentum to accelerate SGD by adding a weighted average of past gradients to the current gradient. +bitsandbytes supports momentum to accelerate SGD by adding a weighted average of past gradients to the current gradient. Nesterov momentum is not supported. ## SGD[[api-class]] diff --git a/tests/test_optim.py b/tests/test_optim.py index 29736311d..a3c7b0b54 100644 --- a/tests/test_optim.py +++ b/tests/test_optim.py @@ -741,3 +741,50 @@ 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.SGD, bnb.optim.SGD8bit, bnb.optim.SGD32bit, bnb.optim.LARS, bnb.optim.LARS8bit, bnb.optim.LARS32bit], + ids=id_formatter("opt"), +) +def test_sgd_lars_reject_nesterov(optim_cls): + # These constructors accepted a `nesterov` argument that the base optimizer never + # reads, so nesterov=True silently produced plain heavy-ball momentum. Reject it + # instead; mirrors the Adam8bit/AdamW8bit guards (relates to #1261). + p = [torch.nn.Parameter(torch.randn(8, 8))] + with pytest.raises(ValueError): + optim_cls(p, lr=1e-3, momentum=0.9, nesterov=True) + # default (nesterov=False) still constructs + optim_cls(p, lr=1e-3, momentum=0.9) + + +@pytest.mark.parametrize( + "optim_cls", + [bnb.optim.SGD, bnb.optim.SGD8bit, bnb.optim.SGD32bit, bnb.optim.LARS, bnb.optim.LARS8bit, bnb.optim.LARS32bit], + ids=id_formatter("opt"), +) +def test_sgd_lars_reject_dampening(optim_cls): + # `dampening` was passed to the base optimizer as betas[1], which no momentum path + # reads, so a non-zero value had no effect on the update. + p = [torch.nn.Parameter(torch.randn(8, 8))] + with pytest.raises(ValueError): + optim_cls(p, lr=1e-3, momentum=0.9, dampening=0.5) + # default (dampening=0) still constructs + optim_cls(p, lr=1e-3, momentum=0.9) + + +def test_pytorch_lars_still_supports_nesterov(): + # PytorchLARS implements nesterov itself and validates it, so it keeps the argument. + # Guards the boundary of this change: a tight and a loose setting must still differ. + def one_step(nesterov): + torch.manual_seed(0) + p = torch.nn.Parameter(torch.randn(32, 32)) + opt = bnb.optim.PytorchLARS([p], lr=1e-1, momentum=0.9, nesterov=nesterov) + torch.manual_seed(123) + for _ in range(2): + p.grad = torch.randn(32, 32) + opt.step() + return p.detach().clone() + + assert not torch.allclose(one_step(False), one_step(True))