Vol. XIV · Parallelism Plate 1 of 11training memory, weights, gradients, optimiser states, activations, bytes per parameter, mixed precisionCE

Where the Memory Goes.

Sixteen bytes per parameter before a single activation is stored — and the activations can be larger still.

Mixed-precision training with Adam keeps, for every parameter, 2 bytes of BF16 weights, 2 of BF16 gradients, 4 of FP32 master weights and 8 for Adam’s two FP32 moments: 16 bytes in all (plate XI.6). On top come the activations saved during the forward pass for use in the backward pass — for each layer and token, about 34 d bytes plus the attention scores if they are materialised (Korthikanti and colleagues, 2023). Activations scale with the tokens in flight, not with the parameters.

Add it up for a published model and the problem of this volume appears. The workshop totals the memory to train each Llama 3 size on one device, for the chosen number of tokens per micro-batch, with or without FlashAttention’s avoidance of the score matrix: even the smallest needs over 100 GB, against 80 GB on an H100.

Below: left, the memory to train the selected model on a single device, split into weights, gradients, master weights, optimiser moments and activations. Right, how the activations grow with tokens per micro-batch.

What you will learn
Workshop
Memory to train the selected Llama 3 configuration on one device: BF16 weights and gradients, FP32 master weights, Adam’s FP32 moments, and activations for the chosen tokens per micro-batch (sequences of 8 192) from Korthikanti et al.’s estimate, with the attention-score term removed when FlashAttention is used. The dashed line is one H100’s 80 GB.

Sixteen bytes and then some

Llama 3 8B has 8.03 billion parameters, so its weights, gradients, master copy and Adam moments take 128.5 GB before a single activation is stored. With 8 192 tokens per micro-batch and FlashAttention, the activations add 36.5 GB, for 165.0 GB in all — 2.1 H100s’ worth. Without FlashAttention the attention-score term dominates: 508.6 GB. For the 405-billion-parameter model the states alone are 6.5 TB, about 88 H100s’ worth with activations.

Nothing fits on one device, and the two parts behave differently. The 16 bytes per parameter are fixed by the model and can only be divided among devices — by sharding them (plate XIV.5) or by splitting the model itself (plates XIV.6–XIV.7). Activations scale with the tokens each device processes at once and can also be traded for compute by recomputing them (plate XIV.2) or split along the sequence (plate XIV.8). Every parallel layout is a way of dividing these terms until each device’s share fits.

The optimiser’s 12 bytes per parameter are the largest fixed term, which is why optimisers with smaller state — 8-bit Adam, or factorised second moments — are attractive for memory-bound training even when they change the arithmetic slightly.

Terms

Model states
Weights, gradients, master weights and optimiser moments: 16 bytes per parameter here.
Activations
Intermediate results saved in the forward pass for the backward pass.
Micro-batch
The tokens one device processes in one forward and backward pass.
Memory budget
The accounting of everything a training step must hold at once.
Training holds 16 bytes per parameter plus activations that grow with the tokens in flight — far more than one accelerator for any large model.

Sources

  1. Rajbhandari et al., “ZeRO”, SC 2020 (16 bytes per parameter for mixed-precision Adam).
  2. Korthikanti et al., “Reducing Activation Recomputation in Large Transformer Models”, MLSys 2023 (activation memory per layer).
  3. Llama 3 configurations, arXiv:2407.21783, Table 3.