enable xpu lars

Signed-off-by: jiqing-feng <jiqing.feng@intel.com>

Author: jiqing-feng

bitsandbytes/backends/default/ops.py

@@ -357,9 +357,11 @@ def _(
 
 name2optimizer_id = {
     "momentum": MOMENTUM,
+    "lars": MOMENTUM,
     "rmsprop": RMSPROP,
     "adagrad": ADAGRAD,
     "adam": ADAM,
+    "lamb": ADAM,
     "lion": LION,
     "ademamix": ADEMAMIX,
 }

bitsandbytes/backends/triton/kernels_optim.py

@@ -24,9 +24,11 @@
 
 name2optimizer_id = {
     "momentum": MOMENTUM,
+    "lars": MOMENTUM,
     "rmsprop": RMSPROP,
     "adagrad": ADAGRAD,
     "adam": ADAM,
+    "lamb": ADAM,
     "lion": LION,
     "ademamix": ADEMAMIX,
 }
@@ -313,6 +315,10 @@ def _optimizer_update_1state_32bit_triton_kernel(
         "preprocess": _optimizer_precondition_2state_32bit,
         "update": _optimizer_update_2state_32bit_triton_kernel,
     },
+    "lamb": {
+        "preprocess": _optimizer_precondition_2state_32bit,
+        "update": _optimizer_update_2state_32bit_triton_kernel,
+    },
     "ademamix": {
         "preprocess": _optimizer_precondition_2state_32bit,
         "update": _optimizer_update_2state_32bit_triton_kernel,
@@ -321,6 +327,10 @@ def _optimizer_update_1state_32bit_triton_kernel(
         "preprocess": _optimizer_precondition_1state_32bit,
         "update": _optimizer_update_1state_32bit_triton_kernel,
     },
+    "lars": {
+        "preprocess": _optimizer_precondition_1state_32bit,
+        "update": _optimizer_update_1state_32bit_triton_kernel,
+    },
     "rmsprop": {
         "preprocess": _optimizer_precondition_1state_32bit,
         "update": _optimizer_update_1state_32bit_triton_kernel,
@@ -1065,9 +1075,11 @@ def _optimizer_update_2state_8bit_blockwise_triton_kernel(
 
 name2optimizer_fn = {
     "momentum": _optimizer_update_1state_8bit_blockwise_triton_kernel,
+    "lars": _optimizer_update_1state_8bit_blockwise_triton_kernel,
     "rmsprop": _optimizer_update_1state_8bit_blockwise_triton_kernel,
     "adagrad": _optimizer_update_1state_8bit_blockwise_triton_kernel,
     "adam": _optimizer_update_2state_8bit_blockwise_triton_kernel,
+    "lamb": _optimizer_update_2state_8bit_blockwise_triton_kernel,
     "lion": _optimizer_update_1state_8bit_blockwise_triton_kernel,
     "ademamix": _optimizer_update_2state_8bit_blockwise_triton_kernel,
 }