Activation Checkpointing

Here is the entire idea in one sentence: don't store what you can cheaply recompute.

That sentence is roughly what most treatments of activation checkpointing give you, and it's true, and it's useless. It doesn't tell you what gets stored, what "cheaply" costs, where the boundaries go, why the naive version wastes most of its own benefit, or how it interacts with dropout in a way that silently corrupts your training run if you implement it yourself. So let's do it properly.

Why anything is stored at all#

Start with why the forward pass leaves a mess behind, because the answer determines exactly what you're allowed to throw away.

Take a single linear layer, $Y = XW$. The backward pass needs two gradients:

Look at the second one. To compute the weight gradient you need $X$ — the layer's input, from the forward pass, which happened a long time ago. So autograd holds onto it. Every layer does this, for every tensor its backward formula references, and the sum of all those retained tensors is the activation memory from chapter 2.

Which tensor gets retained depends on the operation's backward formula, and the differences are instructive:

Two things follow. First, some ops are already frugal — ReLU can get away with one bit per element, which is why it was a memory-efficiency win over GELU before anyone was framing it that way. Second, an op that retains its own output is free to recompute if you kept the next layer's input anyway. That observation is what makes selective recomputation work.

The naive version, and its cost#

Wrap a transformer block in torch.utils.checkpoint, and it does exactly two things: run the forward pass under torch.no_grad() so nothing is retained, and register a backward hook that re-runs the forward pass — this time with grad enabled — to regenerate what's needed, right before it's used.

Here it is written out, because the real implementation is buried in abstraction and the mechanism is genuinely simple:

class Checkpoint(torch.autograd.Function):
    @staticmethod
    def forward(ctx, fn, *args):
        ctx.fn = fn
        ctx.save_for_backward(*args)
        # The whole trick: compute the output, retain nothing in between.
        with torch.no_grad():
            return fn(*args)
 
    @staticmethod
    def backward(ctx, *grad_outputs):
        inputs = [x.detach().requires_grad_(True) for x in ctx.saved_tensors]
        # Second forward pass — this time we keep the tape.
        with torch.enable_grad():
            outputs = ctx.fn(*inputs)
        torch.autograd.backward(outputs, grad_outputs)
        return (None,) + tuple(x.grad for x in inputs)

Applied to every transformer block, this takes per-layer activation memory from $sbh(34 + 5as/h)$ down to $2sbh$ — you keep each block's input and nothing else. For Llama 3 8B at 8k context that's 354 GiB → 2 GiB, a 177× reduction.

The cost is one extra forward pass. Against the standard $6N$ FLOPs per token of training ($2N$ forward, $4N$ backward), an extra $2N$ is:

A third of your training throughput, gone. Which — read against chapter 1 — is a trade you should be suspicious of, because 33% is a lot to pay when the last 5% of that memory saving is doing almost all of the work.

Sublinear memory: the $\sqrt{L}$ result#

Before the transformer-specific refinement, the classical result, because it explains why "checkpoint every layer" isn't obviously the right granularity.

You have $L$ layers. Checkpoint every $k$-th one. Then:

• Stored checkpoints: $L/k$ boundary activations.
• Recompute peak: reconstructing any segment needs the intermediates within it, so $k$ layers' worth, transiently.

Total memory is $O(L/k + k)$, minimized at $k = \sqrt{L}$, giving $O(\sqrt{L})$ memory for one extra forward pass.Chen et al., Training Deep Nets with Sublinear Memory Cost, 2016. Predates transformers entirely — it's a statement about backprop through any deep feed-forward graph.

For $L = 32$ that's checkpointing roughly every 6 layers rather than every one. In practice almost nobody does this, and the reason is worth understanding: with pipeline parallelism you only hold $L/d$ layers per device anyway, and the per-layer boundary is the natural unit for the recompute-vs- communicate tradeoff. The $\sqrt{L}$ schedule optimizes a constraint that sharding already relaxed. It's the right answer to the 2016 question.

Selective recomputation: where the actual win is#

Now the part that matters, and the part the high-level treatments skip.

Go back to the two terms:

At 8k context on Llama 3 8B, the quadratic term is 90% of activation memory. And here's the thing: the operations that produce it — $QK^\top$, softmax, the dropout mask — are also cheap relative to the block's matmuls, because at 3.2% of block FLOPs (we computed this in chapter 1) attention is a small slice of the compute.

So you have a region that is 90% of the memory and 3% of the compute. Recompute only that, and keep everything else.

That's selective recomputation.Korthikanti et al., MLSys 2023, §4. They report it as ~5× activation memory reduction for 2.7% overhead on a 22B model — the exact ratio depends on $s$, $a$, and $h$. It captures most of the benefit for a small fraction of the cost:

Selective recompute buys memory at roughly eleven times the efficiency of full recompute. If you take one operational thing from this chapter: selective first, full only if selective isn't enough.

Toggle between the three strategies here, with flash attention on and off, and watch which combinations are redundant:

Two ways to get this wrong#

Dropout and RNG state. Recomputation must reproduce the forward pass exactly. If the block contains dropout, the second forward pass draws fresh random numbers, uses a different mask, and computes gradients for a network that never ran. Training doesn't crash. It degrades, subtly, and you spend two weeks blaming your learning rate schedule.

PyTorch's checkpoint handles this by default — preserve_rng_state=True stashes the CUDA RNG state on the way in and restores it before recomputing. It isn't free (a device sync), and it's the first thing to check if you ever hand-roll this. Modern LLM training mostly sidesteps the issue by not using dropout at all during pretraining, but fine-tuning stacks often do.

Non-reentrant vs reentrant. The classic implementation (use_reentrant=True) re-enters autograd via a nested backward() call, and breaks in ways that surprise people: it can't handle blocks with no requires_grad inputs, it doesn't compose with torch.autograd.grad(), and it mishandles some hook patterns. The newer non-reentrant implementation uses saved-tensor hooks instead and is strictly better. It's now the default, but plenty of code in the wild still passes use_reentrant=True explicitly, and plenty of tutorials still recommend it.

from torch.utils.checkpoint import checkpoint
 
# Old code you will encounter, and should not copy:
h = checkpoint(block, h, use_reentrant=True)
 
# What you want:
h = checkpoint(block, h, use_reentrant=False)

MFU, HFU, and honest accounting#

This is where the two utilization metrics from chapter 1 earn their keep.

Recomputation adds FLOPs the model never asked for. So:

• MFU counts only the $6N$ per token the math requires. Recomputation makes MFU go down, because you're taking longer to do the same modelling work.
• HFU counts every FLOP the hardware executed, recomputation included. Recomputation makes HFU go up.

A run with full recomputation should show HFU ≈ 1.33 × MFU. If someone reports an impressive utilization number while checkpointing aggressively, check which one they're quoting — HFU flatters recomputation-heavy configurations, and the gap between the two is precisely the tax.

The deeper point: recomputation is only worth it if the memory it frees lets you do something that recovers more than 33% throughput. Usually it does — a bigger micro-batch means larger GEMMs and better arithmetic intensity, and often it's the only way to fit the model at all. But it's a trade to make deliberately, not a default to switch on.

In practice the decision usually gets made in the opposite order from how it's taught. You don't start from "how much recompute do I want" — you start from a target global batch size the run needs for convergence, discover it doesn't fit, and add exactly as much recomputation as it takes to make it fit and not a scrap more. Selective is the first lever because it's nearly free; full comes out only when selective plus a smaller micro-batch still OOMs. The failure mode to watch for is reaching for full recompute reflexively "to be safe" — it's a 33% tax that a moment of memory arithmetic would have shown you didn't need to pay.

What to actually do#

1. Use flash attention. It removes the quadratic term at zero FLOP cost and makes the selective-recompute question mostly moot.
2. If you still don't fit, add full recomputation on the blocks, and expect to pay ~33%.
3. Prefer more micro-batch over less recompute. Larger batches improve arithmetic intensity, and a compute-bound step at 33% overhead often beats a memory-bound step at 0%.
4. Consider per-layer granularity. Frameworks let you checkpoint only the first $k$ layers — enough to fit, no more. Pipeline parallelism makes this especially useful, since the first stage holds the most in-flight micro-batches (chapter 9).