Gradient checkpointing, what does it trade, and roughly by how much?

**Gradient checkpointing** is a training-time technique that reduces GPU memory by trading extra compute for less activation storage. The baseline situation: during a normal forward pass, the model keeps all intermediate activations (layer outputs, attention scores, etc.) in GPU memory because backpropagation needs them to compute gradients. For a deep transformer with long sequences, these activations can be the largest single memory term, often larger than the weights themselves. Gradient checkpointing works by partitioning the model into segments and storing activations only at the segment boundaries. Activations inside each segment are discarded as soon as the forward pass exits the segment. When the backward pass reaches a segment, the framework re-runs that segment's forward pass to regenerate the discarded activations on demand, then uses them to compute gradients, then discards them again. The practical tradeoff: - Memory savings: roughly 60 to 70 percent less activation memory for typical transformer setups. This is the difference between OOM and a clean run on long context fine-tuning. - Compute cost: roughly 30 percent more wall-clock per step, because each segment's forward pass runs twice (once during the original forward, once during the recomputed backward). The feature is essentially free to enable in modern frameworks (`gradient_checkpointing=True` in Hugging Face Trainer) and is almost always on for any fine-tuning run with sequence length above 2048 or model size above 7B.

**Gradient checkpointing trades extra compute for activation memory.** During the forward pass it discards most intermediate activations between checkpointed segment boundaries; during the backward pass it re-runs each segment's forward to regenerate the discarded activations, then uses them for gradient computation. The motivation is that activation memory scales as O(L times S times H) for a transformer of L layers, S sequence length, and H hidden size, often exceeding weight memory for long-sequence fine-tuning. A 7B model in bf16 has 14 GB of weights but easily 30+ GB of activations at 8K sequence length and batch 4. Checkpointing collapses the activation term by a factor of roughly the segment count, at the cost of roughly one extra forward pass per backward. Typical magnitudes for a standard transformer with per-layer checkpointing: - **Activation memory**: roughly 60 to 70 percent reduction in peak training memory. - **Wall-clock per step**: roughly 30 percent increase from the extra forward pass. - **Throughput in tokens per second**: roughly 20 to 25 percent decrease. The exact numbers depend on the model and the checkpoint granularity. Coarser segments (one checkpoint per N layers) save less memory but cost less compute; finer segments (one per layer, the default in most frameworks) maximise memory savings at the larger compute cost. The feature composes with most other training techniques. It stacks cleanly with mixed precision, gradient accumulation, ZeRO sharding, and FSDP. It interacts subtly with FlashAttention (which has its own activation-memory optimisations) and with stage-3 ZeRO (which re-gathers parameters per checkpoint, potentially redundantly). In modern fine-tuning recipes for any model above 7B at any sequence length above 2K, gradient checkpointing is on by default.