Gradient Checkpointing — AI Glossary

Reference details

Related terms

• Mixed Precision Training
• DeepSpeed
• FSDP (Fully Sharded Data Parallel)

How it works

1. During the forward pass the model is divided into segments, and only the activation at each segment boundary is retained while everything inside is discarded as it is produced.
2. The backward pass then reaches a segment and has no activations for it, so the segment's forward computation is re-run from the stored boundary to regenerate them.
3. Those regenerated activations are used for that segment's gradients and discarded again immediately, so peak memory holds one segment's worth rather than the whole model's.
4. Implementations achieve this by detaching the segment from the autograd graph during the first forward pass and re-enabling gradient tracking only for the recomputation.
5. Anything stochastic inside a segment must therefore produce identical values both times, so dropout and similar operations need their random state captured and restored — otherwise the gradients are computed against activations the forward pass never actually used.
6. That requirement is handled by the framework rather than the user, but it is why checkpointing interacts with custom layers that manage their own randomness.

Example: Store a few, recompute the rest

Backpropagation normally keeps every layer's forward activations so the backward pass can use them, which scales with depth. Checkpointing keeps only a subset and recomputes the others on the way back. Segmenting an 80-layer model at roughly √80 ≈ 9 means holding about 9 checkpoints and re-deriving at most 9 layers at a time — on the order of 18 activation sets instead of 80. The price is one extra forward pass over the recomputed segments, commonly around a third more compute. It is the cleanest trade in the training stack: a large, predictable memory saving for a modest, predictable slowdown.

Commonly confused with

This targets activation memory, which grows with batch size and sequence length. Sharded strategies like ZeRO and FSDP target parameter, gradient and optimizer memory. They are complementary, not alternatives — a run can be limited by either, and the fix differs, so measure which one is actually filling the device before choosing.

When to use it

Reach for it when:

• Long sequences or large batches, where activations dominate memory
• Fitting a training run onto the hardware you have rather than acquiring more
• As the first lever when out-of-memory happens during the backward pass

Reach for something else when:

• When compute-bound and memory is comfortable — you would pay the slowdown for nothing
• Inference, which does not retain activations for a backward pass
• Before checking whether the optimizer states, not the activations, are the constraint

Primary source

• Chen et al., Training Deep Nets with Sublinear Memory Cost (arXiv)

Continue your research

Use these connected ModelRefs sections to compare alternatives, inspect implementation paths, and review the evidence and governance boundaries relevant to Gradient Checkpointing — AI Glossary.

• Browse models
• Review providers
• Inspect benchmarks
• Explore workflows
• Learn AI concepts
• Read decision guides
• Review trust and methodology
• Start a decision path