Training Memory Footprint

A 7-billion-parameter model needs about 14 GB of memory to store its weights in fp16, comfortably fitting on a single high-end GPU for inference. Training that same model can need well over 100 GB across a GPU cluster. The gap is not a rounding error, it is four separate categories of memory that inference does not need: gradients, optimiser state, activations, and temporary buffers. Understanding the breakdown explains most of the engineering decisions in large-scale training, from mixed precision to activation checkpointing to sharded optimisers.

The four categories

Weights. The parameters themselves, typically stored in fp16 or bf16 during mixed-precision training: 2 bytes per parameter.

Gradients. One gradient value per parameter, computed during the backward pass, generally kept in the same precision as the weights: another 2 bytes per parameter.

Optimiser state. Adam, the near-universal optimiser for transformer training, keeps two running statistics per parameter (first and second moment estimates), and for numerical stability these are usually kept in fp32 regardless of the training precision: 8 bytes per parameter. A full fp32 master copy of the weights is often kept as well for the same reason, adding another 4 bytes.

bytes_per_param ≈ 2 (weights, fp16)
                + 2 (gradients, fp16)
                + 4 (fp32 master weights)
                + 8 (Adam moments, fp32)
                ≈ 16 bytes per parameter

That is roughly eight times the 2 bytes needed to merely store the weights for inference, and it accounts for none of the fourth category.

Activations: the part that scales with more than parameter count

Activation memory is every intermediate tensor produced during the forward pass that must be kept around because the backward pass needs it to compute gradients (recall the chain rule requires the local derivative at each step; see the-backward-pass-gradient-flow). Unlike weights, gradients, and optimiser state, activation memory does not depend on parameter count alone, it scales with batch size, sequence length, and depth, and it can dwarf the other three categories at long context lengths, since attention's intermediate score matrix alone is O(batch * heads * seq_len^2).

Korthikanti et al., 2022 give a widely cited estimate for a standard transformer layer's activation memory:

activations_per_layer ≈ seq_len * batch * d_model * (34 + 5 * n_heads * seq_len / d_model)

The dominant behaviour to internalise, not the exact constant: activation memory grows linearly with depth (one term per layer, unless mitigated) and grows with seq_len^2 through the attention term, which is why long-context training is disproportionately memory-hungry compared to long-context inference.

Activation checkpointing: trading compute for memory

The standard mitigation is activation checkpointing (also called gradient checkpointing): store only a subset of activations (for instance, the input to each transformer block) and recompute the rest during the backward pass rather than keeping them all in memory. This turns an O(L) memory cost per layer into roughly O(sqrt(L)) at the price of a second forward pass through the recomputed portions, typically 20 to 30 percent more compute for a large reduction in peak memory (Korthikanti et al., 2022 also cover selective recomputation, checkpointing only the most memory-hungry parts of each layer rather than everything, to reduce the compute overhead further).