Memory

Sixteen bytes per parameter. With mixed-precision Adam (chapters 5 and 9), every parameter costs 2 bytes of BF16 weight, 2 of BF16 gradient, 4 + 4 for Adam's first and second moments and 4 for the FP32 master copy: 16 bytes, the accounting of ZeRO (Rajbhandari et al. 2020, §3.1, $K = 12$ bytes of optimiser state). Plain FP32 Adam also costs 16 (4 + 4 + 8). For LLaMA 7B (6,738,415,616 parameters by this calculator; Touvron et al. 2023, Table 2, lists 6.7B) that is 107.8 GB before a single activation.

Activations. The backward pass needs, at every operation, what the forward pass computed there (chapter 4: "which activations must be saved, and why"). Korthikanti et al. (2022) count them for a GPT layer with 16-bit activations: $s\,b\,h\,(34 + 5as/h)$ bytes for sequence length $s$, micro-batch $b$, width $h$ and $a$ heads. The $5as/h$ term is attention's $s \times s$ score and weight matrices, one per head (Decoder, chapter 3): it grows with the square of the sequence. For LLaMA 7B at sequence 4,096 and one sequence per step that is 104.2 GB over 32 layers, about as much again as the weights and optimiser together, and it doubles with every doubling of the micro-batch. The Decoder's model keeps 2,320 values per sequence per block for its backward pass (counted in the reference trainer), against 2,496 by the formula for the same shape; the formula assumes a feed-forward layer four times the model's width (the Decoder's is twice) and dropout masks (the Decoder's has none), so the two counts differ in detail but agree in size. Either way, it is the activations, not the weights, that grow with the sequence and the batch.

Activation checkpointing (Chen et al. 2016) trades compute for that memory: keep only each layer's input (2 $sbh$ bytes) and, during the backward pass, recompute the layer's forward pass to get the rest back. It costs one extra forward pass per step: roughly a third more compute, since a backward pass costs roughly twice a forward one. Korthikanti et al.'s selective recomputation keeps everything but the attention scores and weights (the $5as/h$ term), which are cheap to recompute and expensive to store: $34\,sbh$. Pick "full" or "selective" above and watch the activations shrink and a recomputed layer appear in each backward step. For LLaMA 7B the peak falls from 199.5 GB with no recomputation to 116.1 GB (selective) and 110.7 GB (full).

Still too much for one device. Even with full recomputation, the 16 bytes per parameter remain. The next step is to stop every device holding everything: ZeRO and FSDP shard the optimiser states, gradients and weights across the data-parallel workers (chapter 11), and tensor and pipeline parallelism split the model itself (chapter 12).

Run it yourself on a GPU. Notebook 10 (open in Colab) measures peak memory (torch.cuda.max_memory_allocated) and step time with and without checkpointing a 12-layer model at three sequence lengths (about 10 minutes on any GPU). The notebook is one of the planned GPU replications listed under Runs to come; the calculator above comes from the in-browser model.