Fix Lion (and Adagrad) decoupled weight decay in the Triton 8-bit blockwise kernel - #2
Merged
Merged
Conversation
Follow-up to bitsandbytes-foundation#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 (bitsandbytes-foundation#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 bitsandbytes-foundation#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) <noreply@anthropic.com>
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
1 participant
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.
Consolidates the Triton 8-bit blockwise Lion+Adagrad weight-decay fix into fork main.
_optimizer_update_1state_8bit_blockwise_triton_kernelgated its decoupled-decay branch onOPTIMIZER_ID == 2(ADAGRAD) instead of4(LION) -- a copy/paste on the id under a# LIONcomment. So Lion got coupled decay (corrupting its sign update) and Adagrad got decoupled decay instead of the L2 fold it expects. Changing2->4fixes both with one character.Adds
test_lion8bit_blockwise_weight_decay(validated on CPU via a coupled-bug negative control: fp32 separates ~160x, fp16 ~15x; bf16 excluded -- mantissa rounds the effect away; MPS skipped -- 8-bit blockwise op unimplemented there). Triton kernel is XPU-only. Fork-internal; not for upstream.