STAGE 0 · 2.5·4 · 3 NODES

Optimizers that fit

AdamW is the default optimizer, and it works — but it quietly stores two extra numbers for every weight in the model. Newer optimizers keep less, freeing GPU memory and sometimes converging faster. Here's the trade-off, one knob at a time.

THE 3 STEPS

01AdamW's tax 02Cheaper state 03The trade-off

NODE 01 / 03

AdamW's memory tax

You met the optimizer in Optimizers — the rule that nudges each weight downhill. AdamW is the default, and to do its job it keeps two running numbers per weight.

Momentum (m) — a running average of the gradient, so updates keep their direction instead of jittering.
Variance (v) — a running average of the gradient's size, used to give each weight its own step size.
Those two numbers are the green rows — 8 of the 16 bytes per weight. The optimizer's state is half your memory — as big as the fp32 master copy, the weights, and the gradients combined.
Scale it up: a 70B model spends about 560 GB just on m and v — before a single activation.

MEMORY PER WEIGHT · ADAMW

Drag the model size — the two green rows are the optimizer's own, and they grow right along with it.

MODEL SIZE8B

Momentum + variance are stored per weight, so the optimizer's bill scales one-to-one with model size — and it's the biggest single slice.

MEMORY_BUDGET
TOTAL128 GB
OPT STATE64 GB
SHARE50%
Training memory for weights, gradients and optimizer state at this model size. Optimizer state is momentum + variance.

NODE 02 / 03

Cheaper state

The newer optimizers ask a simple question: do we really need two full numbers per weight? Each keeps less — and the memory you save is real VRAM back.

Lion keeps one number — momentum only. It updates using just the sign of it, so the variance term (and its 4 bytes) disappears.
Adafactor keeps almost none — instead of storing the full variance grid for a weight matrix, it stores one number per row and one per column and reconstructs the rest.
Muon keeps momentum only, but adds a cheap orthogonalize step that rescales the whole 2D weight-matrix update at once — the trick behind its speed.

OPTIMIZER STATE · BYTES PER WEIGHT

Tap an optimizer — its bar is the state it keeps per weight; AdamW stays as the reference.

Lion — momentum only, and it uses just its sign. Half the state, ~4 bytes per weight.

STATE_COST · Lion
PER WEIGHT4 B/wt
vs ADAMW−50% saved
AT 70B280 GB
Optimizer-state bytes per weight, how much that saves against AdamW, and the total optimizer memory it would cost on a 70B model.

NODE 03 / 03

The trade-off

Saving memory isn't free. The variance term AdamW pays for is also what makes it stable — so cheaper state usually means more careful tuning, and sometimes different convergence.

Muon can converge faster — its orthogonalize step makes better use of each big-matrix update, reaching a target loss in fewer steps.
Lion and Adafactor need care — with the variance gone or approximated, they want a smaller learning rate and warm-up, or training can wobble.
AdamW is the safe default — you compare everything to it, and reach for the leaner ones when memory is the binding constraint.

STEPS TO REACH A TARGET LOSS

Tap an optimizer — green is it, grey is AdamW. The dotted line is the target loss; the crossings are the step counts.

Muon — reaches the target in noticeably fewer steps on the big matrices, which is why speed-runs use it.

TRADE_OFF · Muon
STEPS68k
vs ADAMW−31k
STABILITYsolid
Illustrative steps to reach the target loss, the change against AdamW, and how touchy the optimizer is to tune.

EXPLAIN IT BACK

Adafactor saves almost all of AdamW's optimizer memory. So why is AdamW still the default for most training runs?

NEXT: PARAMETER-EFFICIENT FINE-TUNING

You've made the optimizer cheaper. The other way to cut training memory is to stop training most of the weights at all — freeze the model and learn a tiny add-on. Next: parameter-efficient fine-tuning (LoRA).

← Back to distributed training Continue to fine-tuning →

Language: English