Memory Management › Summary

Memory management in transformer training centers on four categories of consumption, each with distinct characteristics:

• Parameters cost $2P$ bytes in FP16, setting a hard floor on GPU memory requirements. They scale with the architecture's depth and width, and they persist for the entire training run.
• Gradients cost $4P$ bytes in FP32 for numerical stability, matching the parameter shape exactly. They exist only during backpropagation and are freed after each optimizer step.
• Optimizer states cost $8P$ bytes for Adam (two FP32 moments), making the optimizer the single largest memory consumer. Alternatives like Lion (4 bytes), 8-bit Adam (2 bytes), and AdaFactor (sub-linear) reduce this significantly at varying convergence tradeoffs.
• Activations cost $O(L \times B \times T \times d)$ bytes, growing quadratically with sequence length due to attention weight matrices. Gradient checkpointing reduces this to $O(B \times T \times d)$ at the cost of approximately 33% more compute.

The combined static memory for mixed-precision Adam training follows the "14x rule": $14P$ bytes. For a 7B model, this is approximately 98 GB before activations. Memory estimation before starting a training run prevents wasted hours on predictable OOM crashes. When crashes do occur, a systematic approach of profiling checkpoints, testing at minimal batch size, and progressively applying reduction techniques provides a clear path to recovery.

When you encounter memory pressure, apply techniques in order of their cost-effectiveness: mixed precision first (free performance improvement), gradient checkpointing second (33% compute overhead), memory-efficient optimizers third (no compute overhead but may affect convergence), and distributed sharding or CPU offloading last (high complexity but unlimited scalability). Combining all three of the first techniques typically makes 7B-class models trainable on a single A100 80GB GPU, which represents the boundary between single-GPU and multi-GPU training for the current generation of hardware.

The next chapter extends these memory concepts into data parallelism, where replicating the model across GPUs multiplies parameter and optimizer state copies while enabling larger effective batch sizes.