Fix decoupled weight decay ordering in the CUDA Adam/AdEMAMix kernels

[yentur]

Closes #2010.

The CUDA 2-state optimizer kernels apply weight decay after adding the update:

p = (p + step_size*update) * (1 - lr*wd)

which expands to p*(1 - lr*wd) + step_size*update*(1 - lr*wd). The extra factor on the update term is not in AdamW Algorithm 2 (Loshchilov and Hutter, arXiv:1711.05101) or in the AdEMAMix update rule (Pagliardini et al., arXiv:2409.03137). Both take the decay and the gradient-based update from the same previous parameter, which rearranges to p*(1 - lr*wd) + step_size*update. torch.optim.AdamW does the same, and so do the cpu, default and triton backends here.

Three sites in csrc/kernels.cu: the ADAM and ADEMAMIX branches of kOptimizer32bit2State, and the 2-state branch of kOptimizerStatic8bit2StateBlockwise. The change is a reorder, 9 insertions and 8 deletions. The 1-state kernels already use coupled L2 for momentum/rmsprop/adagrad and decoupled decay for Lion, matching the Python backends, so they are untouched.

@ErenAta16 reported this and worked out in the thread that CUDA is the side that diverges. @egeozkoc confirmed it against the paper, pointed out that the 8-bit blockwise kernel is affected too, and offered to write the expanded weight-decay tests. I went ahead because the issue had been sitting for about a month. Happy to close this if either of you would rather carry it.

This changes results for existing users

CUDA runs using Adam, AdamW, LAMB or AdEMAMix with weight_decay > 0 will follow different training trajectories after this. The per-step difference is lr*wd*|update|. Runs with weight_decay = 0 are unaffected, as are the cpu, default, triton and XPU backends, which already used this ordering.

Why the current tests do not catch it

They run these optimizers at the default weight_decay of 0, where the two orderings coincide. At the lr=1e-3, wd=0.01 an Adam user would typically pick, 20 steps leave the orderings 1.7e-7 apart, inside the 1e-6 tolerance. The entire existing test_optim.py passes on the unfixed kernel:

$ BNB_TEST_DEVICE=cuda pytest tests/test_optim.py -q -k "not weight_decay"   # unfixed kernel
213 passed, 12 skipped, 35 deselected, 194 warnings in 61.97s

The two orderings separate as lr*lr*wd, while the kernels' own arithmetic error grows only as lr. So test_optimizer32bit_weight_decay runs at lr=0.1, wd=0.1, where the orderings are 1.0e-3 apart and agreement with the reference is 3.9e-5.

test_optimizer8bit_weight_decay cannot use a full run at all: the state quantization error is 6.6e-4 per step while the orderings differ by 1.6e-6, so the signal is buried. It instead takes one step from a fresh optimizer, where the state is still zero and therefore quantizes exactly, and the update reduces to -lr*g/(|g| + eps) (AdEMAMix also mixes in alpha*(1-b3)*g). At those hyperparameters the orderings are 1.0e-3 apart, and the fixed kernels sit 3.0e-6 from the closed form, which is __powf and __fdividef rather than the decay. Both tests are fp32; the 16-bit paths in test_optimizer32bit resynchronise the parameters every step, so at most one step of divergence is visible there and it stays inside the looser 16-bit tolerances.

Verification

RTX 3090, sm_86, CUDA 12.4, torch 2.5.1+cu124, built from source with -DCOMPUTE_CAPABILITY=86. Only csrc/kernels.cu was reverted between the two runs; the tests were identical.

$ BNB_TEST_DEVICE=cuda pytest tests/test_optim.py -q -k weight_decay     # kernels.cu reverted
E  Greatest absolute difference: 0.0010029971599578857 at index (7, 2) (up to 0.0001 allowed)
E  Greatest absolute difference: 0.0010035037994384766 at index (329, 20) (up to 1e-05 allowed)
8 failed, 252 deselected in 16.44s

$ BNB_TEST_DEVICE=cuda pytest tests/test_optim.py -q -k weight_decay     # fix applied
8 passed, 252 deselected in 18.81s

$ BNB_TEST_DEVICE=cuda pytest tests/test_optim.py -q                     # fix applied, full module
221 passed, 12 skipped, 27 deselected, 194 warnings in 83.19s

Isolating the ordering on one step, same hardware, with the fix applied:

AdamW-32bit: |p-decay_first|=3.040e-06   |p-decay_after|=9.971e-04
AdamW-8bit:  |p-decay_first|=3.040e-06   |p-decay_after|=9.971e-04

On cpu and mps, which already had the correct ordering, the new tests pass unchanged (8 passed on cpu, 4 passed, 4 skipped on mps; the 8-bit ones skip because optimizer_update_8bit_blockwise has no MPS kernel). To check they are not vacuous without a GPU, I also re-registered the cpu kernels with the CUDA ordering and confirmed all 8 fail.

pre-commit run --all-files passes.

Not covered

bnb.optim.Adam applies decoupled decay in every backend, so with weight_decay > 0 it matches torch.optim.AdamW rather than torch.optim.Adam (on cpu, 60 steps at lr=1e-2, wd=0.1: max 5.3e-1 from torch.optim.Adam, 1.5e-7 from torch.optim.AdamW). That is the separate point you raised in the issue and it is not touched here. The skip_zeros path still skips the decay along with the update when a gradient is exactly zero, which is existing behaviour I left alone. I verified on sm_86 only.