The Problem
Training a transformer stores, for each layer, the inputs to every op that is differentiated in backward: the attention inputs, the Q/K/V projections, the softmax output, the FFN inputs, the norm outputs, and the residual stream. For a layer with hidden size d, sequence length L, batch B, this is on the order of 12 * B * L * d floats per layer.
For d=8192, L=8192, B=1, that's about 805M values, or 1.6 GB per layer in BF16. A 64-layer model is about 103 GB of activations — and that's before you multiply by microbatch size, before you add attention-softmax intermediates (L^2 per head), and before you factor tensor-parallel partial copies.
The two-sided bill: BF16 weights plus optimizer state might fit in 80GB, but activations push you past. Gradient checkpointing (aka activation recomputation) is the standard fix. Drop most activations; redo the forward during backward to get them back. Cost: extra FLOPs. Benefit: memory drops by the ratio of checkpoint segments to total layers.
Done naively, checkpointing costs roughly 33% more FLOPs per training step. Done well — selective checkpointing per the "smart selection" of Korthikanti et al. — you save 5x memory for under 5% FLOP overhead. And with FP8 matmuls, FSDP offload, and expert-parallel MoE this really matters: you can't afford either the memory or the wasted compute.
Unlock the full lesson
You've read the first 2 sections. The rest of this lesson covers The Concept, Build, Check Yourself, Check Yourself, Key Terms & Next — plus a hands-on lab, quiz, and project artifact.
Create a free account — the opening phases of 29 of 35 courses are free, no credit card.