diff --git a/bitsandbytes/backends/default/ops.py b/bitsandbytes/backends/default/ops.py index 33bdd7ed3..521802922 100644 --- a/bitsandbytes/backends/default/ops.py +++ b/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 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])))); diff --git a/tests/test_optim.py b/tests/test_optim.py index 0a4b3d6af..d9c24bdb5 100644 --- a/tests/test_optim.py +++ b/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)