Selective Activation Recomputation

Checkpointing asks which layers; selective asks which tensors

Classic gradient checkpointing partitions a network into segments, keeps only the segment boundaries live, and replays each segment’s forward pass during backward. The knob is where to cut, the answer is the sqrt(L) rule, and the price is roughly one extra forward pass — about 33% of a training step’s FLOPs.

That framing hides an assumption: that an activation is an activation. It is not. Inside one transformer layer the tensors autograd keeps differ by orders of magnitude in both bytes occupied and arithmetic needed to reproduce them. A LayerNorm output and a down-projection output can be the same shape, [s·b, h], and cost the same memory — while rebuilding one takes a handful of FLOPs per element and the other takes thousands.

Selective recomputation drops the segment abstraction and decides per tensor. That needs one number ranking tensors by how good a deal each is. That number is the whole subject.

The ratio: FLOPs paid per byte freed

Define, for a saved tensor T:

ρ(T) = (FLOPs needed to rebuild T) / (bytes freed by dropping T)

units: FLOPs per byte

Low ρ is a bargain: you hand back a lot of memory for very little arithmetic. High ρ is a bad trade. Because every candidate is measured in the same units, ρ gives a total ordering, and the optimal policy under a memory budget is simply: sort ascending by ρ, drop tensors from the front until you fit, stop.

Two things make this more than bookkeeping. First, ρ has a closed form for the operators a transformer is built from — you compute it symbolically rather than profile it. Second, the values are not spread smoothly; they cluster into two tight groups separated by a factor of a hundred or more. When a ranking has a gap that wide you do not need an optimizer, you need a threshold. (ρ is a close cousin of arithmetic intensity, which counts all traffic in and out of an op; ρ counts only the bytes you stop storing.)

Why ρ is the producing op’s FLOPs per output byte

Rebuilding T means re-running the op that produced it, so ρ is that op’s FLOP count divided by the bytes of its output. Do this for a matmul. A projection X·W with X: [N, k] and W: [k, d] costs 2Nkd FLOPs and emits N·d elements, or 2Nd bytes in bf16:

ρ_matmul = 2Nkd / (2Nd) = k       FLOPs per byte

  QKV / attn-out projection:  k = h    = 768
  FFN up-projection:          k = h    = 768
  FFN down-projection:        k = 4h   = 3072

The N and d cancel: a matmul’s ρ is just its reduction depth. Now do an elementwise op. GELU on [N, 4h] costs on the order of 8 FLOPs per element and frees 2 bytes per element, so ρ ≈ 4. RMSNorm or LayerNorm, two passes plus an affine, lands near ρ ≈ 4 as well. Dropout and softmax sit in the same single-digit neighbourhood.

So 768 versus 4 — a factor of roughly 190. That gap, derived rather than asserted, is the selective-recompute argument.

Ranking the transformer’s tensors

Applying the formula across one pre-norm layer at h = 768, a = 12 heads, head dim d_h = 64:

The interesting row is QK^T. It is a matmul, yet its ρ is only d_h = 64, because attention reduces over the head dimension, not the model dimension — while its output is the layer’s largest tensor, scaling as s². Biggest bytes, shallowest reduction: the best trade in the network, and one no layer-granularity policy can express.

Chains and anchors: you free a span, not a tensor

A subtlety that per-tensor ranking hides: you cannot drop an elementwise output in isolation and expect to rebuild it for 4 FLOPs per byte, because rebuilding it requires its input to still be live. If that input was also dropped, the true cost is the whole chain back to the nearest surviving tensor.

So the real unit of decision is a span between two anchors. Pick tensors to keep (the anchors), and everything strictly between two anchors is free to discard; the cost is the span’s total forward FLOPs and the saving is every intermediate inside it. Evaluate ρ over the span:

anchors: Q, K, V (stored)
span:    S = QK^T/√d_h  →  P = softmax(S)  →  P’ = dropout(P)

cost  ≈ 2bas²d_h  +  ~5bas²  +  ~2bas²      (matmul + softmax + mask)
freed ≈ 3 · 2bas²                             (three s² tensors, bf16)

ρ_span ≈ (2·64 + 7) / 6 ≈ 22 FLOPs per byte

Still an order of magnitude under any projection matmul. Anchoring on Q, K and V is what makes the span cheap: they are small, linear in s, and they cut the chain exactly where the quadratic tensors begin.

Worked example: a 12-layer CPU-trainable SLM

Take L = 12, h = 768, a = 12, s = 2048, b = 1, bf16. Two unit sizes:

linear unit  sbh   = 2048 · 1 · 768   = 1.57M elems  →  3.15 MB
quad   unit  bas²  = 1 · 12 · 2048²  = 50.3M elems  →  100.7 MB

per layer:  ~17 linear-unit tensors  →   53.5 MB
            3 quadratic tensors      →  302.0 MB
whole stack (×12):                      4.27 GB   (642 MB + 3.62 GB)

The s² tensors are 85% of activation memory. Now the bill for the attention span, per layer:

span FLOPs ≈ 2 · 50.3M · 64 + 7 · 50.3M ≈ 6.8 GFLOP
layer fwd  ≈ 24sbh² + 4bas²d_h ≈ 29.0 + 12.9 = 41.9 GFLOP
step (fwd+bwd) ≈ 3 × fwd ≈ 125.7 GFLOP per layer

overhead = 6.8 / 125.7 ≈ 5.4%

3.6 GB removed for a 5% slowdown. Full checkpointing removes the last 642 MB too — and charges 33%.

The frontier: greedy on ρ is steep, then flat

Sort every candidate span by ρ and plot cumulative memory freed against cumulative FLOPs paid. Because you spend the cheapest bytes first, the curve is concave: enormous early returns, then a long flat tail. For the example above, three tiers:

The first two tiers deliver about 94% of the achievable saving for about 5% of a step. The third tier costs five times as much compute as the first two combined and buys a rounding error. That is not a close call, and it is why the correct default is not “checkpoint everything” but “checkpoint down to ρ ≈ 50 and stop.”

Boundaries move with shape: as s grows the quadratic tier swells while its ρ stays pinned at d_h.

Where the published recipes land on the curve

The best-known concrete policy — selective recomputation as described by Korthikanti and colleagues for Megatron-LM — checkpoints straight through the attention block (scores, softmax, dropout) and stores everything else. In the language here, that is exactly “take tier two, skip tier three,” and it is reported at low single-digit overhead on large models. It is a specific point on the frontier, not a law.

Reading it as a point rather than a rule matters, because the point moves. Their configurations have large h and moderate s, so the quadratic tier is a smaller slice of the total and the overhead lands near 3–5%. In the small-h, long-s regime above, the same policy costs more and saves far more. Push the sequence further and the attention tier stops being optional at all. Memorize the ratio, not the recipe: recompute the tiers for your own (h, s, a, b) and the right policy falls out.

What fusion and FlashAttention already recompute for you

Some of this decision was made upstream. A fused kernel computing norm, matmul and activation in one pass never materializes the intermediates, so there is nothing to checkpoint — the tier-one saving is already banked, invisibly, by the kernel author.

FlashAttention goes further and deletes tier two outright. By tiling the attention computation and keeping only the running softmax statistics (m and ℓ, two scalars per row), it never writes the [b, a, s, s] score matrix to memory at all, forward or backward. Its backward pass recomputes the tile-local scores on the fly — which is selective recomputation, implemented inside a kernel instead of by an autograd hook.

The practical consequence: stacking a checkpoint wrapper on a flash-attention layer can be pure loss — you pay the replay and free bytes that were never allocated. Measure the baseline first. If the s² term is already absent from your profile, every remaining candidate is tier one or tier three — a much less interesting menu.

CPU, small models, and the pitfalls

On CPU the arithmetic tilts further toward recomputing. Elementwise and normalization kernels are bandwidth-bound, so replaying them costs closer to a memory sweep than to their nominal FLOP count — cheaper still if the data is already warm in L2 from the backward traversal. Matmul replay, by contrast, competes for the scarce compute that is already your bottleneck. Rank by measured time per byte freed and the two classes separate even more sharply than the table suggests.

Four things that break the accounting. Nondeterminism: dropout and any stochastic op must replay with the identical RNG state or your gradients are silently wrong — frameworks fork and restore the generator for exactly this reason. Nested wrappers: a checkpoint inside a checkpoint replays quadratically. Fragmentation: freed bytes that the allocator cannot coalesce are not usable bytes. Anchor drift: a refactor that moves a stored tensor can silently lengthen a span and multiply its cost.