Gradient checkpointing

Setup#

pip install torch

The first demo runs anywhere, CPU included. The memory measurement needs a GPU; those numbers were captured with torch 2.5.1 on an NVIDIA RTX A6000.

Watch the recomputation happen#

The cleanest proof that forward runs twice — a counter:

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint

calls = {"count": 0}

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = nn.Linear(64, 64)
    def forward(self, x):
        calls["count"] += 1
        return torch.relu(self.layer(x))

torch.manual_seed(0)
block = Block()
x = torch.randn(8, 64, requires_grad=True)

calls["count"] = 0
block(x).sum().backward()
print("normal:        forward ran", calls["count"], "time(s)")

calls["count"] = 0
checkpoint(block, x, use_reentrant=False).sum().backward()
print("checkpointed:  forward ran", calls["count"], "time(s)")

normal:        forward ran 1 time(s)
checkpointed:  forward ran 2 time(s)

checkpoint(block, x) runs the block without storing its internal activations. When backward() reaches that region, the block runs forward again — the counter catches it red-handed — and this time the activations are kept, used, and dropped.

use_reentrant=False selects the modern implementation. Pass it explicitly: the old default has sharp edges around unused inputs and RNG state, and newer PyTorch versions warn until you choose.

Measure the saving#

For a stack of layers, checkpoint_sequential splits it into segments and stores only segment boundaries:

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint_sequential

if not torch.cuda.is_available():
    raise SystemExit("needs a GPU to measure memory")

def peak_mb(use_ckpt):
    torch.manual_seed(0)
    layers = nn.Sequential(*[
        nn.Sequential(nn.Linear(1024, 1024), nn.ReLU()) for _ in range(16)
    ]).cuda()
    x = torch.randn(4096, 1024, device="cuda", requires_grad=True)
    torch.cuda.reset_peak_memory_stats()
    if use_ckpt:
        out = checkpoint_sequential(layers, segments=4, input=x, use_reentrant=False)
    else:
        out = layers(x)
    out.sum().backward()
    return torch.cuda.max_memory_allocated() / 1024**2

print(f"without checkpointing: {peak_mb(False):7.1f} MB peak")
print(f"with 4 segments:       {peak_mb(True):7.1f} MB peak")

without checkpointing:   420.3 MB peak
with 4 segments:         276.4 MB peak

A third of the peak gone, on a model this shallow. The saving grows with depth: the deeper the network, the larger the share of memory that is skippable in-between results. Real transformer training checkpoints each block and commonly cuts activation memory several-fold.

Common mistakes#

Checkpointing modules with randomness or state, carelessly. Dropout runs twice; the modern implementation replays RNG so the two runs match, but custom modules with side effects (counters, caches, prints) will visibly execute twice. BatchNorm updates running statistics on each call — a checkpointed BatchNorm block updates them twice per step in older reentrant mode. Know what is inside the block you wrap.

Wrapping the whole model in one checkpoint. Then nothing in between is stored, and backward recomputes the entire forward — maximum time cost, and the peak is dominated by the single giant recomputed segment. Several mid-sized segments beat one huge one.

Expecting a speedup. This lesson's tool spends time to buy memory. Steps get roughly 20–35% slower. The win is indirect: the memory freed may allow a batch size with better GPU utilisation.

Passing inputs that do not require grad. With nothing to differentiate, older reentrant checkpointing silently skips gradient tracking for the segment. The use_reentrant=False path handles it correctly — one more reason to always pass the flag.

Try it yourself#

In ckpt_mem.py, sweep segments over 2, 4, 8 and 16, printing the peak each time. Find the sweet spot, and notice it is not 16 — boundary storage starts to dominate.

What to learn next#

• torch.compile — the compute-side counterpart: same model, faster kernels.
• Fast attention with scaled_dot_product_attention — kills the biggest recomputable tensor outright.
• FSDP: sharding a model across GPUs — when activations are fine but parameters are not.