From ca274a975ef866a0ad126c3888921c82e3da9cf3 Mon Sep 17 00:00:00 2001 From: Eric Eaglstun Date: Tue, 14 Jul 2026 17:27:42 -0600 Subject: [PATCH] Fix Lion to use decoupled weight decay in the Triton 32-bit kernel Follow-up to #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 #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 #1993) is parametrized over get_available_devices(), so it already exercises this kernel on XPU hardware in CI. Like #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) --- bitsandbytes/backends/triton/kernels_optim.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/bitsandbytes/backends/triton/kernels_optim.py b/bitsandbytes/backends/triton/kernels_optim.py index d7abc4af1..f7eb2e213 100644 --- a/bitsandbytes/backends/triton/kernels_optim.py +++ b/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