Fix decoupled weight decay ordering in the CUDA Adam/AdEMAMix kernels - #2040
Fix decoupled weight decay ordering in the CUDA Adam/AdEMAMix kernels#2040yentur wants to merge 1 commit into
Conversation
The CUDA kernels applied weight decay after adding the optimizer update, computing p = (p + step_size*update) * (1 - lr*wd). Expanded, that scales the update term by an extra (1 - lr*wd) that the reference does not have. AdamW Algorithm 2 in Loshchilov & Hutter (arXiv:1711.05101) and the AdEMAMix update rule in Pagliardini et al. (arXiv:2409.03137) both take the decay and the gradient-based update from the same previous parameter, which rearranges to p = p*(1 - lr*wd) + step_size*update. torch.optim.AdamW does the same, and so do the cpu, default and triton backends. Only the CUDA kernels differed. Three sites: the ADAM and ADEMAMIX branches of kOptimizer32bit2State and the 2-state branch of kOptimizerStatic8bit2StateBlockwise. The 1-state kernels already match, and Lion is unaffected. This changes training trajectories for CUDA runs that use Adam, AdamW, LAMB or AdEMAMix with weight_decay > 0. The per-step difference is lr*wd*|update|. Runs with weight_decay = 0 are unaffected. The existing tests could not see this. They run these optimizers at the default weight_decay of 0, and at lr=1e-3, wd=0.01 the two orderings stay 1.7e-7 apart, inside the 1e-6 tolerance. On an RTX 3090 the whole of test_optim.py passes on the unfixed kernel. The separation grows with lr*lr*wd while the kernels' own arithmetic error only grows with lr, so the new 32-bit test 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. The 8-bit test cannot use a full run, because the state quantization error is 6.6e-4 per step against a 1.6e-6 ordering difference, so it checks one step from a fresh optimizer, where the state is still zero and quantizes exactly and the update has a closed form.
ErenAta16
left a comment
ErenAta16
left a comment
There was a problem hiding this comment.
This fixes the right side, and I want to be explicit about that because #2010 pointed at the other one.
Which ordering is correct
I settled this against a source neither implementation consults, torch.optim itself. In torch 2.11.0, AdamW is Adam(..., decoupled_weight_decay=True), and _single_tensor_adam does:
if weight_decay != 0:
if decoupled_weight_decay:
# Perform stepweight decay
param.mul_(1 - lr * weight_decay)That runs before the moment updates and before the final addcdiv_, so the parameter is shrunk first and the gradient-based update is added afterwards, unscaled:
p_new = p_old * (1 - lr*wd) + step_size * update
which is Algorithm 2 of Loshchilov & Hutter and exactly what this PR moves the CUDA kernels to. The old CUDA ordering multiplied the whole thing, update included, by (1 - lr*wd), so the effective step size shrank with the decay rate.
Correcting my own issue
#2010 framed the CUDA kernel as the reference and the default/triton backends as the deviation. The divergence I reported was real and the arithmetic in that issue is right, but the conclusion about which side to change was mine to get right and I got it backwards. default and triton already matched torch; CUDA was the outlier. Fixing the kernels, as you have done here, is the correct direction, and the Python backends should be left alone.
The fix is complete, which I checked rather than assumed
kernels.cu applies weight_decay in four places. This PR changes two of them. The other two should stay exactly as they are, so nobody needs to wonder whether they were missed:
kOptimizer32bit2State 693, 706 changed here
kOptimizer32bit1State 866, 885 correct as is
kOptimizerStatic8bit2StateBlockwise 1094 changed here
kOptimizerStatic8bit1StateBlockwise 1227 correct as is
The 1-state kernels split by optimizer rather than applying one rule:
if (weight_decay > 0.0f && OPTIMIZER != LION)
g_vals[j] = (float)g_vals[j] + (((float)p_vals[j]) * weight_decay);MOMENTUM, RMSPROP and ADAGRAD fold the decay into the gradient, which is coupled L2 decay and the correct form for those optimizers; there is no ordering question there because the decay never touches the parameter directly. LION is excluded from that branch and applies decoupled decay to the parameter before its sign update, which is already the ordering this PR is establishing for the 2-state kernels. The 8-bit 1-state kernel at 1227 makes the same split with the same result.
So the two-of-four count is the right count, not a partial fix.
Why the existing suite could not have caught it
Your docstring argues the separation grows with lr*lr*wd while the kernel's own arithmetic error grows with lr. That reproduces. Running both orderings side by side, 50 steps over 4096 parameters, and taking the max absolute divergence:
lr=0.001 wd=0.01 -> 2.8e-07 inside a 1e-6 tolerance
lr=0.01 wd=0.01 -> 2.8e-05 outside
lr=0.1 wd=0.1 -> 2.1e-02 outside
My absolute figures differ from the ones in your docstring because the harness differs, step count and gradient scale included, so treat this as corroborating the shape rather than reproducing the exact number. The conclusion holds either way: at the hyperparameters an Adam user actually picks, the two orderings sit inside the tolerance test_optimizer32bit uses, and that suite runs these optimizers at weight_decay=0 anyway. Choosing lr=0.1, wd=0.1 for the new test is what makes it discriminating rather than a second copy of the existing one.
Citing the papers next to the reference implementations in the test file is worth keeping. The next person to look at this will have the same "which side is right" question I had, and the answer is now in the tree instead of in a review thread.
Nothing blocking from me. Thanks for picking this up, and for taking it in the direction the reference actually supports rather than the one the issue suggested.
ErenAta16
left a comment
ErenAta16
left a comment
There was a problem hiding this comment.
I filed #2010, so this is the fix for my report. Please carry it. You have a CUDA card and I do not, and the GPU verification here is the half I could not have done.
I checked the ordering independently rather than re-read the diff, using torch.optim.AdamW as the referee since neither ordering consults it. Float64, lr=1e-3, wd=0.01, 20 steps, same gradients into both:
decay BEFORE update (this PR) max|diff to torch.optim.AdamW| = 0.000e+00
decay AFTER update (current kernel) max|diff to torch.optim.AdamW| = 1.126e-07
Exactly zero, not merely closer. That is the part I would put in the PR description: the two orderings are not two approximations of AdamW where one is more accurate, one of them is AdamW and the other is a different algorithm. The extra (1 - lr*wd) on the gradient term is a factor Algorithm 2 does not contain, and no tolerance argument makes it correct.
The one-step algebra lands where you say:
decay_after = (p - ss*u)*(1 - lr*wd)
decay_first = p*(1 - lr*wd) - ss*u
gap = ss*u*lr*wd predicted 1.500e-06, measured 1.500e-06
Your scaling claim is the load-bearing one and it holds. This is what explains why the suite never caught it, so it is worth having measured rather than argued:
wd=0.01, 20 steps
lr=1e-4 gap=1.126e-09
lr=1e-3 gap=1.126e-07
lr=1e-2 gap=1.125e-05
lr=1e-1 gap=1.113e-03
A hundredfold per tenfold in lr, so the gap really does grow as lr^2 * wd while the kernels' own error grows as lr. That is exactly why test_optimizer32bit_weight_decay has to sit at lr=0.1 and why the 8-bit test has to take a single step from zeroed state: at the hyperparameters a user actually picks, the defect is smaller than the arithmetic noise, and any test written at those values would be measuring nothing.
My number is 1.126e-07 where you quote 1.7e-07 for the same configuration. Different bias-correction placement in my scratch reimplementation, most likely; same order and same conclusion, so I mention it only so nobody reads the mismatch later as a disagreement.
Two things I would keep exactly as they are:
Re-registering the cpu kernels with the CUDA ordering to confirm the new tests fail is the step most people skip. Without it "8 passed on cpu" only tells you the tests run, not that they discriminate.
Scoping out bnb.optim.Adam matching torch.optim.AdamW rather than torch.optim.Adam is right. That is a naming and contract question and folding it into a kernel reorder would make both harder to review.
One small thing, non-blocking. The decay now sits inside the if (!skip_zeros || ...) guard in the ADAM branch, which it also did before, so skip_zeros=True still skips the decay for an exactly-zero gradient. Unchanged behaviour and correctly left alone, but now that the decay is the first statement rather than the last it reads more like a deliberate coupling than it did. One line saying it is inherited, not chosen, would save the next reader the archaeology.
@egeozkoc since you confirmed the 8-bit blockwise site, the reorder there covers both the ADEMAMIX and ADAM branches because it moved above the if (OPTIMIZER == ADEMAMIX).
matthewdouglas
added
CUDA
|
Hi @yentur, thanks for the PR. I wanted to explain why the issue remained open for a while. In the issue #2010, I had stated:
So, I had not yet committed to this change, just indicated where I was leaning. I wasn't prepared to introduce a change like this into our release cycle at the time. Then I went on vacation. I will however consider this PR for v0.51.0. |
3 participants
Closes #2010.
The CUDA 2-state optimizer kernels apply weight decay after adding the update:
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 top*(1 - lr*wd) + step_size*update.torch.optim.AdamWdoes the same, and so do the cpu, default and triton backends here.Three sites in
csrc/kernels.cu: theADAMandADEMAMIXbranches ofkOptimizer32bit2State, and the 2-state branch ofkOptimizerStatic8bit2StateBlockwise. 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 > 0will follow different training trajectories after this. The per-step difference islr*wd*|update|. Runs withweight_decay = 0are 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_decayof 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 existingtest_optim.pypasses on the unfixed kernel:The two orderings separate as
lr*lr*wd, while the kernels' own arithmetic error grows only aslr. Sotest_optimizer32bit_weight_decayruns 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_decaycannot 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 inalpha*(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__powfand__fdividefrather than the decay. Both tests are fp32; the 16-bit paths intest_optimizer32bitresynchronise 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. Onlycsrc/kernels.cuwas reverted between the two runs; the tests were identical.Isolating the ordering on one step, same hardware, with the fix applied:
On cpu and mps, which already had the correct ordering, the new tests pass unchanged (
8 passedon cpu,4 passed, 4 skippedon mps; the 8-bit ones skip becauseoptimizer_update_8bit_blockwisehas 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-filespasses.Not covered
bnb.optim.Adamapplies decoupled decay in every backend, so withweight_decay > 0it matchestorch.optim.AdamWrather thantorch.optim.Adam(on cpu, 60 steps at lr=1e-2, wd=0.1: max 5.3e-1 fromtorch.optim.Adam, 1.5e-7 fromtorch.optim.AdamW). That is the separate point you raised in the issue and it is not touched here. Theskip_zerospath 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.