Activation Checkpointing › Selective Checkpointing

Full per-layer checkpointing is maximally conservative: it recomputes everything. Selective checkpointing applies checkpointing only to the parts of the computation that are most memory-expensive, leaving cheap activations in memory. The goal is to recover most of the memory benefit while adding less than the 33% compute overhead of full checkpointing.

The key insight is that not all activations are equally expensive to store, and not all recomputations are equally expensive in terms of compute. By identifying which activations are large but cheap to recompute, you can selectively discard exactly those tensors while keeping smaller or more expensive-to-recompute activations in memory.

The most common form of selective checkpointing distinguishes between the attention sublayer and the feed-forward sublayer within each transformer block.

Attention sublayers with long sequences produce the expensive $[B, H, T, T]$ weight matrices. Recomputing attention is inexpensive because the stored Q, K, and V projections make the recompute path short: one matrix multiply and softmax to regenerate the attention weights. So attention sublayers are high-value checkpoint targets: large tensors to save, cheap to recompute.

Feed-forward sublayers have large intermediate tensors (shape $[B, T, 4d]$) but also have short recomputation paths. The input to the FFN is available from the post-attention residual, and the intermediate activation is just one matrix multiply and GELU away. So FFN intermediates are also good candidates for checkpointing.

Layer norm activations (the mean and variance per token) are tiny by comparison and not worth checkpointing. The Q, K, and V projection matrices themselves are moderate in size and are needed for the attention recomputation, so they are typically kept.

A selective policy might checkpoint at the beginning of each attention block and each FFN block, skipping layer norm intermediate states and small projection activations that are cheap to store relative to their recomputation cost. This configuration typically achieves 70-80% of full-checkpointing memory savings at only 10-15% compute overhead rather than 33%.

In models that process variable-length inputs, activation memory scales as $O(T^2)$ for attention weight matrices but only $O(T)$ for feed-forward activations. Selective checkpointing can trigger more aggressively for long sequences and relax for short ones, keeping the memory budget balanced without imposing fixed overhead regardless of sequence length.

Concretely, you might implement a threshold: if the sequence length exceeds some value (say, 2048 tokens), enable attention checkpointing; if below, skip it. This is useful for fine-tuning datasets with highly variable sequence lengths, where most short sequences would waste compute overhead if uniformly checkpointed.

FlashAttention (Dao et al., 2022) implements a tiled attention algorithm that never materializes the full $[B, H, T, T]$ attention weight matrix in memory. Instead, it tiles the computation into blocks that fit in the GPU's SRAM (fast on-chip memory), computes softmax numerically stably in blocks, and recomputes the weight tiles during the backward pass from the stored query and key tiles. This is conceptually identical to activation checkpointing applied specifically to attention weights, but implemented at the CUDA kernel level for maximum efficiency.

The SRAM-level implementation makes FlashAttention's "checkpointing" much faster than explicit framework-level checkpointing. SRAM access is orders of magnitude faster than HBM (high-bandwidth memory), so recomputing from Q and K tiles that fit in SRAM adds minimal overhead. By contrast, framework-level checkpointing must read the checkpoint tensor from HBM, adding memory bandwidth cost on top of compute cost.

When using FlashAttention, the $O(T^2)$ attention memory cost is already handled. Remaining activation memory is $O(T)$, and further checkpointing provides smaller marginal benefit. In such settings, practitioners often apply checkpointing selectively to only the feed-forward intermediate activations, accepting the smaller memory cost of other activations to reduce compute overhead.

The combined FlashAttention plus selective FFN checkpointing configuration achieves roughly 80% of the memory savings of full per-layer checkpointing at only about 12% compute overhead. This is the configuration used in most modern large-scale training runs.

An orthogonal technique, sometimes combined with checkpointing, stores activation tensors at lower precision. If intermediate activations are quantized to int8 before being saved as checkpoints (and dequantized before reuse), checkpoint memory is halved compared to float16. The accuracy impact depends on the layer and gradient magnitude, but for many training scenarios the quantization error in saved activations is small relative to gradient noise. This approach is sometimes called "activation compression."

Quantized activation storage is particularly appealing for the residual stream checkpoint tensors. These large $[B, T, d]$ tensors are saved at layer boundaries and held for the duration of the backward pass. Quantizing them to int8 reduces their storage cost by 2x with minimal impact on gradient quality, since the gradient magnitudes in the residual stream are typically much larger than the quantization error.

The memory-compute tradeoff for different selective strategies can be visualized as a frontier curve. Each point represents a different checkpointing policy, from no checkpointing (maximum memory, minimum compute overhead) to full per-layer checkpointing (minimum memory, maximum compute). Selective strategies occupy the interior of this frontier.