Gradient checkpointing

1. Intuition

What backward needs. To compute the gradient of a linear layer $y = xW$ with respect to $W$, backward needs the input $x$. To push the gradient through a ReLU it needs to know which inputs were positive. So the forward pass keeps, for every layer, the intermediate tensors that backward will read. Autograd does this for you.

Why it adds up. Those saved tensors scale with tokens times hidden size, for every layer. A 7B-style model with 32 layers at 2,048 tokens keeps about 9 GB of them. At 8,192 tokens it is about 37 GB, more than the bf16 weights.

The trade. Don't keep everything. Keep only the input to each segment of layers (a checkpoint). During backward, when you reach a segment, run its forward again from the saved input to rebuild the intermediates, then do its backward and throw them away. Only one segment's intermediates exist at a time.

The price. Each layer's forward runs twice. A backward pass costs about twice a forward, so a step goes from $1 + 2 = 3$ units to $1 + 1 + 2 = 4$ units: about 33% more compute. The gradients are exactly the same.

Smarter options.

• Selective recomputation keeps the cheap-to-store, expensive-to-recompute tensors and recomputes only the ones that are big but cheap to redo (classically the attention scores). Korthikanti et al. (2022) report it recovers most of the memory for an overhead of a few percent.
• FlashAttention already recomputes attention scores tile by tile in its backward, so the $T \times T$ score matrix is never stored even without checkpointing.
• Offloading copies activations to CPU memory and brings them back for backward, trading PCIe bandwidth for recompute.

2. Visualization

Synthetic 7B-style shape (32 layers, hidden size 4096). Estimates from the rule of thumb in section 3, the same one the LoRA lab uses; real numbers depend on the kernels.

Try it (predict first, then run code/27-training-at-scale/gradient_checkpointing.py and edit it):

1. In the toy network, predict the peak floats for k = 8 (one segment covering all 8 layers) compared with full storage. Then check: why is it slightly higher?
2. Change 34 to 60 in activation_bytes (a model without fused attention kernels). Predict whether k = 1 is still the best segment size.
3. Set Lt = 80 and D = 8192 (a 70B-style shape). Predict the no-checkpointing activation memory at 2,048 tokens before running.
4. Count the recomputed layer forwards for every k. Predict why it is always 8, even though peak memory changes.

3. The math

Symbols

Memory

Without checkpointing every layer keeps its activations:

With a checkpoint every $k$ layers, we keep $\lceil L/k \rceil$ segment inputs of size $b \cdot T \cdot d$, and during backward one segment's full activations, $k \cdot A \cdot T \cdot d$:

The $A = 34$ rule of thumb (Korthikanti et al. 2022, for a standard transformer layer with fused attention, ignoring the $T^2$ term) is rough; it varies with the MLP width and the kernels. The point is the ratio: a layer's input is $b = 2$ bytes per value, its full saved state about 17 times more.

Minimizing over k. If $b$ and $A$ were equal, the two terms would balance at $k = \sqrt{L}$: the classic result of Chen et al. (2016). Here $A$ is 17 times $b$, so the second term dominates and the best $k$ is small. For $L = 32$: $k = 1$ gives $32 \cdot 2 + 1 \cdot 34 = 98$ units of $T d$ bytes, $k = 2$ gives $16 \cdot 2 + 2 \cdot 34 = 100$, $k = 6$ gives $6 \cdot 2 + 6 \cdot 34 = 216$. Frameworks checkpoint every transformer block ($k = 1$) for this reason.

Worked example: 7B-style model, 2,048 tokens

1. $T \cdot d = 2{,}048 \times 4{,}096 = 8{,}388{,}608$.
2. No checkpointing: $32 \times 34 \times 8{,}388{,}608 = 1{,}088 \times 8{,}388{,}608 = 9{,}126{,}805{,}504$ bytes $\approx 9.13$ GB.
3. Checkpointing every layer, stored inputs: $32 \times 2 \times 8{,}388{,}608 = 536{,}870{,}912$ bytes.
4. One layer's full activations during its backward: $34 \times 8{,}388{,}608 = 285{,}212{,}672$ bytes.
5. Total: $536{,}870{,}912 + 285{,}212{,}672 = 822{,}083{,}584$ bytes $\approx 0.822$ GB, about 11 times less.

This 0.822 GB is the activation term in the fine-tuning lesson's 108.63 GB full fine-tuning estimate.

At 8,192 tokens everything is 4 times bigger: $9.13 \times 4 = 36.5$ GB without checkpointing, $0.822 \times 4 = 3.29$ GB with it.

Compute

Per layer, forward costs $F$ and backward about $2F$ (one matrix multiply for the input gradient, one for the weight gradient). With checkpointing, every layer's forward runs once more:

This is the same whatever $k$ is: every layer except possibly the last gets recomputed once. Segment size changes memory, not compute. (Some implementations skip recomputing the final segment, saving $k/L$ of the overhead.)

4. Implementation

Build it from scratch

A residual MLP block stands in for a transformer layer. Its forward returns the output plus the cache backward needs: the input x, the pre-activation h, and the post-ReLU a.

def layer_forward(x, w1, w2):
    """y = x + relu(x W1) W2   (a residual MLP block). Returns y and what backward needs."""
    h = x @ w1
    a = np.maximum(h, 0)
    return x + a @ w2, (x, h, a)

def layer_backward(gy, cache, w1, w2):
    x, h, a = cache
    gw2 = a.T @ gy
    gh = (gy @ w2.T) * (h > 0)
    gw1 = x.T @ gh
    gx = gy + gh @ w1.T          # residual path + MLP path
    return gx, (gw1, gw2)

Predict: which of x, h, a could backward rebuild from the others without re-running a matrix multiply?

The ordinary step keeps every layer's cache alive until backward reaches it:

def backward_full(x, params, gout):
    caches, stored = [], 0
    for w1, w2 in params:
        x, cache = layer_forward(x, w1, w2)
        caches.append(cache)
        stored += sum(t.size for t in cache)       # every intermediate stays alive
    grads = [None] * len(params)
    for i in reversed(range(len(params))):
        gout, grads[i] = layer_backward(gout, caches[i], *params[i])
    return grads, stored, 0

Predict: doubling the batch doubles stored. What does doubling the number of layers do?

The checkpointed step keeps only the input of every k-th layer. In backward, it walks segments from last to first, re-runs each segment's forward from its saved input, and back-propagates through it.

def backward_checkpointed(x, params, gout, k):
    L = len(params)
    boundaries = {}                                   # segment start index -> its input
    for i, (w1, w2) in enumerate(params):
        if i % k == 0:
            boundaries[i] = x                          # the only tensors kept from the forward pass
        x, _ = layer_forward(x, w1, w2)                # cache thrown away
    grads, recomputed, peak = [None] * L, 0, 0
    for start in reversed(range(0, L, k)):
        end = min(start + k, L)
        h, caches = boundaries[start], []
        for i in range(start, end):                    # re-run this segment's forward
            h, cache = layer_forward(h, *params[i])
            caches.append(cache)
            recomputed += 1
        live = sum(b.size for b in boundaries.values()) + sum(t.size for c in caches for t in c)
        peak = max(peak, live)
        for j in reversed(range(start, end)):
            gout, grads[j] = layer_backward(gout, caches[j - start], *params[j])
    return grads, peak, recomputed

Predict: the recomputed forward uses exactly the same weights and inputs. Will the gradients match the full version to the last bit, or only approximately?

Use the library

In PyTorch, wrap each block's forward in torch.utils.checkpoint.checkpoint:

from torch.utils.checkpoint import checkpoint

class Model(nn.Module):
    def forward(self, x):
        for block in self.blocks:
            x = checkpoint(block, x, use_reentrant=False) if self.training else block(x)
        return x

Hugging Face models expose the same thing as model.gradient_checkpointing_enable() (or gradient_checkpointing=True in TrainingArguments). Megatron-Core offers full and selective recomputation, and PyTorch's FSDP2 and DeepSpeed can offload activations to CPU instead. use_reentrant=False is the recommended implementation; it supports keyword arguments and inputs that don't require gradients.

Watch out for randomness: dropout in the recomputed forward must use the same random mask. PyTorch's checkpoint saves and restores the RNG state by default (preserve_rng_state=True).

Run it

uv run code/27-training-at-scale/gradient_checkpointing.py

Key lines:

  full storage: 4,608 floats kept, 0 layers recomputed
  k = 1: peak 1,088 floats, 8 layer forwards recomputed, max grad diff 0.0e+00
  k = 2: peak 1,408 floats, 8 layer forwards recomputed, max grad diff 0.0e+00
  k = 8: peak 4,672 floats, 8 layer forwards recomputed, max grad diff 0.0e+00
  no checkpointing: 1 + 2 = 3 units; checkpointing: 1 + 1 + 2 = 4 units -> +33.3% per step
  no checkpointing: 32 x 34 x 2048 x 4096 = 9.127 GB
  checkpoint every  1 layers:  0.822 GB
  checkpoint every  8 layers:  2.349 GB
  at 8,192 tokens: none 36.51 GB, k = 1 3.29 GB

The gradient difference is exactly zero: recomputation is deterministic. With k = 8 (one segment) the peak is slightly above full storage because the segment's input is stored on top of all the caches: checkpointing everything into one segment is all cost and no saving.

5. Engineering

When to turn it on. When activations, not weights or optimizer state, are what doesn't fit: long sequences, large micro-batches, or after you have already sharded the weights and optimizer (see ZeRO and FSDP). Most fine-tuning recipes turn it on by default because sequence length is the usual limit.

Spend the memory. Checkpointing frees memory; use it for a bigger micro-batch or a longer context. A bigger micro-batch can win back part of the 33%, because larger matrix multiplies run the GPU more efficiently.

Selective first. If your framework supports selective recomputation, try it before full checkpointing: most of the memory saving for a few percent of compute. With FlashAttention the attention scores are already not stored, so the remaining activations are mostly linear-layer inputs and MLP intermediates.

Offload when bandwidth is idle. Activation offload to CPU costs no extra FLOPs but needs PCIe transfers to overlap with compute. It helps most when the GPU is memory-bound but the host link is underused.

Interactions. With pipeline parallelism, each stage stores activations for every micro-batch in flight, so checkpointing per stage is common. With tensor parallelism, recompute repeats the layer's communication too. With FP8 training, recomputation must reuse the same scaling factors as the original forward.

Common mistakes

• Expecting checkpointing to reduce weight, gradient, or optimizer memory: it only touches activations.
• Checkpointing the whole model as one segment: you store the input plus everything else during backward, so you save nothing.
• Forgetting RNG state: if dropout masks differ between the forward and the recomputation, gradients are wrong. Keep preserve_rng_state=True.
• Leaving checkpointing on at inference or evaluation: there is no backward pass, so it only costs time. Guard it with self.training.
• Assuming the cost is free: plan for about a third more compute per step with full recomputation.

6. Knowledge check

Practice exercise. A synthetic 13B-style model has $L = 40$ layers and $d = 5{,}120$. Using $A = 34$ and $b = 2$, compute activation memory at 4,096 tokens with no checkpointing and with checkpointing every layer.

Implementation challenge. Add an offload variant to the script: in the forward pass, put each layer's cache in a separate list called cpu (standing in for host memory) and count the floats "transferred". In backward, take caches back from cpu. Report peak GPU floats and floats transferred, and compare with k = 1 checkpointing.

def backward_offload(x, params, gout):
    cpu, moved = [], 0
    for w1, w2 in params:
        x, cache = layer_forward(x, w1, w2)
        cpu.append(cache)                              # "copy to host"
        moved += sum(t.size for t in cache)
    grads = [None] * len(params)
    for i in reversed(range(len(params))):
        cache = cpu.pop()                              # "copy back"
        moved += sum(t.size for t in cache)
        gout, grads[i] = layer_backward(gout, cache, *params[i])
    return grads, moved

Summary

• Backward needs each layer's saved intermediates, so activation memory scales with layers × tokens × hidden size: about 9.13 GB for a 7B-style model at 2,048 tokens.
• Checkpointing keeps only segment inputs and recomputes each segment's forward during backward; gradients are identical.
• Checkpointing every layer brings that 9.13 GB down to 0.822 GB, for about 33% more compute per step whatever the segment size.
• The $\sqrt{L}$ rule assumes stored inputs cost as much as full activations; in transformers they are about 17 times cheaper, so per-block checkpointing wins.
• Selective recomputation, FlashAttention, and offloading get most of the saving for less compute.

Next: Collective operations from scratch

Related: Backpropagation · Fine-tuning in practice · ZeRO and FSDP

Go deeper, code-first: AI Engineering from Scratch, gradient checkpointing.

Parts of this lesson are adapted from AI Engineering from Scratch by Rohit Ghumare (MIT license, notice): Gradient Checkpointing and Activation Recomputation.

Interview angle