Fix Lion to use decoupled weight decay (default + CUDA 32-bit)#1993
Merged
matthewdouglas merged 3 commits intoJul 9, 2026
Conversation
The shared 32-bit optimizer path in the `default` backend folded weight decay into the gradient (coupled L2) for Lion, which corrupts the sign update. Lion uses *decoupled* (AdamW-style) weight decay per Chen et al. 2023: shrink the param directly (p *= 1 - lr*wd), outside the sign. This is what the cpu backend and the CUDA 8-bit blockwise kernel already do; the default backend (used by MPS and any device without a dedicated kernel) was the outlier. - Drop LION (id 4) from the coupled-decay group; keep MOMENTUM/RMSPROP/ADAGRAD. - Apply decoupled decay in the LION branch, matching cpu/ops.py. - Add test_lion32bit_weight_decay: the existing test_optimizer32bit never caught this because it constructs optimizers with the default weight_decay=0, where coupled and decoupled coincide. New test fails without the fix, passes with it. NOTE: the CUDA 32-bit kernel (csrc/kernels.cu::kOptimizer32bit1State) and the Triton 1-state kernel (backends/triton/kernels_optim.py) carry the same coupled- decay bug for Lion and need a separate upstream fix (they can't be built/tested on Apple Silicon). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…sandbytes-foundation#1992) The shared 32-bit optimizer kernel (kOptimizer32bit1State) folded weight decay into the gradient (coupled L2) for every optimizer, including Lion, which corrupts Lion's sign update. Lion requires *decoupled* (AdamW-style) decay applied to the parameter directly (p *= 1 - lr*wd), outside the sign, per Chen et al. 2023. The CUDA 8-bit blockwise kernel, the cpu backend, and now the default backend already do this; the 32-bit kernel was the outlier. - Exclude LION from the coupled gradient fold. - Apply decoupled decay to the param in the LION branch, before the sign update, matching kOptimizerStatic8bit1State. - Update the NOTE in the default backend now that CUDA is fixed (only the Triton XPU 1-state kernel still carries the coupled-decay bug). test_lion32bit_weight_decay (added with the default-backend fix) is parametrized over get_available_devices(), so it exercises this kernel on CUDA hardware in CI. It cannot be built or run on Apple Silicon. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
matthewdouglas
added
ROCm
Optimizers
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Co-authored-by: Matthew Douglas <38992547+matthewdouglas@users.noreply.github.com>
matthewdouglas
approved these changes
Jul 9, 2026
matthewdouglas
left a comment
matthewdouglas
left a comment
Member
There was a problem hiding this comment.
Looks good, thanks for the PR! Appreciate that it is well scoped and has a reasonable test addition. All of the CUDA tests are passing.
If you'd like, feel free to open an equivalent PR for the Triton backend. I wouldn't expect it to be particularly invasive and probably safe to merge whether we can validate it or not.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #1992.
Problem
Lion with
weight_decay > 0applies coupled (L2) weight decay, folding decay into the gradient before thesign()in the 32-bit optimizer paths. Lion requires decoupled (AdamW-style) decay applied to the parameter directly (p *= 1 - lr*wd), outside the sign update (Chen et al. 2023).Coupling decay into the gradient changes the input to
sign(), so it corrupts the update direction, not just its magnitude.The smoking gun that this is a bug and not a design choice: the CUDA backend is already inconsistent with itself. The 8-bit blockwise kernel and the cpu backend do decoupled decay correctly; only the 32-bit paths got it wrong.
Fix
backends/default/ops.pycsrc/kernels.cu::kOptimizer32bit1StatekOptimizerStatic8bit1Statebackends/cpu/ops.pyThe CUDA change mirrors the existing 8-bit kernel exactly: exclude
LIONfrom the coupled gradient fold, and applyp *= (1 - lr*wd)in theLIONbranch before the sign update (un-scaled by themax_unormclip, matching the reference).I verified
kOptimizer32bit1Stateis the only 32-bit CUDA site that applies Lion decay: the 2-state kernel is ADAM/ADEMAMIX (already decoupled), and the 1-state precondition kernel only accumulates the update norm for Lion.Test
Adds
test_lion32bit_weight_decaytotests/test_optim.py, comparingbnb.optim.Lionagainst thelion-pytorchreference withweight_decay=0.1. It's parametrized overget_available_devices(), so it exercises the CUDA 32-bit kernel on GPU CI. The existingtest_optimizer32bitnever caught this because it uses the defaultweight_decay=0, where coupled and decoupled coincide.Under the bug the params diverge far past the error budget; with the fix they track the reference.
Notes / scope
backends/triton/kernels_optim.py) carries the same coupled-decay bug, but it's XPU-only with no CI hardware here. Left as a follow-up, happy to open a companion PR if someone can validate on XPU.