Add LARS to str2optimizer32bit dictionary (#1855)

* Add LARS to str2optimizer32bit dictionary

LARS optimizer was missing from str2optimizer32bit, causing KeyError
when using LARS32bit optimizer. LARS uses momentum-based kernels since
it's essentially SGD with momentum plus layerwise adaptive learning rates.

Fixes #1810

* Fix 32-bit error message and add LARS test coverage

Fix the error message in _optimizer_update_32bit_impl that incorrectly
displayed str2optimizer8bit_blockwise keys instead of str2optimizer32bit
keys. Add LARS to the parametrized 32-bit optimizer tests using
PytorchLARS as the reference implementation, with bf16 skip since
momentum kernels lack a bf16 variant.

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

---------

Co-authored-by: Tim Dettmers <tim.dettmers@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: 3 people

bitsandbytes/backends/cuda/ops.py

@@ -578,6 +578,10 @@ def _gemv_4bit_impl(
         lib.cademamix32bit_grad_fp16,
         lib.cademamix32bit_grad_bf16,
     ),
+    "lars": (
+        lib.cmomentum32bit_grad_32,
+        lib.cmomentum32bit_grad_16,
+    ),
 }
 
 str2optimizer8bit_blockwise = {
@@ -637,7 +641,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