Recomputation (Activation Checkpointing)

Definition

Recomputation (a.k.a. activation/gradient checkpointing) trades FLOPs for memory: instead of storing every forward activation for the backward pass, you store a sparse set of checkpoints and re-run the forward segments to regenerate the needed activations on demand. It is a core IO-aware technique — spend cheap, often-idle compute to relieve scarce memory (capacity and, on memory-bound ops, bandwidth). FlashAttention’s backward pass is the canonical specialized instance.

How it works

Backprop needs forward activations to form gradients. Storing all of them costs $O(L\cdot \text{activation size})$ HBM for an $L$-layer model — often the dominant memory term at long sequence/large batch. Checkpointing saves activations only at layer (or segment) boundaries; during backward it recomputes the intermediate activations of each segment from its stored input. With $\sqrt{L}$-spaced checkpoints, activation memory drops from $O(L)$ to $O(\sqrt{L})$ at the cost of one extra forward pass (~+33% training FLOPs). This is the gap between HFU (hardware FLOPs, including recompute) and MFU (model FLOPs utilization, useful-only) — recomputation raises HFU above MFU.

FlashAttention backward. Naive attention backprop stores the $N\times N$ probability matrix $P$ — reintroducing the $O(N^2)$ HBM footprint the forward pass eliminated. Instead, FlashAttention stores only the per-row logsumexp $L_i=m_i+\log {\ell }_i$ (an $O(N)$ vector) and recomputes each block’s scores/probabilities in SRAM: $P_{ij}=\exp (S_{ij}-L_i)$, exact and without re-running the softmax reduction. The extra $S$-recompute FLOPs run on Tensor Cores that were idle waiting on HBM anyway, so the recomputed backward is not only memory-frugal but faster than store-and-reload — because attention is memory-bound.

Why it matters

Recomputation is what lets frontier training fit long-context, large-batch activations in fixed HBM, directly enabling bigger effective batch sizes and sequence lengths per device. The subtle point: on a memory-bound op, “wasting” FLOPs to avoid bytes is a speedup, not just a memory saving — the same FLOPs-for-bytes bargain that defines FlashAttention and, more broadly, the co-design mindset for LLM systems. The trade is tunable (checkpoint granularity, selective recompute of only cheap ops) and interacts with parallelism strategy.

Taught in

• 06-memory-wall-and-flashattention — §5.3 the FlashAttention backward and the FLOPs-for-bytes bargain.
• 03-performance-modeling-roofline — the HFU/MFU gap.

See also

• flash-attention — recomputation in the backward pass
• io-aware-algorithms — the FLOPs-for-bytes principle
• memory-wall — why trading FLOPs for bytes pays
• mfu — HFU vs MFU
• mixed-precision-training · gpu-memory-hierarchy · arithmetic-intensity
• gpu-systems-for-llms