Activation Checkpointing › Historical Development and the Path to Modern Checkpointing

Activation checkpointing did not emerge from deep learning research: it arrived from scientific computing. Researchers working on optimal control and parameter estimation in the 1980s faced a structurally identical problem. They needed to compute gradients of an objective function through long sequences of operations, and storing all intermediate states was prohibitively expensive on the hardware of the time. Andreas Griewank's 1992 paper "Achieving logarithmic growth of temporal and spatial complexity in reverse automatic differentiation" established the theoretical foundation, showing that a divide-and-conquer checkpoint placement strategy could reduce the temporal complexity of gradient computation to $O(n \log n)$ at the cost of $O(\log n)$ memory, where $n$ is the number of time steps. This is the same trade that modern activation checkpointing exploits, applied to the layer dimension rather than the time dimension.

The first direct application to neural network training appeared in the context of recurrent networks. Training long recurrent neural networks with backpropagation through time (BPTT) required storing a hidden state for every time step, which made very long sequences impractical. Researchers applied checkpointing along the time axis: save hidden states every $k$ steps, recompute intermediate states on demand during the backward pass. This "truncated BPTT with checkpointing" was used in practice well before the transformer era.

Chen et al.'s 2016 paper "Training Deep Nets with Sublinear Memory Cost" brought the idea to feedforward networks and proved the $O(\sqrt{n})$ memory bound for equal-spacing checkpointing with $k = \sqrt{n}$. This paper was influential because it showed that the technique applied to arbitrary feedforward architectures, not just recurrent ones, and provided a clean theoretical analysis of the tradeoff. The paper also demonstrated that memory savings of 10x or more were achievable for networks with tens of layers at less than 33% compute overhead, which matched what practitioners were finding empirically.

PyTorch incorporated torch.utils.checkpoint in version 0.4.0 (2018), and the API has been stable since. The initial implementation used the reentrant approach, which had correctness limitations. The non-reentrant implementation added in later versions fixed these issues and is now the recommended interface. Jax, the other major deep learning framework used for large-scale training, provides jax.checkpoint (also called jax.remat) with similar semantics and has been used extensively in training models like PaLM and Gemini.

The scaling of transformer models through 2020-2023 made activation checkpointing essential rather than optional. Models like GPT-3 (175 billion parameters), Megatron-Turing NLG (530 billion parameters), and the Chinchilla family could not have been trained within practical memory budgets without combining activation checkpointing with tensor and pipeline parallelism. The Megatron-LM codebase, which underlies most large-scale transformer training at NVIDIA and Microsoft, has activation checkpointing enabled by default for any model above a modest size threshold.

The discovery of FlashAttention in 2022 represents the most significant recent advance in the checkpointing ecosystem. By fusing the attention computation into a single CUDA kernel that tiles across SRAM, FlashAttention achieves the memory profile of activation checkpointing for the attention sublayer while running at or above the speed of standard attention on hardware with favorable compute-to-bandwidth ratios. FlashAttention-2 and FlashAttention-3 further extended these gains with improved parallelism and hardware-specific optimizations. The practical effect was to make long-context training tractable: 128K-token context windows that would have required petabytes of activation memory with naive attention are now trainable on standard GPU clusters.

The next frontier is automatic checkpoint placement. Current frameworks require the user to specify where to checkpoint, either at the per-layer granularity or via segment counts. Automatic differentiation systems that can analyze the computation graph and determine optimal checkpoint placement without user intervention would eliminate this manual tuning step. Research systems like Checkmate have demonstrated that optimal placement can provide 20-30% better memory-compute tradeoffs than equal-spacing heuristics for heterogeneous networks. Incorporating this into production training frameworks remains an active research direction.