Activation Checkpointing: Trading Recomputation for Stored Tensors

Backward through a layer needs that layer’s input, and for some layers it needs intermediates computed in the forward pass as well. The straightforward approach keeps all of them, which makes activation memory grow linearly with depth and with the number of stored tensors per layer. Checkpointing replaces that storage with a second forward pass: keep only the inputs to a few chosen boundaries, and when backward reaches a segment, rerun its forward computation to reconstruct the internals it needs.

The memory saving is close to the segment count; the cost is one additional forward pass over the model, plus the bookkeeping to know which segments to rerun and when to free them.

The dependency that forces storage

Autograd records, for each operation, which tensors are needed to compute the vector-Jacobian product. For a matrix multiply $y = Wx$, the backward pass needs $x$ and $W$. For a softmax, it needs the output. For a normalisation, it needs the saved mean and variance. The framework retains these because recomputing them would require the input to the operation, which in turn requires the input to the previous operation, and so on back to the start of the network.

Storing everything is the memoised version of the chain. Checkpointing is the observation that the chain can be cut at chosen points: if the input to a segment is retained, everything inside the segment can be recomputed from it on demand. The retained inputs form the checkpoint set, and their number controls the memory-time trade.

import torch
from torch.utils.checkpoint import checkpoint

def segment(x, blocks):
    for blk in blocks:
        x = blk(x)
    return x

def forward_with_checkpoints(x, blocks, per_segment=4):
    for i in range(0, len(blocks), per_segment):
        group = blocks[i:i + per_segment]
        # Only the group's input survives; internals are recomputed in backward.
        x = checkpoint(segment, x, group, use_reentrant=False)
    return x

The recomputation runs with gradient tracking disabled during the backward pass, so the reconstructed activations do not themselves build a graph. Only the segment boundaries remain as graph nodes.

The memory law for uniform segmentation

Split $L$ layers into $g$ segments of $L/g$ layers each. Retaining one input tensor per segment costs $g$ tensors; retaining nothing inside a segment means the peak inside a segment is the transient working set of a single layer, not the accumulated set of all layers in it.

Peak memory is therefore the checkpoint set plus the largest transient inside the segment currently being recomputed. Written as a function of $g$: memory is $O(g + L/g)$ tensors. That function has a minimum at $g = \sqrt{L}$, giving $O(\sqrt{L})$ memory rather than $O(L)$.

The square-root result is the reason the technique is described as sublinear memory cost. It is also the reason the returns diminish quickly: moving from $g = 1$ to $g = 2$ halves the stored set and adds one extra forward pass over the whole model; moving from $g = \sqrt{L}$ to $g = 2\sqrt{L}$ saves a smaller absolute amount for the same added compute. In practice the segment count is chosen from the memory residual, not from the optimum of the formula, because the formula ignores the fixed per-layer workspace that does not shrink with segmentation.

What the extra forward pass actually costs

Recomputing a segment repeats its forward FLOPs. If every segment is recomputed exactly once, the total is one extra forward pass over the model: the step goes from $6ND$ to $8ND$ FLOPs, a 33 percent increase in arithmetic, as the step accounting makes explicit.

That figure assumes the recomputation happens once per segment. A naive implementation that recomputes a segment for every operation inside it would do far worse, which is why the checkpoint is placed around a group of layers rather than around each operation. The overhead also depends on whether the recomputation is itself checkpointed: nested checkpointing multiplies the recompute factor, and recursive schemes can reach $O(\log L)$ passes if the nesting is deep.

The wall-clock cost is usually below the FLOP ratio, for two reasons. The recomputed forward pass runs without autograd bookkeeping, so it does not pay the graph-construction cost. And the recomputation is interleaved with backward, which keeps the device busy with a mix of forward and backward work rather than idling during a separate phase. On hardware where the backward pass is already the bottleneck, the added forward work fits into slack that the memory-bound backward pass leaves.

The cost can exceed the ratio in one case: when the recomputed segment’s activations are themselves large enough that the recomputation evicts useful data from cache. This is a second-order effect and rarely dominant.

Which layers are worth checkpointing

Not all stored tensors cost the same. A checkpoint around a group of cheap layers saves little and adds their recompute time; a checkpoint around an attention block with a long sequence saves the score matrix, which is the largest single tensor in the model.

The selection criterion is bytes saved per FLOP of recomputation. A layer that stores $k$ tensors of $B \cdot S \cdot H$ elements and costs $c$ FLOPs to recompute has a ratio of $2k B S H$ bytes per $c$ FLOPs. Layers with many stored intermediates relative to their arithmetic, such as normalisation followed by a nonlinearity, are good candidates. Layers that are already cheap in storage, such as a pure elementwise residual add, are not.

Attention deserves separate treatment. The score matrix grows with $S^2$, so at long sequence lengths the recompute-versus-store decision for attention alone can determine whether the batch fits. Two mechanisms address it: recompute the scores from the stored queries and keys, which is what a checkpointed attention block does; or use an attention implementation that never materialises the full matrix in the first place, computing the softmax in tiles. The second avoids both the storage and the recomputation, at the cost of a more complex kernel. Sequence-length effects on memory explain why this term dominates at long context.

Correctness hazards

Checkpointing reruns the forward computation, and anything in that computation that is not a pure function of its inputs will differ between the two runs. Dropout is the canonical case: the mask drawn during the original forward pass is discarded, and the recomputation draws a new one. The gradient is then computed against activations that never existed in the first place. The result is not an error in the mathematical sense but a different, noisier estimator, and frameworks handle it by restoring the random state at recomputation time or by treating the mask as part of the checkpointed state.

Batch normalisation has the same problem from the other direction: the running statistics update during the recomputation, so a checkpointed layer updates its buffers twice per step. For layers with momentum this is a small bias; for layers whose buffers feed evaluation, it changes what evaluation sees. The standard mitigation is to update such buffers only during the original forward pass.

A third hazard is state mutated inside the checkpointed function. If the segment appends to a list, increments a counter, or writes to a module attribute, the recomputation repeats the mutation. Checkpointing assumes the segment is a function of its tensor inputs, and violations produce duplicated side effects that are invisible until the state is inspected.

Where it sits among the memory levers

Checkpointing reduces the activation term and nothing else. It does not touch parameters, gradients, or optimiser state, so it cannot rescue a run whose memory is dominated by those. When the activation term is the largest, it is often the cheapest lever available, because it requires no change to the model architecture and no additional communication.

It composes with the other levers without conflict. A smaller micro-batch reduces the activation term proportionally; checkpointing reduces it to roughly the square root of its depth-dependent part. Combining both gives a product of savings rather than a sum, which is why a run that needs a very large effective batch typically ends up using accumulation for the batch dimension and checkpointing for the depth dimension at the same time. The remaining term, if memory is still short, is the sharded state described in ZeRO partitioning.