Activation Checkpointing › Activation Checkpointing

Training large language models is fundamentally a memory-constrained problem. A single forward pass through a transformer stores activations for every layer so that backpropagation can compute gradients. These stored tensors are not small: for a 7B-parameter model with a batch size of 4 and a sequence length of 2048, the activations alone can occupy over 30 GB of GPU memory, often more than the model weights themselves. When memory runs out, you cannot increase batch size, cannot train larger models, and cannot fit more context without buying more hardware.

Activation checkpointing (also called gradient checkpointing or rematerialization) solves this problem by trading computation for memory. Instead of storing every activation produced during the forward pass, the technique stores only a subset of them at strategic points in the network, called "checkpoints." During the backward pass, when gradients need to flow through a layer whose activations were discarded, those activations are recomputed on the fly from the nearest checkpoint. You pay a one-time recomputation cost but avoid storing the full activation buffer. For most architectures, this exchange is highly favorable: recomputing one forward pass per checkpoint adds roughly 33% more compute, but memory usage drops by a factor proportional to the number of layers between checkpoints.

The idea predates modern deep learning, tracing back to reverse-mode automatic differentiation research in the 1980s. In their foundational work on adjoint methods, researchers noticed that storing all intermediate values in a long chain of computations creates prohibitive memory requirements. The "checkpointing" strategy they developed for scientific computing is essentially the same concept: save only selected states, recompute the rest on demand. When deep learning rediscovered this idea for neural networks in the mid-2010s, tools like the Theano-based memory-efficient backpropagation library and later Chen et al.'s 2016 paper "Training Deep Nets with Sublinear Memory Cost" gave it renewed attention. Checkmate and similar analysis frameworks later formalized the optimal checkpoint placement problem as a graph-theoretic optimization. Today, frameworks like PyTorch expose activation checkpointing as a single function call via torch.utils.checkpoint.checkpoint, making it accessible without manual implementation.

This chapter explains how activation checkpointing works, why the memory-compute tradeoff is so favorable, how to choose which layers to checkpoint, and how to implement selective checkpointing to minimize overhead while maximizing memory savings.