← Splitting the budget: Chinchilla⌂ A Language Model, From ScratchFlashAttention: the same maths, different memory traffic →
Part 5 · The Maths III: Training at Scale
The memory wall
FLOPs are rarely what stops you; bytes are. A mixed-precision AdamW run holds, per parameter: 2 bytes of bf16 weights, 2 of gradients, and an fp32 master copy plus two optimizer moments at 4 bytes each.
2 + 2 + 4 + 4 + 4 = 16 bytes per parameter
A 70B model therefore needs about 1.1 TB before a single activation is stored — roughly fourteen 80 GB accelerators just to hold state. ZeRO 1910.02054 attacks this by sharding rather than replicating: stage 1 partitions optimizer state, stage 2 adds gradients, stage 3 adds the parameters themselves, taking per-device state from 16Ψ down to 16Ψ/Nd at the cost of extra collective communication.