Activation Recomputation Math

What activations cost in the forward pass

Backpropagation needs the forward pass’s intermediate results to compute gradients, so by default a framework keeps every one of them alive until the backward pass consumes it. For a standard transformer layer in half precision, the stored activations add up to a well-known estimate:

A_layer ≈ s·b·h · (34 + 5·a·s/h)  bytes
s = sequence length, b = batch size,
h = hidden size,     a = attention heads

The 34sbh part collects the linear-size tensors: inputs to the QKV projections, the attention output, the two MLP matmuls and the GeLU input, the LayerNorm inputs, and the dropout masks — roughly seventeen tensor-equivalents at 2 bytes each. The 5as²b part is different in kind: it is the attention score matrix QK^T, its softmax, and the associated dropout mask, each of shape [b, a, s, s]. That term grows quadratically in sequence length, and at long s it dwarfs everything else.

The recomputation bargain: FLOPs for bytes

Recomputation (also called activation checkpointing) changes what the forward pass promises to the backward pass. Instead of “every intermediate is stored,” the promise becomes “a few checkpoint tensors are stored, and anything between two checkpoints can be rebuilt on demand by re-running that segment’s forward code.”

The canonical checkpoint is the input to each transformer layer — a single [s, b, h] tensor costing 2sbh bytes in fp16. During the backward pass, when gradients reach layer ℓ, the framework re-runs layer ℓ’s forward from its stored input, materializes the full A_layer set for that one layer, computes the gradients, and frees it all before moving to layer ℓ−1. Peak activation memory falls from all layers at once to checkpoints plus one live layer:

no recompute:   M = L · A_layer
full recompute: M ≈ L · 2sbh  +  A_layer   (one segment live)

Full recomputation costs one extra forward pass

The compute side of the bargain has a clean closed form. Per token and per parameter, a training step costs roughly 2 FLOPs for the forward pass and 4 for the backward (the backward differentiates through two matmul operands, doubling the work). So a normal step is about 6 FLOPs per parameter per token — a 1 : 2 forward : backward split.

Full recomputation re-runs each layer’s forward exactly once more during the backward sweep. The step becomes forward + (forward again) + backward:

overhead = extra / baseline = 2 / (2 + 4) = 1/3 ≈ 33%

That one-third figure is the number to remember: full recomputation never costs more than about 33% extra FLOPs, no matter how deep the model, because you replay each forward once and forwards are half the price of backwards. In wall-clock terms the observed slowdown is often a bit less than 33%, because the recomputed forward runs with better cache locality and the memory saved can buy a larger, more efficient batch.

Checkpoint placement and the sqrt(L) rule

Checkpointing every layer is one point on a curve, not the only one. Suppose you divide an L-layer network into k segments and checkpoint only segment boundaries. You store k boundary tensors, and during backward you must re-materialize one whole segment of L/k layers at a time. Peak memory is proportional to:

M(k) ∝ k + L/k
dM/dk = 1 − L/k² = 0  ⇒  k* = √L
M(k*) ∝ 2√L

This is the classic result of Chen et al. (2016): with a single level of checkpointing you can train an L-layer network in O(√L) activation memory for one extra forward pass. For L = 36, six segments of six layers store roughly 12 layer-inputs’ worth instead of 36 full layers’ worth. Recursive checkpointing pushes memory to O(log L), but each level replays another forward, and past one level the returns rarely justify the FLOPs — which is why practical systems stop at per-layer or √L-style schemes.

Worked example: a 24-layer model at s = 1024

Take a GPT-2-medium-shaped model: L = 24, h = 1024, a = 16, s = 1024, b = 8, fp16. First the per-layer pieces:

sbh = 1024 · 8 · 1024 ≈ 8.39e6
linear term: 34·sbh          ≈ 285 MB per layer
score term:  5·a·s²·b = 5·16·1024²·8 ≈ 671 MB per layer
A_layer ≈ 956 MB;  all 24 layers ≈ 22.9 GB

Nearly 23 GB of activations for a ~350M-parameter model — more than the weights, gradients, and Adam states combined. With full recomputation you keep 24 layer inputs at 2sbh ≈ 16.8 MB each (~0.4 GB) plus one live layer (~0.96 GB): about 1.4 GB peak, a 16× reduction, for ~33% more FLOPs. Notice also which term dominates: at s = h = 1024 the ratio 5as/h = 80 far exceeds 34, so the quadratic score term is 70% of the bill — a fact the next section exploits.

Selective recomputation: drop only the s^2 term

Full recomputation treats all activations alike, but they are not alike. The attention-score tensors are enormous in bytes yet almost free to rebuild: recomputing QK^T, the softmax, and the attention-weighted sum costs only the 4s²bh-ish FLOPs of those ops, a small slice of the layer’s total (the 24sbh² projection and MLP matmuls dominate). Meanwhile the matmul inputs in the linear term are cheap to store but expensive to recompute.

Selective recomputation (Korthikanti et al., 2022) therefore checkpoints straight through the attention-score block and stores everything else. The stored bytes drop from 34sbh + 5as²b to 34sbh per layer. In the worked example that is 6.8 GB instead of 22.9 GB — 70% of the memory win of full recomputation — for measured overheads in the low single digits (~3–5% of step FLOPs) rather than 33%. When one lever removes the only super-linear term for near-zero cost, pull that lever first.

What recomputation does not touch

It is worth being precise about which line of the memory budget this technique edits. Total training memory is roughly:

M_total = weights + gradients + optimizer states + activations
        ≈ 2P + 2P + 12P + M_act    (bytes, fp16 + Adam fp32)

Recomputation shrinks only M_act. The 16 bytes per parameter of state are untouched — those need ZeRO-style sharding, offload, or 8-bit optimizers. This tells you when recomputation is the right tool: it wins when M_act dominates, which happens with large batch, long sequence, or modest parameter count — exactly the fine-tuning regime. In our example, 350M parameters cost ~5.6 GB of states while activations cost 22.9 GB; recomputation is clearly the first move. For a 7B model at b = 1, s = 512, the states (~112 GB) dominate and no amount of activation thrift will save you.

Recompute or offload? The bandwidth arithmetic

The main alternative to rebuilding activations is shipping them to cheaper memory (host RAM, NVMe) and shipping them back for the backward pass. Which is faster is a two-line calculation. Offload cost is bytes over bus bandwidth, both directions; recompute cost is forward FLOPs over achieved throughput:

t_offload ≈ 2 · A_layer / BW_bus
t_recompute ≈ F_fwd / FLOPS_achieved
example layer: F_fwd = 24sbh² + 4s²bh ≈ 2.45e11 FLOPs
t_recompute @ 50 TFLOPS ≈ 4.9 ms
t_offload: 2 · 956 MB / 25 GB/s (PCIe 4) ≈ 76 ms

On an accelerator, recomputation wins by an order of magnitude unless the transfers overlap perfectly with compute. The balance point shifts toward offload as compute gets slower relative to the bus — which is precisely the situation on a CPU, where “offload” often just means “leave it in RAM.”

FlashAttention changes the accounting

Fused attention kernels quietly perform selective recomputation for you. FlashAttention never materializes the [b, a, s, s] score matrix at all: it computes attention in tiles, keeps running softmax statistics (the row max and the log-sum-exp), and in the backward pass recomputes the tiles from Q, K, V plus those statistics. The 5as²b term simply never exists on the memory ledger.

Two accounting consequences follow. First, with FlashAttention enabled, per-layer activation memory is already down to roughly 34sbh, so adding framework-level selective recomputation of the score block is redundant — there is nothing left to drop. Second, the marginal value of full recomputation falls: it now saves ~32sbh per layer instead of 34sbh + 5as²b, while still costing the same one-third FLOPs premium. Always redo the arithmetic after fixing the kernel stack; a memory optimization priced against a naive baseline is often priced wrong.

CPU and SLM implications

On CPU, the trade flips in both directions at once. Compute is scarce — a desktop CPU sustains tens of GFLOPS to a few hundred, not tens of TFLOPS — so the recomputed forward that cost 5 ms on a GPU costs seconds, and a 33% FLOPs premium is a 33% premium on an already slow step. But memory is abundant: 64 GB of RAM comfortably holds the 22.9 GB activation set that overwhelmed a 24 GB GPU. For a small language model fine-tuned on CPU at moderate sequence lengths, the right default is usually no recomputation at all.

The exception is the quadratic term. Push s to 8K on the same model and the score tensors alone reach 5·16·8192²·1 ≈ 5.4 GB per layer at b = 1 — 129 GB across 24 layers. Because CPU training is bandwidth-bound anyway, recomputing that block costs relatively little extra wall-clock while keeping the job inside RAM. On CPU, selective recomputation of the s² term earns its keep; blanket per-layer checkpointing rarely does.

Pitfalls and sharp edges

The math assumes the replayed forward is identical to the original, and that assumption has teeth. Dropout and other RNG ops must replay with the same random mask, or gradients are computed against activations that never existed; frameworks snapshot and restore RNG state inside checkpointed regions, which is correct but adds overhead and is a classic source of silent bugs in custom kernels. In-place operations inside a checkpointed segment can corrupt the saved boundary tensor the replay depends on.

Watch for double recomputation: wrapping a FlashAttention layer in full checkpointing replays a kernel that will internally recompute again, stacking overheads. Measure peak memory, not average — the live segment plus checkpoints sets the peak, and an oversized segment can quietly restore the old high-water mark. And remember non-reentrant vs reentrant implementations differ in how they interact with gradient accumulation and no_grad regions; when a loss curve changes after enabling checkpointing, suspect the replay, not the math.