From 6a81fa93e9e6fefe86236a97718831b33b6f8380 Mon Sep 17 00:00:00 2001 From: Eric Eaglstun Date: Wed, 15 Jul 2026 15:58:40 -0600 Subject: [PATCH] Fix Lion to use decoupled weight decay in the CUDA 32-bit kernel Mirror of the merged upstream fix (#1993): kOptimizer32bit1State folded weight decay into the gradient (coupled L2) for every optimizer including Lion, which corrupts Lion's sign update. Exclude LION from the coupled fold and apply decoupled decay (p *= 1 - lr*wd) in the LION branch, matching the 8-bit kernel and the cpu/default backends. The default backend + test_lion32bit_weight_decay are already on this fork's main; this adds the remaining CUDA kernel piece. Co-Authored-By: Claude Opus 4.8 (1M context) --- csrc/kernels.cu | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/csrc/kernels.cu b/csrc/kernels.cu index 0d313c8d7..df0f093df 100644 --- a/csrc/kernels.cu +++ b/csrc/kernels.cu @@ -858,7 +858,11 @@ __launch_bounds__(TH, 1) __global__ void kOptimizer32bit1State( #pragma unroll 4 for (unsigned int j = 0; j < NUM_PER_THREAD; j++) { g_vals[j] = gnorm_scale * ((float)g_vals[j]); - if (weight_decay > 0.0f) + // Coupled (L2) weight decay folds decay into the gradient. Correct for + // MOMENTUM/RMSPROP/ADAGRAD, but NOT for LION, which uses decoupled + // (AdamW-style) decay applied to the param directly (see the LION branch + // below and Chen et al. 2023). This matches the CUDA 8-bit blockwise kernel. + if (weight_decay > 0.0f && OPTIMIZER != LION) g_vals[j] = (float)g_vals[j] + (((float)p_vals[j]) * weight_decay); } @@ -875,6 +879,10 @@ __launch_bounds__(TH, 1) __global__ void kOptimizer32bit1State( p_vals[j] = ((float)p_vals[j]) + update_scale * (-lr * (s1_vals[j])); break; case LION: + // Decoupled weight decay: shrink the param directly, outside the + // sign update (Chen et al. 2023), matching the 8-bit blockwise kernel. + if (weight_decay > 0.0f) + p_vals[j] = ((float)p_vals[j]) * (1.0f - lr * weight_decay); p_vals[j] = ((float)p_vals[j]) - update_scale * (lr * sgn(((float)s1_vals[j]) * beta1 + ((1.0f - beta1) * ((float)g_vals[j]))));