ARCHITECTURE & FORMULATIONModule 4: Distributed Infrastructure & 3D Parallelism
The Training Memory Bill: Weights, Gradients, Adam, & Activations
A common misconception is that GPU memory only stores model weights. In FP16/BF16 mixed-precision AdamW training, static optimizer states and gradients consume 8x the memory of model parameters alone.
The 16-Bytes-per-Parameter Rule (State)
For parameter count N: Model Weights (BF16) = 2N bytes. Gradients (BF16) = 2N bytes. Adam Momentum (FP32) = 4N bytes. Adam Variance (FP32) = 4N bytes. Master Weights (FP32) = 4N bytes. Total static memory = 16N bytes!
Dynamic Activations
In addition to the 16N static bytes, intermediate layer outputs (activations) stored for backprop grow linearly with sequence length N, batch size B, hidden size d, and layer count L.
Why Sharding Is Mandatory
A 70B model requires 70B * 16 bytes = 1,120 GB static memory — impossible to fit on a single 80GB GPU without distributed sharding.
MATHEMATICAL FORMULATION · STATIC TRAINING VRAM BREAKDOWN
Static_Memory = N × ( 2_weights + 2_grads + 4_momentum + 4_variance + 4_master ) = 16 N bytes
Calculating raw static bytes before adding activation tensors and temporary workspace buffers.
REAL-WORLD PRODUCTION ENGINEERING
ZeRO-Stage 1 shards the 12N bytes of Adam state across data-parallel ranks, reducing per-GPU memory by up to 8x.