From 7333266d5f2222689066fcac5d16feb9994530a0 Mon Sep 17 00:00:00 2001 From: Eric Eaglstun Date: Tue, 14 Jul 2026 17:30:16 -0600 Subject: [PATCH] Fix Lion (and Adagrad) weight decay in the Triton 8-bit blockwise kernel Follow-up to #1993 and the Triton 32-bit fix. _optimizer_update_1state_8bit_ blockwise_triton_kernel gated its decoupled-decay branch on OPTIMIZER_ID == 2, which is ADAGRAD, not LION (id 4) -- a copy/paste error on the id, right under a `# LION` comment. The consequences: - Lion (id 4) fell through to the coupled (L2) `elif` and folded weight decay into the gradient before sign(), corrupting the update direction. Lion needs *decoupled* decay applied to the param directly (Chen et al. 2023). - Adagrad (id 2) wrongly got Lion's decoupled decay instead of the coupled L2 fold it expects (and that the cpu backend applies for it). Correcting the id to 4 fixes both optimizers with a one-character change: Lion gets decoupled decay, Adagrad falls through to the coupled fold. Adds test_lion8bit_blockwise_weight_decay, the 8-bit companion to test_lion32bit_weight_decay (#1993). It runs Lion8bit against the lion-pytorch reference with weight_decay=0.1 and no per-step resync, so the coupled-vs- decoupled divergence accumulates into a reliable signal. Validated by injecting the coupled bug into the CPU 8-bit kernel: the test fails under the bug and passes with decoupled decay (fp32 separates ~160x, fp16 ~15x). bf16 is excluded -- its mantissa rounds the bug's per-element update difference away, so it cannot resolve coupled vs decoupled regardless of correctness. MPS is skipped (the 8-bit blockwise op is unimplemented there). Like #1993's CUDA change and the Triton 32-bit fix, this 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 | 7 ++- tests/test_optim.py | 46 +++++++++++++++++++ 2 files changed, 51 insertions(+), 2 deletions(-) diff --git a/bitsandbytes/backends/triton/kernels_optim.py b/bitsandbytes/backends/triton/kernels_optim.py index d7abc4af1..8a61ba9d0 100644 --- a/bitsandbytes/backends/triton/kernels_optim.py +++ b/bitsandbytes/backends/triton/kernels_optim.py @@ -907,8 +907,11 @@ def _optimizer_update_1state_8bit_blockwise_triton_kernel( s1 = dequant_8bit_blockwise_kernel_util(state1_ptr, offsets, qmap1_ptr, absmax1_ptr, mask, BLOCK_SIZE_N) # 3. Optimizer-specific updates - # LION - if weight_decay > 0.0 and OPTIMIZER_ID == 2: + # LION (id 4) uses decoupled weight decay: shrink the param directly, outside the + # sign update (Chen et al. 2023). This was previously gated on OPTIMIZER_ID == 2, + # which is ADAGRAD -- so Lion got coupled decay (corrupting its sign update) and + # Adagrad got decoupled decay instead of the L2 fold it expects. + if weight_decay > 0.0 and OPTIMIZER_ID == 4: p *= 1.0 - lr * weight_decay # Apply weight decay for momentum, rmsprop, adagrad elif weight_decay > 0.0: diff --git a/tests/test_optim.py b/tests/test_optim.py index d9c24bdb5..94b0f27ee 100644 --- a/tests/test_optim.py +++ b/tests/test_optim.py @@ -292,6 +292,52 @@ def test_lion32bit_weight_decay(dim1, dim2, gtype, device): p2.copy_(p1.data) +# bf16 is excluded: its ~8-bit mantissa rounds the bug's per-element (2 * lr) update +# difference away on param write-back, so coupled and decoupled Lion8bit are numerically +# indistinguishable in bf16 regardless of correctness. fp32/fp16 resolve it cleanly. +@pytest.mark.parametrize("gtype", [torch.float32, torch.float16], 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_lion8bit_blockwise_weight_decay(dim1, dim2, gtype, device): + """Lion8bit must also use *decoupled* weight decay, like the 32-bit path. + + Companion to test_lion32bit_weight_decay, covering the 8-bit blockwise kernels. The + Triton 1-state blockwise kernel gated its decoupled-decay branch on OPTIMIZER_ID == 2 + (ADAGRAD) rather than 4 (LION), so Lion fell through to the coupled (L2) fold and its + sign update was computed from a decay-polluted gradient. + + Params are *not* resynced between steps: the coupled-vs-decoupled difference is a + small per-step signal that only becomes reliably measurable once it accumulates. The + budgets sit well above the correct path's 8-bit quantization noise and far below the + error the coupled-decay bug produces (measured ~1.3e-3 in fp32, ~1.5e-3 in fp16). + """ + if device == "mps": + # The 8-bit blockwise optimizer op is not implemented for MPS (the existing + # test_optimizer8bit hits the same gap); there is nothing to exercise here. + pytest.skip("optimizer_update_8bit_blockwise is not implemented for the MPS device") + + weight_decay = 0.1 + err_budget = 1e-4 if gtype == torch.float32 else 4e-4 + + 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.Lion8bit([p2], weight_decay=weight_decay) + + 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() + + torch_optimizer.step() + bnb_optimizer.step() + + assert (p1 - p2.float()).abs().mean().item() < err_budget + + @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)