Activation Checkpointing › Checkpoint Selection

Which activations to checkpoint is not arbitrary. The goal is to minimize total memory while minimizing recomputation cost. Several factors determine whether a layer is a good checkpoint candidate.

The natural granularity for checkpointing in transformer models is the transformer layer. Each transformer block is a closed computation unit: it takes one residual stream tensor as input and produces one residual stream tensor as output. Checkpointing at layer boundaries stores only these residual tensors, which are small relative to the internal activations of the layer.

For a model with hidden dimension $d = 4096$, batch size $B = 4$, and sequence length $T = 2048$, the residual stream tensor at each layer boundary has shape $[B, T, d] = [4, 2048, 4096]$. At float16 (2 bytes), this is:

$$4 \times 2048 \times 4096 \times 2 \approx 67 \text{ MB}$$

Storing 32 such tensors for a 32-layer model requires only about 2.1 GB, compared to the 64 GB needed for all intermediate activations: a 30x reduction for the checkpoint storage cost alone.

The reason the layer boundary tensor is so much smaller than the full layer activations is that it does not include the attention weight matrix. For the same configuration with 32 attention heads and sequence length 2048, the attention weight matrix has shape $[B, H, T, T] = [4, 32, 2048, 2048]$, occupying:

$$4 \times 32 \times 2048^2 \times 2 \approx 1.07 \text{ GB per layer}$$

Summed across 32 layers, attention weights alone account for about 34 GB. The residual stream checkpoints are tiny by comparison.

Within a single transformer layer, several activations are particularly expensive to store:

Attention weight matrices. The attention computation produces a weight matrix with shape $[B, H, T, T]$, where $H$ is the number of attention heads. The memory cost scales quadratically with sequence length:

$$\text{Memory}_{\text{attn}} = B \times H \times T^2 \times 2 \text{ bytes}$$

For $B = 4$, $H = 32$, $T = 4096$, this is $4 \times 32 \times 4096^2 \times 2 \approx 4$ GB per layer. For long-context models, this term dominates layer activation memory. Because the attention weights can be recomputed from the stored Q and K tensors (one matrix multiply to form $QK^T$ and one softmax), they are excellent candidates for not storing.

Feed-forward intermediate activations. The FFN sublayer expands the hidden dimension by a factor of 4, producing an intermediate tensor with shape $[B, T, 4d]$. At 4x the width of the residual stream, this is the largest single activation within the FFN sublayer. But unlike the attention weights, FFN intermediate activations grow only linearly with sequence length, so their relative importance increases for shorter sequences.

Layer norm statistics. Layer normalization stores per-token mean and variance for backpropagation. These have shape $[B, T]$ (one scalar per token), making them negligible compared to the above.

The cost-benefit of checkpointing within a layer depends on the ratio of activation size to recomputation cost. Feed-forward activations are expensive to store but cheap to recompute (one matrix multiply and GELU). Attention weights are very expensive to store for long sequences and also cheap to recompute from the stored Q and K tensors.

The visualization below shows how attention weight memory grows quadratically with sequence length, while feed-forward memory grows only linearly. For sequences beyond 2K tokens, attention memory completely dominates.

Finding the globally optimal checkpoint placement is an NP-hard problem in general, requiring solving a graph scheduling problem over the full computation graph. The Checkmate system (Jain et al., 2019) framed this as an integer linear program and showed that optimal placement can differ substantially from naive equal-spacing heuristics for irregular networks. However, for regular transformer architectures where all layers have identical structure, the heuristics work very well in practice.

The commonly used approaches are:

Equal spacing. Checkpoint every $k$ layers, with $k$ chosen based on the memory budget and acceptable overhead. Checkpointing every layer ($k = 1$) is a common, simple default because it minimizes the within-segment activations, although the boundary-plus-segment model above can favor a larger $k$ when the accumulated boundary tensors matter. Full recomputation adds roughly 33% to the total forward-plus-backward compute cost.

Recomputation cost weighting. Weight checkpoints toward computation-cheap layers. If two adjacent layers have very different compute costs, it is cheaper to recompute the cheaper layer and checkpoint after it.

Memory pressure-aware scheduling. In pipeline-parallel settings (as discussed in the Pipeline Parallelism chapter), different pipeline stages may have different numbers of micro-batches in flight. Stages with more micro-batches need more aggressive checkpointing, and per-stage checkpoint policies can tune the balance independently. This is important for Megatron-LM and similar frameworks that use interleaved pipeline schedules.

Activation-size-proportional checkpointing. Rather than spacing by layer count, space checkpoints by cumulative activation memory. This produces uniform memory usage per segment even when layers vary in their activation footprint, which is useful for models with heterogeneous layer sizes.

In practice, most large-scale training runs use one of two configurations: full per-layer checkpointing (when memory is extremely tight) or selective checkpointing with FlashAttention (when a good balance of memory and compute is needed). The theoretical optimal for irregular networks is rarely computed in production because transformer layers are regular enough that equal-spacing gives near-optimal results.