Interactive explainer · Phase 3 — Training & systems
Efficiency & GPU systems
A training run is a budget of bytes moved and FLOPs spent, and every efficiency trick in the field is one of two moves: move fewer bytes, or waste fewer FLOPs. This page is the bookkeeping — where the memory actually goes, when the GPU computes versus waits, and what FlashAttention and parallelism really buy. Every number is computed from published formulas and H100 spec sheets.
— every op pays whichever bill is bigger. that's the whole chapter.
§ 1 · Where the memory goes
The optimizer is the tenant
Ask someone untrained where a 7B model's memory goes and they'll say "the weights." Wrong by 8×. AdamW mixed-precision training carries 16 bytes per parameter — weights, gradients, fp32 master copies, and two Adam moments — before a single activation is stored. Slide the model size and watch which cards survive.
mode
total: 144.5 GB — OOM on every single card — ~2× 80GB worth of memory. This is why ZeRO, offload, and recomputation exist.
The states alone are 112.0 GB — 16 bytes riding on every parameter, ~8× what bf16 inference needs. The optimizer, not the model, is the tenant.
fig. 1 — shape derived from N ≈ 12·L·d² (d ≈ 128·L); activations use Megatron's no-recomputation estimate. real configs differ in detail, not in moral.
§ 2 · Compute-bound or memory-bound
One plot decides
For every operation, divide the arithmetic it does by the bytes it moves. That single number — arithmetic intensity — against the GPU's two ceilings (compute peak, memory bandwidth) tells you whether the silicon is working or waiting. Drag the matmul dimensions; find the ridge.
operation
intensity 683 FLOP/B · delivers 989 TFLOP/s · 100% of peak
Compute-bound. The multipliers are saturated — this op earns the GPU its price tag. Big matmuls live here; almost nothing else does. MFU (model-FLOPs-utilization) is this same ratio measured for a whole training run — real runs land at 30–50% because attention, norms, and everything that isn't a big matmul drags the average down.
fig. 2 — the roofline model. H100 SXM stated specs: 989 TFLOP/s dense bf16, 3.35 TB/s HBM3. everything left of the gold ridge is waiting on memory.
§ 3 · The FlashAttention idea
Fewer bytes, not fewer FLOPs
Attention's score matrix is T×T. At long context that matrix — a temporary, used once — becomes the biggest thing on the card. FlashAttention's insight: it's a memory-bound op, so stop writing the matrix. Compute it in on-chip tiles, carry running softmax statistics, and never let the T² bytes touch HBM.
at T = 32,768 · naive 2.1 GB vs flash 9 MB — per head
The trick is tiling + online softmax: stream Q and K through fast on-chip SRAM in blocks, keep a running row-max and row-sum, and rescale the partial output as each block arrives. The T×T matrix is still computed — it's just never written to HBM. Same FLOPs, exact same answer (it is not an approximation), a fraction of the bytes — which is the whole game on a memory-bound op.
fig. 3 — one head, bf16. multiply by head count for the full bill; the shape of the argument doesn't change.
§ 4 · Splitting the model
Three ways to cut a network
When one GPU isn't enough, there are exactly three axes to cut along: the batch, the matmuls, or the depth. Each trades a different thing over the wire.
- replicated
- the whole model, on every GPU
- sharded
- the batch — each GPU sees different examples
- communicated
- gradients, all-reduced once per step
- reach for it when
- the default. Cheap to reason about, scales until the model itself no longer fits on one device.
On one GPU you need the reading-level grasp, not the ops experience. One connection worth keeping: ZeRO / FSDP is data parallelism with the optimizer states sharded — it attacks exactly the 16-bytes-per-parameter tenant from §1.
§ 5 · The whole chapter
Bytes and FLOPs, nothing else
memory = optimizer states
Training carries 8× the memory of bf16 inference. ZeRO, offload, LoRA — all attacks on this line.
intensity decides
Below the ridge the GPU waits on HBM; above it, it computes. Big matmuls are the only citizens above.
flash = fewer bytes
Same FLOPs, exact same answer. Speed came from traffic, not arithmetic — memorize this shape of win.
When you profile your own training run in Phase 3 — and you should, that's the milestone — this page is the map you read the trace against: is the time in matmuls (fine), in memory-bound soup (fuse it), or in waiting on data (fix the loader)? For the real depth, CS336's systems assignments make you build the Triton kernel.