Gradient checkpointing is a training technique that saves memory by not storing every intermediate activation during the forward pass, and instead recomputing the missing ones during backpropagation. You trade extra compute for a smaller memory footprint. In practice, it lets you train models or use batch sizes that would otherwise exceed GPU memory, at the cost of roughly one additional forward pass over the checkpointed part of the network.
Why Activation Memory Grows During Training
Neural network training needs the intermediate outputs of each layer, called activations, to compute gradients in the backward pass. Backpropagation applies the chain rule: to get the gradient for a layer's weights, you multiply the gradient flowing back from the next layer by the derivative of that layer's function evaluated at its input. That input, the stored activation, must still be in memory when the backward pass reaches the layer. So during the forward pass, a typical training loop keeps every activation alive until its gradient has been computed, which means memory scales with depth, batch size, sequence length, and hidden dimension all at once.
For a transformer, this cost is dominated by attention and feed-forward blocks, and it grows quickly with context length: each layer holds tensors proportional to batch size times sequence length times hidden size, and the attention mechanism adds score matrices whose size scales with the square of the sequence length in the unfused case. Parameters themselves are usually a smaller fixed cost by comparison. When you hit an out-of-memory error on a large model, activations, not weights, are often the culprit.
How Gradient Checkpointing Works: Store Less, Recompute More
The core idea is selective storage. Instead of saving every activation, you divide the model into segments and store only the activations at segment boundaries, the "checkpoints." During the backward pass, when the gradient reaches a segment, you rerun the forward pass for just that segment from its stored boundary activation, producing the intermediate values on demand, use them to compute gradients, then discard them. Peak memory drops because only one segment's worth of intermediates exists at a time, plus the boundary activations.
The compute cost comes from those recomputations. If you checkpoint a fraction of the network, the backward pass costs one extra forward pass over the checkpointed portion. A common estimate: a full forward-backward step is roughly three forward-pass-equivalents in compute (one forward, two for backward), so recomputing everything adds about one more, making the step roughly a third more expensive, or around 1.33x. That figure depends on the model and hardware; treat it as an order-of-magnitude estimate, not a precise benchmark. Because it is a training-only technique, it changes nothing about the trained model or inference.
Choosing Checkpoint Segments and Measuring the Trade-off
Where you place checkpoints controls the memory-compute balance. Checkpointing every layer maximizes savings but maximizes recomputation. Checkpointing only large blocks, such as whole transformer layers, is usually the practical choice: it captures most of the memory win while keeping recompute overhead modest. The granularity matters because storing boundary activations also costs memory, so too many tiny segments adds storage without much benefit.
Measure rather than guess. Compare peak GPU memory and step time with checkpointing off, on for selected blocks, and on everywhere, using a fixed batch size. Typical questions to answer: does the memory freed let you increase batch size enough to offset the slower step, and does throughput per example stay acceptable? If memory is not your bottleneck, checkpointing only adds overhead. If it is, the trade is usually worth it.
When to Use It: Transformers, Long Context, and Limited GPUs
Gradient checkpointing pays off when activations dominate memory: deep transformers, long-context training where attention and sequence tensors are large, and any setup where you are memory-bound on a single GPU or a small pool. It is less useful for small models or short sequences where weights and optimizer state are the main cost, and it does nothing for inference. A sensible default is to enable it for the largest blocks first, confirm the memory savings, and expand coverage only if you still need headroom.
FAQ
What is gradient checkpointing in simple terms? During training, the model normally keeps every intermediate result so it can use them later to compute gradients. Gradient checkpointing throws most of them away and saves only a few, then recalculates the thrown-away values when needed. You use less memory but do more math.
How much memory does gradient checkpointing actually save? It depends on how much you checkpoint. Checkpointing every layer can cut activation memory dramatically, often by a large factor, because only boundary activations and one segment's intermediates are live at once. Checkpointing only some blocks saves proportionally less. The exact number depends on model architecture, sequence length, and batch size, so measure peak memory on your own setup.
Does gradient checkpointing slow down training, and by how much? Yes, because recomputation adds work. If you checkpoint the whole network, a reasonable estimate is roughly a 20 to 40 percent increase in step time, corresponding to about one extra forward pass on top of the normal forward-backward cost. Partial checkpointing costs less. Actual overhead varies with hardware and implementation.
Can I combine gradient checkpointing with mixed precision or LoRA? Yes. It is independent of precision: recomputed activations follow the same forward path, so mixed precision works alongside it. It also composes with parameter-efficient methods like LoRA, which reduce the number of trainable weights. Combining them is common when memory is tight, but each adds its own overhead, so measure the combined effect rather than assuming the savings simply add up.
Where does gradient checkpointing fit into the backpropagation stage? It changes what happens between the forward and backward passes. The forward pass stores only checkpoint boundaries; the backward pass, as it reaches each segment, reruns that segment's forward computation to rebuild the intermediates, then computes and propagates gradients through it before freeing the rebuilt tensors. The gradient math is unchanged; only when activations are materialized changes.
Choosing Checkpoint Segments and Measuring the Trade-off
A checkpoint segment is the span of layers between two points where you keep activations. Its size controls the trade. Checkpoint every layer and you store almost nothing but recompute almost everything; checkpoint every tenth layer and you save modestly.
The mechanics: during the forward pass you store only the activations at segment boundaries, and discard the intermediate tensors inside each segment as you go. During backpropagation, when the gradient needs those intermediates, you rerun the forward pass for that segment from its stored boundary, then compute gradients. You pay roughly one extra forward pass per segment, but only when that segment is reached.
Worked example. Assume a 24-layer model where each layer's stored activations cost 100 MB, so full storage is 2,400 MB. Suppose you checkpoint every 4 layers, giving 6 segments. You store 6 boundaries at 100 MB each: 600 MB. During backward, each segment is recomputed once, adding the equivalent of one forward pass over the whole model (all 24 layers are recomputed exactly once across the 6 segments). So activation memory falls from 2,400 MB to 600 MB, a 75% reduction, at the cost of one extra forward pass. Since a training step is typically one forward plus one backward (backward usually costing about twice a forward), the step goes from about 3 forward-equivalents to 4, roughly a 33% increase. That is an estimate, not a benchmark; the real ratio depends on your model and hardware.
Formula for the trade:
segments = ceil(layers / segment_size)
stored = segments * per_layer_activation_mb
saved = 1 - stored / (layers * per_layer_activation_mb)
overhead ≈ extra_forward / (forward + backward)
Checklist before enabling it:
- Confirm you are actually activation-memory bound, not weight- or optimizer-memory bound.
- Pick the largest segment size that still fits your target batch.
- Verify batch size can grow enough to justify the slowdown.
- Re-run a short training slice and compare throughput and loss curve.
- Keep the same segment boundaries across steps for consistency.
When to Use It: Transformers, Long Context, and Limited GPUs
Transformers are the classic case. Attention activations scale with sequence length times batch, so long-context training fills memory fast, and checkpointing the attention and feed-forward blocks is usually the first lever. It pairs naturally with long context: doubling sequence length roughly doubles activation memory per layer, so shrinking the stored fraction matters more, not less.
Use it when a bigger batch or longer sequence would otherwise not fit, and when the extra compute is cheaper than the alternative (smaller batch, gradient accumulation, or more GPUs). Skip it when you have memory headroom and compute is your bottleneck, since it only adds work.
Combine it with mixed precision and parameter-efficient methods like LoRA (low-rank adaptation, which trains small added matrices instead of all weights). These reduce weight and optimizer memory, not activation memory, so they stack well with checkpointing. Note that checkpointing interacts with precision: recomputed activations must match the original forward pass's numerics, so keep the same precision settings on both passes.
A practical rule: enable checkpointing on the largest submodules (attention blocks, feed-forward layers) rather than individual operations, and tune segment size until the memory saving stops being worth the throughput cost.
Frequently Asked Questions
What is gradient checkpointing in simple terms?
Gradient checkpointing is a memory-saving trick for training neural networks. During the forward pass, instead of keeping every intermediate activation (the tensors each layer produces) in memory for later use, it saves only a few of them at chosen points called checkpoints. When the backward pass needs an activation that wasn't stored, it recomputes it on the fly from the nearest checkpoint. You trade extra computation for lower memory use.
How much memory does gradient checkpointing actually save?
The activation memory for a checkpointed segment drops from roughly proportional to the number of layers to proportional to the number of checkpoints plus the cost of recomputing one segment. If you checkpoint every layer individually (the most aggressive common setting), activation storage approaches the cost of a single layer's activations rather than all of them. A reasonable estimate is that activation memory falls by a factor on the order of the segment length you checkpoint over, so checkpointing every k layers saves roughly a factor of k in that portion of memory, minus the checkpointed boundaries themselves. Model weights, gradients, and optimizer state are not reduced.
Does gradient checkpointing slow down training, and by how much?
Yes. Each recomputed segment adds one extra forward pass through those layers during backpropagation, so the added cost scales with how much you recompute. In the common case of checkpointing every layer, the extra work is roughly one additional forward pass over the checkpointed part of the network, which typically translates to something like a 20–35% increase in step time as an estimate, depending on how forward-heavy the model is relative to backward. Checkpointing larger segments recomputes more per segment and tends to cost more; checkpointing fewer layers costs less but saves less memory.
Can I combine gradient checkpointing with mixed precision or LoRA?
Yes, and they are frequently used together. Mixed precision (running most computation in 16-bit floats to cut memory and speed up math) reduces the size of stored activations, so checkpointing compounds the savings, though recomputation must respect the same precision settings as the original forward pass. LoRA (Low-Rank Adaptation) freezes most weights and trains small added matrices, which shrinks optimizer memory but not activation memory, so checkpointing still helps when activations dominate. Just verify that your framework's checkpointing implementation supports the precision and adapter setup you're using.
Where does gradient checkpointing fit into the backpropagation stage?
It lives inside the backward pass. The forward pass runs normally but discards most intermediate activations, keeping only the checkpoints. When backpropagation reaches a checkpointed segment, it reruns the forward pass for that segment starting from its checkpoint to rebuild the needed activations, then computes gradients through it. This means the recomputation is triggered during the backward sweep, one segment at a time, and the freshly rebuilt activations can be freed again as soon as that segment's gradients are done.