Activation checkpointing stores selected intermediate activations and recomputes the others during backpropagation. Recomputing activations is a standard way to fit the training of large models into GPU memory (Korthikanti et al., 2023); PyTorch (torch.utils.checkpoint) and JAX (jax.checkpoint) both provide it.
How to study this pagePredict · Check · Generalize · Continue
PredictSay what you expect.
Pick a concrete input and say what the equation should give before you read the math.
CheckCheck against the math and the papers.
Work through the math, then read the cited papers for the original argument.
GeneralizeName what stays the same.
Identify the part of the idea that survives a change of notation or model, so you can recognize it in another paper or system.
ContinueFollow the connections.
Use the prerequisites and related concepts at the end of the page to fill the next gap.
Why It Matters for Modern Models
Activation checkpointing stores selected intermediate activations and recomputes the others during backpropagation. Recomputing activations is a standard way to fit the training of large models into GPU memory (Korthikanti et al., 2023); PyTorch (torch.utils.checkpoint) and JAX (jax.checkpoint) both provide it.
Chen et al. (2016) showed that an n-layer network can be trained with O(√n) activation memory for the cost of one extra forward pass per mini-batch.
Recomputing everything is wasteful: Korthikanti et al. (2023) combined selective recomputation with sequence parallelism to cut activation memory by a factor of 5 while removing over 90% of the time overhead of full recomputation.
Questions to Think About
Why do √L segments of √L layers minimize peak memory, and how does the best split change when layers have very different activation sizes?
Why does recomputation cost about one extra forward pass, roughly a third more compute, rather than doubling the cost of training?
Core Math
Key Equation
M(k)≈kL+k⇒k∗=L
Why activations take memory. Backpropagation needs the activations from the forward pass to compute gradients, so a standard implementation keeps the output of every layer until the backward pass reaches it. For L layers this costs memory proportional to L, and each layer's activations also grow with batch size and sequence length.
Checkpointing stores only some activations and recomputes the rest (Chen et al., 2016):
Forward: split the L layers into L/k segments of k layers and keep only the input of each segment.
Backward: before back-propagating through a segment, rerun its forward pass from the stored input.
Counting in units of one layer's activations, peak memory is about L/k stored segment inputs plus the k activations of the segment being processed:
M(k)≈kL+k,minimized at k=L, where M≈2L.
The price is one extra forward pass. A backward pass costs about twice a forward pass, so training compute rises from about three forward passes' worth to four, roughly a third more.
Applying the idea recursively reduces memory to O(logL) at the cost of O(LlogL) forward computation (Chen et al., 2016), and Griewank & Walther (2000) give schedules that are optimal for a given number of checkpoints.
Selective recomputation recomputes only the parts of a transformer layer that need much memory but little arithmetic, such as the attention softmax and its dropout, and stores the rest (Korthikanti et al., 2023).