Fix Lion to use decoupled weight decay (default + CUDA 32-bit)

bitsandbytes/backends/default/ops.py

@@ -450,7 +450,11 @@ def _optimizer_update_32bit(
 
     p_vals = p.float()
     g_vals = (gnorm_scale * g).float()
-    if optimizer_id in [0, 1, 2, 4] and 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 in [0, 1, 2] and weight_decay > 0.0:
         g_vals = g_vals + p_vals * weight_decay
 
     update_scale = 1.0
@@ -515,6 +519,12 @@ def _optimizer_update_32bit(
         state1.copy_(s1_vals)
 
     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 cpu backend, the CUDA
+        # 8-bit blockwise kernel, and the Lion paper (Chen et al. 2023).
+        if weight_decay > 0.0:
+            p_vals = p_vals * (1.0 - lr * weight_decay)
+
         momentum_update = state1 * beta1 + (1.0 - beta1) * g_vals
         update_val = update_scale * lr * torch.sign(momentum_update)
         p_vals = p_vals - update_val

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]))));

tests/test_optim.py

@@ -243,6 +243,55 @@ def test_optimizer32bit(dim1, dim2, gtype, optim_name, device):
             assert bnb_optimizer.state[p2]["unorm_vec"] > 0.0
 
 
+@pytest.mark.parametrize("gtype", [torch.float32, torch.float16, torch.bfloat16], ids=describe_dtype)
+@pytest.mark.parametrize("dim1", [1024], ids=id_formatter("dim1"))
+@pytest.mark.parametrize("dim2", [32, 1024], ids=id_formatter("dim2"))
+@pytest.mark.parametrize("device", get_available_devices(), ids=id_formatter("device"))
+def test_lion32bit_weight_decay(dim1, dim2, gtype, device):
+    """Lion must use *decoupled* weight decay (p *= 1 - lr*wd), matching the Lion paper
+    and the lion_pytorch reference.
+
+    Regression test for a coupled-vs-decoupled weight-decay bug: the shared 32-bit
+    optimizer path (used by the `default` backend, i.e. MPS and any device without a
+    dedicated kernel) folded weight decay into the gradient (coupled L2) for Lion, which
+    corrupts the sign update. The existing test_optimizer32bit never caught it because it
+    constructs optimizers with the default weight_decay=0, where the two forms coincide.
+    """
+    weight_decay = 0.1
+
+    p1 = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.1
+    p2 = p1.clone()
+    p1 = p1.float()
+
+    torch_optimizer = Lion([p1], weight_decay=weight_decay)
+    bnb_optimizer = bnb.optim.Lion([p2], weight_decay=weight_decay)
+
+    if gtype == torch.float32:
+        atol, rtol = 1e-6, 1e-5
+    elif gtype == torch.bfloat16:
+        atol, rtol = 1e-3, 1e-2
+    else:
+        atol, rtol = 1e-4, 1e-3
+
+    for i in range(k):
+        g = torch.randn(dim1, dim2, device=device, dtype=gtype) * 0.01
+        p1.grad = g.clone().float()
+        p2.grad = g.clone()
+
+        bnb_optimizer.step()
+        torch_optimizer.step()
+
+        # Lion's sign update is noisy at boundaries, so allow a few mismatches (matches
+        # the tolerance policy in test_optimizer32bit). Under the coupled-decay bug this
+        # diverges far beyond the error budget; under the correct decoupled decay it holds.
+        assert_most_approx_close(p1, p2.float(), atol=atol, rtol=rtol, max_error_count=15)
+
+        if gtype != torch.float32:
+            # keep 16-bit params from drifting apart across steps (see test_optimizer32bit)
+            p1.data = p1.data.to(p2.dtype).float()
+            p2.copy_(p1.data)
+
+
 @pytest.mark.parametrize("dim1", [1024], ids=id_formatter("dim1"))
 @pytest.mark.parametrize("dim2", [32, 1024, 4097], ids=id_formatter("dim2"))
 @pytest.mark.parametrize("gtype", [torch.float32, torch.float16], ids=describe_dtype)