fix(optim): reject nesterov and dampening in SGD and LARS

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,18 +95,26 @@ 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):
                 The minimum number of elements of the parameter tensors for 8-bit optimization.
             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,18 +155,26 @@ 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):
                 The minimum number of elements of the parameter tensors for 8-bit optimization.
             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__(

bitsandbytes/optim/sgd.py

@@ -29,18 +29,26 @@ 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`):
                 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("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__(

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]]
 

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))