ZeRO Optimization › The Memory Anatomy of Training

Before understanding ZeRO, you need to understand precisely what occupies GPU memory during training. The memory footprint has two main categories: model states and residual states.

Model states are the tensors directly associated with the model's learned parameters. These are the most important category because their size is fully determined by model architecture, and they grow linearly with parameter count.

The three components of model state are:

• Parameters ($\Psi$): The model weights themselves. At 16-bit (fp16 or bf16), each parameter takes 2 bytes. A 7B-parameter model needs 14 GB just for parameters.
• Gradients: One gradient tensor per parameter, matching the size of the parameters. Another 14 GB for a 7B model. These are computed during the backward pass and must persist until the optimizer step completes.
• Optimizer states: The most expensive component when using Adam. Adam maintains two additional tensors per parameter: the first moment $m$ (the exponential moving average of gradients) and the second moment $v$ (the exponential moving average of squared gradients). Both are stored in 32-bit precision to maintain numerical stability across many update steps. This is $2 \times 4 = 8$ bytes per parameter, or 56 GB for a 7B model. Mixed-precision training also requires a 32-bit master copy of the parameters, adding another 4 bytes per parameter or 28 GB.

The reason Adam's moments are stored at 32-bit rather than 16-bit precision is subtle but important. The moment estimates accumulate small incremental updates over potentially millions of training steps. At 16-bit precision, the dynamic range is insufficient to represent the ratio between large and small gradient contributions without losing the small ones to rounding. The 32-bit master copy is similarly necessary because the optimizer step computes a small parameter delta, and adding a small 16-bit delta to a 16-bit parameter can lose precision at a critical scale where training dynamics matter.

Residual states include activations, temporary buffers, and memory fragmentation. Activations are the intermediate outputs computed during the forward pass that must be retained for computing gradients during the backward pass. For a transformer with a sequence length of 2048 and a batch size of 4, the activation memory per layer can reach several gigabytes, and activations must be stored for all layers simultaneously. Gradient checkpointing (activation recomputation) trades activation memory for extra compute by discarding most activations and recomputing them during the backward pass. Since ZeRO targets model states specifically, it complements rather than replaces gradient checkpointing.

For mixed-precision training with Adam, the breakdown for a model with $\Psi$ parameters is:

• 16-bit parameters: $2\Psi$ bytes
• 16-bit gradients: $2\Psi$ bytes
• 32-bit parameter copy for optimizer update: $4\Psi$ bytes
• 32-bit Adam first moment: $4\Psi$ bytes
• 32-bit Adam second moment: $4\Psi$ bytes

This totals to $16\Psi$ bytes per model. For a 7B model, that is 112 GB. For a 70B model, 1.12 TB. No single GPU can hold this, and in naive data parallelism, every GPU in the fleet holds this full amount. A training cluster with 256 A100s running a 7B model under standard data parallelism stores $256 \times 112 \text{ GB} = 28.7 \text{ TB}$ of model state, of which $28.7 \text{ TB} - 112 \text{ GB} \approx 28.6 \text{ TB}$ is pure redundancy.