The Memory Arithmetic of Training

Sixteen bytes per trainable parameter, the activation bill, and why checkpointing every √L layers is exactly the right thing to do.

Fine-TuningAdvanced⏱ 1 hour9 min readmemoryadammixed precisiongradient checkpointingzero

3 problems →

AssumesLow-Rank Adaptation · Adaptive Methods and Variance Reduction

Persistent state

PropositionSixteen bytes per trainable parameter

Mixed-precision training with Adam stores, per trainable parameter: BF16 weights (2 bytes), a BF16 gradient (2), an FP32 master weight (4), and FP32 first and second moment estimates (4 each) —

16 \text{ bytes per trainable parameter}.

Assumptions

Adam or AdamW with the standard two-moment state; BF16 or FP16 compute with an FP32 master copy, the usual mixed-precision recipe. Variants change the constant and not the structure: plain SGD with momentum needs 4 bytes of state instead of 8, 8-bit optimizers quantize the moments to 1 byte each, and pure FP32 training drops the master copy while doubling the others.

CorollaryWhere LoRA's memory saving actually comes from

Freezing a base model of $P$ parameters and training an adapter of $P_a \ll P$ reduces persistent state from $16P$ to $2P + 16P_a$ bytes — a saving of $14P - 16P_a \approx 14P$ — and leaves activation memory unchanged.

Activations

PropositionActivation memory scales with batch, sequence and depth

For a transformer with $L$ layers, width $d$ , batch $B$ and sequence length $S$ , the activations retained for backpropagation occupy

M_{\text{act}} = \Theta\big(L\,B\,S\,d\big) \text{ bytes},

with a constant depending on how many intermediate tensors per layer are kept, and an additional $\Theta(L\,B\,S^2\,n_{\text{heads}})$ term if attention probability matrices are materialised.

Note that $P$ does not appear: activation memory is independent of parameter count and depends instead on the shape of the input. This is why the same model trains fine at $S=512$ and fails at $S = 8192$ .

Checkpointing

Gradient checkpointing stores activations only at selected layer boundaries and recomputes the rest during the backward pass.

TheoremOptimal uniform checkpointing uses √L segments

Partition $L$ layers into consecutive segments of $k$ layers. Store the activation at each segment boundary, and during backpropagation recompute a segment's interior from its stored boundary. Peak activation storage, in units of one layer's activations, is

m(k) = \frac{L}{k} + k ,

minimised over real $k > 0$ at $k^\star = \sqrt L$ with $m(k^\star) = 2\sqrt L$ , at the cost of one additional forward pass.

Key idea

The trade is $\Theta(L)$ memory for $\Theta(\sqrt L)$ memory at the price of roughly $33\%$ more compute — one extra forward pass against the usual forward-plus-backward cost of about three forward passes. For $L = 32$ this is $2\sqrt{32} \approx 11$ layers' worth of activations instead of 32, a $2.8\times$ reduction. For $L = 80$ it is $8.9\times$ against $80$ — the deeper the model, the better the deal.

What this does not say

Every count here is of tensors the algorithm requires. Real allocators add fragmentation, workspace for fused kernels, communication buffers, and framework overhead that together commonly add tens of percent. Use these formulas to decide whether a configuration is plausible, then measure.

References

PaperP. Micikevicius et al. (2018). Mixed Precision Training
PaperT. Chen, B. Xu, C. Zhang, C. Guestrin (2016). Training Deep Nets with Sublinear Memory Cost
PaperS. Rajbhandari, J. Rasley, O. Ruwase, Y. He (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models

These works are cited as sources to read, not reproduced. All exposition, examples and problems on this site are written independently; where an argument is standard, the citation points to a canonical treatment.

Results stated on this page

PropositionSixteen bytes per trainable parameterproof includedMixed-precision Adam requires 16 bytes of persistent state per trainable parameter, of which only 2 are the weights used in the forward pass.
PropositionActivation memory scales with batch, sequence and depthproof includedStoring activations for backpropagation costs Θ(L·B·S·d) bytes, independent of parameter count.
TheoremOptimal uniform checkpointing uses √L segmentsproof includedStoring every k-th layer's activations minimises peak memory at k = √L, giving Θ(√L) storage for one extra forward pass.
CorollaryWhere LoRA's memory saving actually comes fromproof includedFreezing the base model removes the 14 bytes of optimizer and gradient state per parameter, but leaves activation memory untouched.

Problems (3)

All problems