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).

When it falls down

• The 16-bytes-per-parameter rule assumes Adam and mixed precision. Different optimisers (SGD with momentum, Adafactor) and different precision schemes change the multiplier substantially; Adafactor, for instance, was designed specifically to cut optimiser-state memory by factorising the second moment.
• Sharding changes where the bytes live, not how many there are. Techniques like ZeRO / FSDP partition optimiser state, gradients, and even weights across devices so no single device holds the full 16 bytes per parameter, but the aggregate memory across the cluster is roughly unchanged; the win is distribution, not reduction.
• Peak memory is not average memory. Training memory usage spikes at specific points (right before an optimiser step, or during the backward pass through the largest layer), so profiling average usage badly underestimates the peak that actually causes out-of-memory errors.
• Recomputation is not free of correctness subtleties. Non-deterministic operations (some dropout implementations, certain fused kernels) recomputed during the backward pass must reproduce bit-identical results to the forward pass or the gradient computed is technically wrong, which constrains which operations checkpointing frameworks can safely wrap.

Further reading

• Korthikanti et al., 2022, Reducing Activation Recomputation in Large Transformer Models, arXiv:2205.05198 - the activation memory formula and selective recomputation strategy used above.
• Rajbhandari et al., 2019, ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, arXiv:1910.02054 - the sharded optimiser-state approach referenced in the limitations section.