The memory an optimiser costs before any activation is stored

STATEMENT

Count the copies of the parameter vector each optimiser holds, price them for three model sizes, and identify the size at which the choice of optimiser stops being a convergence question and becomes a hardware one.

GIVEN

A model with $N$ parameters in fp32. Every method holds the weights and the gradient; they differ in what else persists between steps. Sizes: 124M, 1.5B and 7B.

FIND

State per parameter and total copies for SGD, SGD with momentum, RMSProp and Adam; total memory at each size; the optimiser state as a share of the total; and what swapping Adam for SGD with momentum saves at 1.5B.

STRATEGY

Count in copies of $\vec{\theta}$ rather than in bytes, because the ratio is what transfers between precisions and model sizes. Weights and gradients are common to all four, so the whole difference is the third term.

SOLUTION

Step 1 — the state each method keeps.

Adam keeps both moments of (I.7.6): $m$ for direction and $v$ for scale. That is two extra copies, and with mixed precision an fp32 master copy of fp16 weights pushes it to six.

Step 2 — priced.

Step 3 — where it stops being a preference. At 7B, Adam’s four copies come to $112$ GB. The largest single accelerator commonly available holds $80$ GB, so the model cannot be trained on one device with Adam before a single activation is stored — and activations are the larger term at any useful batch, by I.6.B05. SGD with momentum brings it to $84$ GB, which is still over. The decision has moved from “which converges better” to “which fits”, and no amount of gradient checkpointing addresses it, because checkpointing trades activation memory and this is not activation memory.

Step 4 — the share.

Half of Adam’s four-way footprint is state that exists only to choose a step size per coordinate.

Step 5 — the swap. At 1.5B:

A third of the optimiser footprint, for one moment estimate. Whether that is a good trade depends on whether the per-coordinate scaling is buying convergence on the model in question — which is an empirical matter, and is exactly the question the framing of Adam as a default discourages anyone from asking.

Check — sanity

The ratios are precision-independent. Counting in copies means fp16 halves every number in the table and changes none of the ratios. That is why the copy count is the quantity to remember and the GB figures are an illustration.

Adam is exactly twice SGD. Four copies against two, at every size. A useful sanity check on any memory estimate: if switching from SGD to Adam does not roughly double the parameter-side footprint, something else is being counted.

The 124M row is unremarkable and that is the point. $1.98$ GB against $0.99$ GB — a gigabyte, on hardware with tens. The choice genuinely does not matter at that scale, which is why the habit of reaching for Adam forms on small models and then goes unexamined on large ones.