Activation Recomputation: When Saving Memory Costs You More Than You Think

Activation recomputation (also called gradient checkpointing) is one of those techniques that sounds like a free lunch until you've run it in production for a week. The pitch: instead of storing every intermediate activation during the forward pass, you recompute them on demand during the backward pass. Memory pressure drops. You fit a larger batch, or a bigger model, or both. Notebooks reproduce cleanly.

Photo by panumas nikhomkhai on Pexels.

Then you notice training throughput has fallen 20-40% and you start asking questions.

The core tradeoff is real and worth respecting. A standard transformer forward pass stores activations at every layer so the backward pass can compute gradients without redoing work. For a model with $L$ layers and sequence length $S$, that activation buffer scales roughly as $O(L \cdot S \cdot d_{model})$. On long sequences or wide models, this becomes the dominant consumer of GPU memory, crowding out optimizer states and even the weights themselves. Recomputation breaks that scaling: you pay $O(\sqrt{L})$ or $O(1)$ memory (depending on the checkpointing scheme) and accept one extra forward pass per layer during backprop.

The question is never whether the memory savings are real. They are. The question is what you've given up to get them.

Where compute cost actually lands

PyTorch's torch.utils.checkpoint.checkpoint() is the standard entry point. When you wrap a subgraph, PyTorch discards activations after the forward pass and reruns that subgraph during backward. The recomputation happens synchronously on the backward stream, which means it serializes with gradient computation. On a single GPU this is tolerable. Across a pipeline-parallel training job, it's a different story: the recomputation now sits on the critical path between pipeline stages, and any bubble you'd already budgeted gets wider.

Profile this with torch.profiler before assuming the overhead is acceptable. Look at the CUDA kernel timeline. If you see long stretches of recompute kernels interleaved with backward kernels and no overlap with optimizer updates, you've found your regression.

Selective recomputation is the actual tool

Full recomputation, applied naively to every layer, is rarely the right call. Selective checkpointing lets you designate which subgraphs to recompute and which to retain. The heuristic: checkpoint operations that are cheap to recompute (elementwise activations like GeLU, LayerNorm, dropout) and keep the outputs of expensive operations (attention score matrices, large linear projections). Recomputing a softmax over a full attention matrix on a 4096-token sequence costs almost nothing. Recomputing a 16384x16384 matmul does not.

DeepSpeed's ZeRO-R and Megatron-LM both expose granular control here. In Megatron, --recompute-granularity selective combined with --recompute-method uniform lets you checkpoint every $k$-th transformer block rather than every layer uniformly. Tuning $k$ is empirical: profile at $k=1$ (full recompute), $k=4$, and $k=8$ and measure actual step time versus peak memory. The memory-compute Pareto curve is not linear, and the knee of the curve is usually somewhere around $k=2$ to $k=4$ for standard transformer blocks.

graph TD
    A[Forward Pass] --> B{Checkpoint boundary?}
    B -- Yes --> C[/Discard activations/]
    B -- No --> D[Retain activations]
    C --> E[Backward Pass]
    D --> E
    E --> F{Need activation?}
    F -- Discarded --> G((Recompute subgraph))
    F -- Retained --> H[Use cached value]

When recomputation makes things worse end-to-end

Consider a job running on 64 A100s with tensor parallelism degree 8 and pipeline parallelism degree 4. Each pipeline stage already has a bubble at the start and end of each microbatch sequence. Adding full activation recomputation to every stage effectively doubles the bubble cost because the recompute phase cannot overlap with inter-stage communication. You've traded memory headroom for pipeline efficiency, and the net effect on tokens-per-second-per-dollar is negative.

In this configuration, the better move is often to reduce the tensor parallelism degree (freeing memory by reducing the activation slice each rank holds) or to switch to a memory-efficient attention kernel like Flash Attention, which avoids materializing the full $O(S^2)$ attention matrix entirely. Recomputation is one lever among several; reaching for it first because it's the most familiar lever is a habit worth questioning.

Observability you actually need

Track peak reserved memory per rank (torch.cuda.max_memory_reserved()), forward-pass wall time, backward-pass wall time, and step time as separate metrics, not just aggregate throughput. When recomputation is active, you want to see the backward-pass time increase by roughly the cost of one forward pass. If it increases by more, something else is serializing: a synchronization barrier, a host-device copy, or a fragmented memory allocator forcing extra cudaMalloc calls.

Recomputation is a legitimate tool. Use it with accounting.