Fix Lion to use decoupled weight decay in the Triton 32-bit kernel

Follow-up to bitsandbytes-foundation#1993, which fixed the same coupled-vs-decoupled weight-decay bug
for Lion in the default backend and the CUDA 32-bit kernel. matthewdouglas
invited a Triton follow-up when approving bitsandbytes-foundation#1993.

_optimizer_update_1state_32bit_triton_kernel folded weight decay into the
gradient (coupled L2) for every 1-state optimizer, Lion included. Lion requires
*decoupled* (AdamW-style) decay applied to the param directly (p *= 1 - lr*wd),
outside the sign update (Chen et al. 2023); folding decay into the gradient
corrupts the input to sign(), changing the update direction, not just its
magnitude. This mirrors the CUDA 32-bit fix exactly: exclude LION (id 4) from
the gradient fold, and shrink the param in the LION branch before the sign
update.

The existing test_lion32bit_weight_decay (added in bitsandbytes-foundation#1993) is parametrized over
get_available_devices(), so it already exercises this kernel on XPU hardware in
CI. Like bitsandbytes-foundation#1993's CUDA change, the Triton kernel is XPU-only and cannot be built
or run on the available hardware.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: eaglstun

bitsandbytes/backends/triton/kernels_optim.py

@@ -270,7 +270,11 @@ def _optimizer_update_1state_32bit_triton_kernel(
     s1_vals = tl.load(state1_ptr + offsets, mask=mask, other=0.0)
 
     g_vals = gnorm_scale * g_vals
-    if weight_decay > 0.0:
+    # Coupled (L2) weight decay: fold wd into the gradient. This is correct for
+    # MOMENTUM/RMSPROP/ADAGRAD, but NOT for LION (id 4), which uses *decoupled*
+    # (AdamW-style) weight decay applied to the param directly (see the LION branch
+    # below and Chen et al. 2023). LION is intentionally excluded here.
+    if OPTIMIZER_ID != 4 and weight_decay > 0.0:
         g_vals = g_vals + p_vals * weight_decay
 
     update_scale = 1.0
@@ -289,6 +293,12 @@ def _optimizer_update_1state_32bit_triton_kernel(
         p_vals = p_vals + update_val
 
     elif OPTIMIZER_ID == 4:  # LION
+        # Lion uses decoupled weight decay: shrink the param directly (p *= 1 - lr*wd)
+        # rather than folding wd into the gradient. Matches the 8-bit blockwise kernel,
+        # the default/cpu backends, and the Lion paper (Chen et al. 2023).
+        if weight_decay > 0.0:
+            p_vals = p_vals * (1.0 - lr * weight_decay)
+
         momentum_update = s1_vals * beta1 + (1.0 - beta1) * g_vals
         update_val = update_scale * lr * tl.where(momentum_update > 0, 1.0, tl.where(momentum_update < 0, -1.0, 0.0))
         p_vals = p_vals - update_val