Optimisers, and what their state costs

Why not plain gradient descent

The simplest update is w -= lr * grad. It works, badly, on deep networks: different parameters have gradients of wildly different magnitudes, so any single learning rate is too large for some and too small for others.

Adam fixes this per parameter, by keeping two running averages:

• m, an exponential moving average of the gradient itself: momentum. It smooths the noise from batch to batch.
• v, an exponential moving average of the gradient squared: a running estimate of that parameter’s typical magnitude.

The update divides by the square root of v, so a parameter whose gradients are consistently tiny still takes a meaningful step, and one whose gradients are consistently huge takes a small one. This is what makes transformer training work at all with one global learning rate.

m = β₁·m + (1-β₁)·g
v = β₂·v + (1-β₂)·g²
w = w - lr · m̂ / (√v̂ + ε)

The defaults β₁ = 0.9 and β₂ = 0.999 mean m is an average over roughly the last ten steps and v over roughly the last thousand. You will essentially never change these; the one exception is that very short runs sometimes benefit from β₂ = 0.95, because 0.999 takes longer to warm up than the whole run lasts.

The memory bill

Two averages, one per parameter, each stored in fp32: 8 bytes per trainable parameter.

For a 7B model trained fully, that is 56 GB of optimiser state alone, on top of 28 GB of fp32 weights and 28 GB of gradients. This is the single biggest reason full fine-tuning of a 7B model does not fit on an 80 GB card without sharding, and it is the whole motivation for the parameter-efficient methods in the next module.

For a LoRA with 20M trainable parameters, the same formula gives 160 MB. The optimiser stops being a consideration entirely.

AdamW, and why the W

Weight decay pulls parameters towards zero. In the original Adam it was implemented by adding wd · w to the gradient, which then went through the same division by √v, so the amount of decay a parameter received depended on its gradient history. AdamW decouples it: decay is applied directly to the weights, outside the adaptive step.

w = w - lr · m̂ / (√v̂ + ε) - lr · wd · w

This is not a subtlety; it changes results measurably, and AdamW is the default everywhere now. A typical value is 0.01, and for LoRA it is often set to 0 because the trained matrices are small and initialised so that the adapter starts as a no-op anyway.

The cheaper optimisers, and when they are free wins

• 8-bit AdamW (bitsandbytes, optim="adamw_bnb_8bit"). Quantises the two states to 8 bits with block-wise scaling. 8 bytes per parameter becomes 2. On full fine-tunes the saving is large and the quality difference is, in most reported comparisons, within noise. Cheap to try, easy to check.
• Paged AdamW (adamw_bnb_8bit with paging, or paged_adamw_32bit). Uses CUDA unified memory so optimiser state can spill to CPU RAM during a spike instead of raising an out-of-memory error. This is what makes long-sequence QLoRA runs survive their worst batch.
• Adafactor. Factorises v into row and column statistics instead of storing it per parameter, so state falls to roughly 4 bytes per parameter or less. Was standard in the T5 era; still useful when memory is desperate.
• SGD with momentum. One state tensor instead of two, 4 bytes per parameter. It can train transformers, but it needs more careful tuning and more steps, so it is a memory-of-last-resort choice rather than a default.

The practical rule

If you are doing LoRA or QLoRA, use AdamW and never think about this again: 160 MB of state is not your problem. If you are doing full fine-tuning, the optimiser is the first place to look for memory, and 8-bit AdamW is the cheapest 30 GB you will ever recover on a 7B model.

One honest caveat: the quality claims for 8-bit and factorised optimisers come from benchmarks on pretraining and on standard fine-tuning sets. They are well supported but they are not a guarantee about your task. If you switch optimisers to fit in memory, run your evaluation against a smaller run with full-precision AdamW once, so you know what the swap cost you rather than assuming it cost nothing.