Activation checkpointing & recomputation

The third memory-saving axis. FSDP shards the weights; this shards the activations — not across ranks, but across time. Throw most of them away on the forward pass and pay to recompute them on the backward. A constant-factor compute tax buys a √L reduction in activation memory.

The fourth thing in HBM

Lesson 01 counted three consumers of HBM: parameters, gradients, optimizer state. There is a fourth, and at long sequence lengths it is the largest of all — the activations. Every operation in the forward pass produces a tensor that the backward pass needs in order to apply the chain rule. The matmul Y = XW needs X to compute ∂L/∂W = Xᵀ · ∂L/∂Y. So X cannot be freed when forward finishes with it — it has to survive until backward reaches that layer.

For a transformer the per-layer activation memory scales as b · s · h (batch × sequence × hidden), times a constant for the handful of tensors saved inside each block (the attention scores are b · s · s — quadratic in sequence). Multiply by the number of layers L and you get the total:

M_act ≈ L · b · s · (c₁·h + c₂·s)

The thing to notice: parameters, gradients and optimizer state do not depend on batch or sequence. Activations depend on both. Double the context and the weights cost the same while the activations double (and the attention term quadruples). This is why a model that "fits" at 2k context OOMs at 32k — nothing about the weights changed; the activation pile grew.

The idea, in one sentence

You do not have to store an activation to have it at backward time — you can recompute it by re-running the forward from the nearest thing you did keep. Storage and recomputation are interchangeable. So: keep a sparse set of activations (the checkpoints), discard everything between them, and when backward needs a discarded activation, regenerate it by replaying forward from the last checkpoint.

The name "checkpoint" is overloaded in this series. Here it means an activation deliberately kept at a segment boundary so a forward replay has a starting point. Lesson 26's "checkpoint" is the whole-model snapshot written to disk. Same word, unrelated mechanisms.

  WITHOUT checkpointing — store every layer's input, peak = L
    fwd:  x0 ─▶[L0]─▶ x1 ─▶[L1]─▶ x2 ─▶[L2]─▶ x3 ─▶ ... ─▶ xL   (all xᵢ kept)
    bwd:  every xᵢ already in HBM ────────────────────────────── no recompute

  WITH checkpointing, segment size g=2 — keep only x0, x2, x4, ...
    fwd:  x0 ▣ ─▶[L0]─▶(x1 freed)─▶[L1]─▶ x2 ▣ ─▶[L2]─▶(x3 freed)─▶[L3]─▶ x4 ▣
    bwd:  reached L1, need x1 → REPLAY [L0] from x2's segment start → get x1 → backward
          ▣ = activation kept (a "checkpoint")      (xᵢ freed) = discarded, will be recomputed

The √L result — why segmenting wins

Let one layer's stored activation be one "unit". Divide the L layers into segments of length g (a segment size, distinct from the sequence length s in M_act above), so there are L/g segments. Two things sit in HBM at once during backward:

• The checkpoints — one kept activation per segment boundary: L/g units, held for the whole backward.
• One segment's recompute — to backward through a segment you first replay its forward, which transiently materialises up to g activations: g units.

M_peak(g) = L/g + g

Minimise: dM/dg = −L/g² + 1 = 0 ⇒ g = √L, giving M_peak = 2√L. A 64-layer model drops from 64 units of activation memory to 2√64 = 16 — a 4× cut. A 1024-layer model drops 16×. The deeper the model, the bigger the win, because L/√L = √L grows.

And the compute cost? Every discarded activation is recomputed exactly once, during its segment's backward replay. That's one extra forward pass over the whole network, independent of g. A training step is normally one forward + one backward, and backward costs ≈ 2× forward (it computes two gradients per op — w.r.t. input and w.r.t. weight). So a step is ≈ 3F; adding one recompute forward makes it 4F — a flat ~33% compute tax, no matter how you choose g.

Interactive · the memory curve and its √L minimum

Set the depth L and slide the segment size g from 1 (checkpoint every layer — minimal recompute peak, maximal checkpoint count) to L (checkpoint nothing but the entrance — one giant recompute). The blue curve is peak activation memory L/g + g; the dashed line is the no-checkpointing cost L. Watch the curve bottom out at g = √L.

Animated · one segment's store-then-recompute, scrubbable

An 8-layer model, segment size 2 (4 checkpoints). Scrub the timeline. On the forward pass, only segment-boundary activations are kept (solid) — the rest are computed and immediately freed (hollow). On the backward pass, each segment is replayed forward from its checkpoint (the recompute burst), regenerating its internal activations just long enough to run backward through them, then freed again. The bottom strip is the HBM curve: a low baseline of checkpoints with a sawtooth of recompute spikes — compare it to the flat-high line of no-checkpointing.

Selective recomputation — the modern refinement

Full checkpointing recomputes everything in a segment, paying 33% to save all activations. But not all activations are equal. Inside a transformer block, two kinds of tensor compete for that memory:

Selective activation recomputation (Korthikanti et al., Megatron-LM, 2022) keeps the expensive-to-recompute matmul activations and drops only the cheap-to-recompute, memory-heavy attention activations. The result: most of the memory saving of full checkpointing for only ~2–5% extra compute instead of 33%. At long context — where the quadratic b·s·s term dominates — this is close to free memory.

How it composes with the rest of the stack

• + FSDP (lesson 05): the FSDP unit and the checkpoint segment are usually the same transformer block. On backward, FSDP AllGathers the block's weights and checkpointing replays the block's forward — both happen together, just-in-time, then both are freed. This is the standard Llama-class recipe.
• + pipeline parallel (lesson 07): 1F1B keeps several micro-batches' activations alive at once (that's the memory the bubble "costs"). Checkpointing each stage's activations is what makes deep pipelines fit — it shrinks the per-micro-batch footprint that PP multiplies.
• + tensor / sequence parallel (lessons 06, 08): SP exists precisely to shard the activations TP leaves replicated. Checkpointing and SP both target activations and are often used together at long context.

The decision, in order

1. Activations fit → don't checkpoint. The 33% is pure loss.
2. Long context blowing the budget → turn on selective recomputation first; it's nearly free and kills the quadratic term.
3. Still over budget → full checkpointing per transformer block: keep each block's input activation, drop the dozen-odd activations inside the block, recompute them on backward. This is the common default — its win is killing the within-block pile, not hitting the √L optimum (one checkpoint per block is the g=1 end of the curve above).
4. Want the absolute minimum activation memory → tune segment size toward √L, and compose with FSDP + SP.
5. Out of compute budget, not memory → you've over-checkpointed; trade some back for a bigger micro-batch (lesson 11) instead.