Activation Checkpointing › Limitations and Practical Considerations

Activation checkpointing is not universally beneficial, and several practical constraints limit its applicability.

The 33% compute overhead is a reasonable cost when the training bottleneck is GPU memory capacity. But not all training configurations are memory-bound. If you are training a small model on a large GPU where activation memory is a minor fraction of total memory, the added recomputation moves you from memory-bound to compute-bound, wasting time without enabling anything new.

The technique is most valuable when it is the enabling factor for a configuration that would otherwise cause an out-of-memory error or require a reduced batch size. A practical test: if you can already train at your desired batch size and sequence length without checkpointing, the overhead is pure waste. Enable checkpointing only when it allows you to increase batch size, fit a larger model, or extend sequence length in ways that improve training efficiency or model quality.

On modern hardware like H100 GPUs with HBM3 memory, the compute-to-memory-bandwidth ratio is much higher than on older A100 or V100 GPUs. This means recomputation is relatively cheaper in terms of wall time because the GPU executes matrix multiplications much faster relative to memory accesses. On H100 hardware, the effective overhead of full per-layer checkpointing is often closer to 20-25% rather than 33%, because memory bandwidth savings from smaller activation tensors accelerate the overall forward pass.

The use_reentrant=False mode requires that all inputs to a checkpointed segment that require gradient computation are properly handled by PyTorch's autograd graph. Some operations do not play well with non-reentrant checkpointing, including certain custom CUDA extensions and operations that use in-place mutations on tensors that are returned as outputs.

When you encounter unexpected errors in checkpointed regions, the first diagnostic step is to run without checkpointing to isolate whether the error is checkpointing-specific. If the model runs correctly without checkpointing but fails with it, the issue is almost always one of: (1) in-place modification of a tensor that appears in the checkpoint's computation graph, (2) a custom autograd function that does not correctly implement its backward pass under checkpointing, or (3) operations that break the RNG state restoration.

PyTorch 2.0's torch.compile performs whole-graph optimizations that can conflict with the boundaries imposed by checkpoint.checkpoint. The compiler may not be able to fuse operations across checkpoint boundaries as aggressively as it would without them. In some configurations, using torch.compile with checkpointing requires careful configuration of the fullgraph=False option or selective application of compilation to non-checkpointed regions.

The underlying tension is that torch.compile works best when it can see the entire computation graph and optimize across it globally. Checkpoint boundaries break the graph into segments that are individually compiled, preventing cross-boundary fusion. For some operations, this can negate a significant fraction of the compilation speedup. This is an active area of tooling improvement in PyTorch, and the interaction is expected to improve in future releases.

During the backward pass, when a segment is being recomputed, there is a brief period where both the checkpoint input activation (from the saved boundary) and the newly recomputed intermediate activations coexist in memory. This creates a temporary memory spike above the steady-state checkpointed level.

For a checkpointed model with $L/k$ boundary tensors and $k$-layer segments, the steady-state activation memory is approximately $L/k \cdot M_{\text{ckpt}}$ (boundaries) plus $k \cdot M_{\text{layer}}$ (current recomputation segment). During the recomputation spike, you additionally need the previous segment's boundary tensor $M_{\text{ckpt}}$ (to re-run from). So the peak is slightly higher than the steady state by one extra boundary tensor.

For models operating very close to the memory limit, this spike can still cause an out-of-memory error even with checkpointing enabled. The solution is to either increase checkpoint frequency (smaller segments, more boundaries, lower spike) or reduce the batch size slightly as a buffer.

For practitioners setting up activation checkpointing for a new training run, the following decision process covers most scenarios:

Start with a memory profile. Before enabling checkpointing, measure memory usage in your training configuration. Use torch.cuda.memory_allocated() and torch.cuda.max_memory_allocated() at different points in the training step to identify the peak. If the peak occurs during the forward pass (after many layers have been processed), activation memory is the culprit. If the peak occurs at the optimizer step, optimizer state memory is the constraint, and checkpointing will not help.

Enable FlashAttention first. If you are not already using FlashAttention, enable it before considering explicit activation checkpointing. FlashAttention provides 50-70% memory reduction for attention weights with near-zero compute overhead on modern hardware. After enabling it, re-profile to see how much activation memory remains.

Choose checkpointing granularity based on memory budget. If you are still memory-constrained after FlashAttention, enable per-layer checkpointing. Start with full per-layer checkpointing (k=1) to confirm correctness, then experiment with coarser granularity (k=2 or k=4) to recover some compute speed while staying within the memory budget.

Monitor training loss curves for anomalies. After enabling checkpointing, training loss should behave identically to the baseline (no checkpointing) run, since the gradients are mathematically equivalent. If training loss diverges or shows unusual spikes, disable checkpointing and compare: an anomaly that disappears without checkpointing indicates a compatibility issue with a specific operation in the model.

Account for the memory spike. If you are targeting a memory budget at the very limit of your GPU's capacity, leave a 10-15% buffer above your steady-state checkpointed memory usage to accommodate the recomputation spike during the backward pass. Running right at the limit risks out-of-memory errors that can be difficult to reproduce consistently.

Activation checkpointing has been a critical enabling technique throughout the scaling era of language models. GPT-3 (Brown et al., 2020) and subsequent models used gradient checkpointing to fit training within available GPU memory budgets. The training codebase for Megatron-LM (the framework used for GPT-3 and many subsequent large models) has included activation checkpointing as a core feature since its initial release.

The most significant recent development is the interaction between activation checkpointing and FlashAttention. By eliminating the $O(T^2)$ attention memory cost through kernel-level recomputation, FlashAttention made long-context training (sequences of 8K, 32K, or 128K tokens) feasible without proportionally increasing activation memory. For a sequence length of 128K tokens with 32 attention heads and batch size 1, the attention weight matrix alone would occupy:

$$1 \times 32 \times (128 \times 10^3)^2 \times 2 \approx 1 \text{ TB}$$

This is physically impossible to store. FlashAttention's kernel-level recomputation eliminates this memory requirement entirely, making 128K context training feasible on clusters of A100 GPUs.

The combination of FlashAttention (implicit attention checkpointing) plus selective FFN checkpointing plus mixed precision training (discussed in the next chapter) defines the current standard memory optimization stack for training large language models. This stack is what enabled models like Claude, GPT-4, and Gemini to train on very long documents at scale. Activation checkpointing alone was not sufficient to reach these context lengths, but it remains an indispensable component of the full stack.