LLaMA 7B's weights take 13.5 GB, yet training it needs far more. Sixteen bytes per parameter for mixed-precision Adam, activations that grow with sequence and batch, and the trade activation checkpointing makes: a live calculator.
This deck follows chapter 10, memory, where one device’s memory is animated event by event through a training step. Measured numbers: results.md §10; the calculator below is a port of the site’s scripts/training/memory.py.
LLaMA 7B’s weights take 13.5 GB at two bytes each, yet training it on one device needs far more. With mixed-precision Adam (decks 04 and 08) every parameter costs:
That is ZeRO’s accounting (Rajbhandari et al. 2020, §3.1, \(K = 12\) bytes of optimiser state including the master copy). Plain FP32 Adam also costs 16 (4 + 4 + 8). For LLaMA 7B (6,738,415,616 parameters by the calculator; Touvron et al. 2023, Table 2, lists 6.7B) that is 107.8 GB before a single activation.
The backward pass needs, at every operation, what the forward pass computed there (deck 03). Korthikanti et al. (2022) count it for a GPT layer with 16-bit activations:
$$\text{bytes per layer} = s\,b\,h\,\Big(34 + \frac{5\,a\,s}{h}\Big)$$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 everything else, and it doubles with every doubling of the micro-batch. The whole budget per device:
$$M = 16N + L\,s\,b\,h\,(34 + 5as/h) + 6\,s\,b\,V$$LLaMA 7B’s shape (\(L\) 32, \(h\) 4,096, \(a\) 32, SwiGLU FFN 11,008, vocabulary 32,000, untied head, RMSNorm). Peak = the largest total over one step, walked event by event: gradients are written during the backward pass as activations are freed, so the peak is not simply everything at once.
Bar colours: weights, gradients, Adam state, master copy, activations, logits; dashed line = peak; solid line = one 80 GB accelerator.
Sequence 4,096, micro-batch 1, bytes on one device (results.md §10; the calculator above reproduces every peak):
| Recipe | Weights | Gradients | Optimiser | Master | Activations | Logits | Peak |
|---|---|---|---|---|---|---|---|
| mixed-precision Adam, no recomputation | 13.5 GB | 13.5 GB | 53.9 GB | 27.0 GB | 104.2 GB | 0.8 GB | 199.5 GB |
| mixed-precision Adam, selective recomputation | 13.5 GB | 13.5 GB | 53.9 GB | 27.0 GB | 18.3 GB | 0.8 GB | 116.1 GB |
| mixed-precision Adam, full recomputation | 13.5 GB | 13.5 GB | 53.9 GB | 27.0 GB | 1.1 GB | 0.8 GB | 110.7 GB |
| FP32 Adam, no recomputation | 27.0 GB | 27.0 GB | 53.9 GB | 0.0 GB | 104.2 GB | 0.8 GB | 186.3 GB |
Even with full recomputation, the 16 bytes per parameter remain and the peak is still over 100 GB: more than any single accelerator holds. Memory, not compute, is what first forces a model to be split across devices.
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 (Chen et al. 2016). It costs one extra forward pass per step: roughly a third more compute, since backward costs about twice forward.
Korthikanti et al. keep everything but the attention scores and weights (the \(5as/h\) term), which are cheap to recompute and expensive to store: \(34\,sbh\) per layer. Most of the activation saving, recomputing only the cheap part.
For LLaMA 7B the peak falls from 199.5 GB with no recomputation to 116.1 GB (selective) and 110.7 GB (full). With recomputation, the backward pass re-creates one layer’s full activations at a time on top of the checkpoints, which the calculator counts.
The formula can be checked against a real backward pass. The Decoder’s model keeps 2,320 values per sequence per block for its backward pass, counted in the reference trainer’s forward pass; Korthikanti et al.’s formula for the same shape gives 2,496 (results.md §10).
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).
Both count the same things: inputs to every linear layer, GELU’s pre-activation, attention’s weights, each norm’s input. It is the activations, not the weights, that grow with the sequence and the batch.
For the Decoder’s whole model in mixed precision: 5,376 parameters × 16 bytes = 86,016 bytes of state, and a peak of 286,208 bytes per step at 16 sequences of 8: about a quarter of a megabyte.
The next step is to stop every device holding everything:
Split optimiser states, gradients and finally weights across the data-parallel workers: 16 bytes per parameter divided by the number of workers (deck 10).
Tensor parallelism splits each layer; pipeline parallelism splits the stack; sequence parallelism splits the activations (deck 11).
The site’s notebook 10 measures peak memory (torch.cuda.max_memory_allocated) with and without checkpointing a 12-layer model on a GPU: GPU run pending; the calculator is not yet checked against it.
Animated: chapter 10, memory · next: deck 10, data parallelism and ZeRO