Activation Checkpointing › Summary

Activation checkpointing addresses the fundamental memory bottleneck in deep network training by discarding intermediate activations during the forward pass and recomputing them on demand during backpropagation. The key concepts from this chapter are:

• The memory-compute tradeoff. Full per-layer checkpointing reduces activation memory by roughly 10-20x at the cost of approximately 33% more compute, by running a second forward pass through each layer during backpropagation. The overhead is only 33% rather than 100% because the backward pass already costs twice the forward pass.

• Why backpropagation needs activations. Gradient computation at each layer requires knowing the values that were present during the forward pass. Nonlinear operations like GELU cannot be inverted, so their inputs must be stored or recomputed. Every activation dependency creates a memory reservation that lasts until the corresponding backward pass gradient is computed.

• Checkpoint placement. Transformer layer boundaries are natural checkpoint points because they have small, well-defined activation tensors (just the residual stream). Within a layer, attention weight matrices and FFN intermediate activations are the most expensive activations to store, making them the highest-priority targets for selective checkpointing.

• Selective checkpointing. Applying checkpointing only to high-cost activations while retaining cheap activations in memory reduces compute overhead below 33% while recovering most of the memory benefit. The FlashAttention-plus-selective-FFN configuration achieves around 80% memory savings at roughly 12% compute overhead.

• FlashAttention interaction. FlashAttention implements implicit attention-weight checkpointing at the kernel level, handling the $O(T^2)$ memory cost without explicit overhead. When FlashAttention is used, remaining checkpointing targets are primarily feed-forward layers.

• Distributed training interactions. Activation checkpointing interacts with pipeline parallelism (reduces per-micro-batch memory at each stage), tensor parallelism (stacks multiplicatively), and gradient accumulation (each micro-batch's activations are independent). Understanding these interactions is necessary for configuring the full training infrastructure stack.

• Determinism requirement. Recomputed activations must exactly match original forward pass activations. PyTorch handles this automatically by saving and restoring random number generator states at checkpoint boundaries, which is critical for operations like dropout.

• PyTorch API. torch.utils.checkpoint.checkpoint(module, input, use_reentrant=False) implements per-layer checkpointing. checkpoint.checkpoint_sequential(modules, n_chunks, input) implements coarser segmentation for reduced overhead. The use_reentrant=False argument is required for the modern, correct implementation.

• Limitations. Checkpointing adds a temporary memory spike during recomputation, can interact poorly with torch.compile fusing, and is most valuable when activation memory is the binding constraint for a configuration. On hardware with high compute-to-bandwidth ratios (like H100), the practical overhead is often lower than the theoretical 33%.

• Historical significance. The technique originated in scientific computing in the 1980s and was formalized for neural networks by Chen et al. in 2016. It has been indispensable for training models at the scale of GPT-3 and beyond, and remains a core component of every large-scale language model training stack today.

Activation checkpointing is one of those techniques that becomes more important as models scale. A technique that saves 20x memory is helpful at small scales but essential at large scales. As language models continue to grow in size and context length, and as FlashAttention handles the attention memory problem, the remaining challenge is feed-forward memory scaling. Continuing advances in selective checkpointing, activation compression, and automated checkpoint placement will shape how practitioners manage this constraint for the next generation of models.