Learn · Intermediate

Gradient checkpointing: throwing work away so training fits in memory

Gradient checkpointing is a technique that lets you train a model too large to fit in your GPU's memory, by deliberately throwing away most of the intermediate results during the forward pass and recomputing them during the backward pass. It trades roughly 30% more compute for a dramatic reduction in memory: instead of memory growing linearly with the number of layers, it grows with roughly the square root. It is the reason a great many models that "should not fit" on a given GPU nonetheless train on it.

To see why it is needed, you have to understand what backpropagation actually demands. When a network runs forward, each layer takes the previous layer's output, transforms it, and passes it on. To later compute how much each weight should change, the backward pass needs the input that each layer saw. So the standard implementation keeps every single intermediate output in memory from the moment it is computed until the backward pass consumes it -- which, for the first layer, means holding it for the entire duration of the step.

For a 100-layer network with large activations, that stored history dominates memory use. The weights themselves are often the smaller number. This is the counterintuitive fact that trips people up: you can be unable to train a model whose parameters fit comfortably in memory, because the activations do not.

Tianqi Chen and colleagues laid out the fix in Training Deep Nets with Sublinear Memory Cost in 2016. The insight is that intermediate activations are cheap to recreate and expensive to store. They are a deterministic function of the input and the weights, both of which you still have. So instead of storing all of them, store a few -- the checkpoints -- and when the backward pass needs something you discarded, recompute it from the nearest stored checkpoint.

The analogy that makes this click: imagine reading a long novel and needing to answer questions about every chapter afterwards. One approach is to write a detailed summary of every chapter as you read. That is fast to consult and takes an enormous amount of paper. The alternative is to note only where each of the ten major sections begins, and when someone asks about chapter 34, re-read from the start of that section. You do more reading. You carry far less paper.

The arithmetic works out well. If a network has n layers and you place checkpoints every sqrt(n) layers, you store about sqrt(n) checkpoints and, at recomputation time, never need to redo more than about sqrt(n) layers of forward work. Memory goes from O(n) to O(sqrt(n)). Compute goes up by roughly one extra forward pass over the segments being recomputed -- which in practice lands near 30% more time per step, since the backward pass is normally about twice the cost of the forward one.

That exchange rate is usually excellent, because memory is a hard wall and time is a soft one. A step that takes 30% longer is an inconvenience. A step that does not fit is a stop. And the memory you free does not just avoid a crash -- you can spend it on a larger batch size, which often recovers much of the lost throughput and improves gradient quality at the same time.

Gruslys and colleagues generalised the idea in Memory-Efficient Backpropagation Through Time, which uses dynamic programming to find the optimal checkpoint placement for a given memory budget rather than the simple square-root heuristic -- letting you specify how much memory you have and get the fastest schedule that fits.

The same principle shows up in one of the most important systems papers of the modern era. FlashAttention, by Tri Dao and colleagues, avoids ever materialising the full attention matrix -- which grows with the square of sequence length -- by recomputing pieces of it during the backward pass instead of storing it. It is gradient checkpointing applied surgically to the single most memory-hungry operation in a transformer, and it is a large part of why long context windows became practical.

A few things worth knowing before you turn it on. First, the naming collision is genuinely unfortunate: gradient checkpointing has nothing to do with saving model checkpoints to disk, and the two appear in the same configuration files. Second, layers with randomness -- dropout, for instance -- must recompute with the same random values they used originally, or the gradients are wrong. Every serious framework handles this by saving and restoring the random number generator state, but a hand-rolled implementation can get it subtly wrong and produce a model that trains slightly badly rather than obviously badly. Third, checkpoint placement matters: putting them at natural block boundaries, such as transformer layers, is both simpler and usually near-optimal.

It also composes with the other tools in the memory toolkit. Mixed precision training halves activation size. Distributed training parallelism splits activations across devices. Offloading moves data to CPU memory. Gradient checkpointing composes with all three, and in practice serious training runs use several at once.

The deeper lesson generalises past training. When a resource is scarce and a computation is cheap and deterministic, storing the result is a choice, not a requirement. Recomputation is often the better trade -- a principle that shows up again in the KV cache decisions that govern inference memory, where the same question gets asked in the opposite direction.

Ground Truth. (2026, September 3). Gradient checkpointing: throwing work away so training fits in memory. Ground Truth. https://groundtruth.day/learn/gradient-checkpointing.html

@misc{groundtruth:gradient-checkpointing,
  title  = {Gradient checkpointing: throwing work away so training fits in memory},
  author = {{Ground Truth}},
  year   = {2026},
  month  = {sep},
  url    = {https://groundtruth.day/learn/gradient-checkpointing.html}
}

Topics: training · memory · optimization · backpropagation · fundamentals