AI ENGVisual Encyclopedia

MODULE 4: DISTRIBUTED INFRASTRUCTURE · SCENE 13

The state you pay to keep

Weights are the cheap part. Gradients and Adam moments quadruple the bill before a single token flows.

STATIC STATE (16N) · 1.1 TBPER-GPU SHARDED · 1.1 TB

AdamW mixed precision requires 16 bytes per parameter: Weights (2B) + Gradients (2B) + Adam m,v (8B) + Master Weights (4B).

TECHNICAL BREAKDOWNModule 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.