Add LARS to str2optimizer32bit and fix error message (#1810)

LARS optimizer was missing from the str2optimizer32bit dispatch
dictionary, causing a ValueError on optimizer.step(). Map it to
momentum kernels (matching the LAMB→adam precedent). Also fix the
error message in _optimizer_update_32bit_impl which incorrectly
displayed str2optimizer8bit_blockwise keys instead of
str2optimizer32bit keys.

Add LARS to the parametrized 32-bit optimizer tests with
PytorchLARS as the reference implementation.

Fixes #1810

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: TimDettmers

bitsandbytes/backends/cuda/ops.py

@@ -574,6 +574,10 @@ def _gemv_4bit_impl(
         lib.cademamix32bit_grad_fp16,
         lib.cademamix32bit_grad_bf16,
     ),
+    "lars": (
+        lib.cmomentum32bit_grad_32,
+        lib.cmomentum32bit_grad_16,
+    ),
 }
 
 str2optimizer8bit_blockwise = {
@@ -633,7 +637,7 @@ def _optimizer_update_32bit_impl(
     optim_fns = str2optimizer32bit.get(optimizer_name, None)
     if optim_fns is None:
         raise ValueError(
-            f"Unsupported optimizer name: {optimizer_name}. Supported optimizers: {list(str2optimizer8bit_blockwise.keys())}"
+            f"Unsupported optimizer name: {optimizer_name}. Supported optimizers: {list(str2optimizer32bit.keys())}"
         )
     if g.dtype == torch.float32:
         optim_func = optim_fns[0]

tests/test_optim.py

@@ -102,6 +102,11 @@ def rm_path(path):
     lambda pxx: bnb.optim.SGD8bit(pxx, 0.01, 0.9, block_wise=True),
 )
 
+str2optimizers["lars"] = (
+    lambda pxx: bnb.optim.PytorchLARS(pxx, 0.01, 0.9),
+    lambda pxx: bnb.optim.LARS(pxx, 0.01, 0.9),
+)
+
 str2optimizers["rmsprop"] = (
     lambda pxx: torch.optim.RMSprop(pxx, 0.01, 0.9),
     lambda pxx: bnb.optim.RMSprop(pxx, 0.01, 0.9, block_wise=False),
@@ -118,6 +123,7 @@ def rm_path(path):
 str2statenames["lion"] = [("exp_avg", "state1")]
 str2statenames["paged_lion"] = [("exp_avg", "state1")]
 str2statenames["momentum"] = [("momentum_buffer", "state1")]
+str2statenames["lars"] = [("momentum_buffer", "state1")]
 str2statenames["lamb"] = [("exp_avg", "state1"), ("exp_avg_sq", "state2")]
 str2statenames["rmsprop"] = [("square_avg", "state1")]
 
@@ -155,6 +161,7 @@ def rm_path(path):
     "paged_adamw",
     "paged_adam",
     "momentum",
+    "lars",
     "rmsprop",
     "lion",
     "paged_lion",
@@ -181,7 +188,7 @@ def test_optimizer32bit(dim1, dim2, gtype, optim_name, device):
     if optim_name.startswith("paged_") and device == "xpu":
         pytest.skip("Paged optimizers are not supported on XPU currently.")
 
-    if gtype == torch.bfloat16 and optim_name in ["momentum", "rmsprop"]:
+    if gtype == torch.bfloat16 and optim_name in ["momentum", "lars", "rmsprop"]:
         pytest.skip()
     if dim1 == 1 and dim2 == 1:
         return