Technical GUIDE
Gradient Checkpointing
Gradient checkpointing (also called activation checkpointing) is a memory-saving trick that throws away most intermediate activations during the forward pass and recomputes them on the fly during backpropagation.
Listen to this guideAI-narrated · ~2 min
On this page2 min read
Overview
It lets you train deeper, larger networks by trading extra compute for much lower memory use.
Deep Dive
Training normally retains intermediate activations needed for backpropagation. Activation checkpointing saves selected boundary tensors and recomputes internal activations during backward. This trades additional computation for a lower activation-memory footprint; the actual savings and runtime overhead depend on checkpoint placement, operations, model graph, and hardware. In a simplified chain of N similarly sized layers, a textbook schedule with checkpoint boundaries spaced about sqrt(N) layers apart can reduce activation storage to order sqrt(N) checkpoints, with recomputation cost governed by that schedule. This asymptotic example is not a universal prediction for framework implementations. PyTorch exposes activation-checkpointing utilities, and its current guidance also describes selective rematerialization and compiler-assisted choices. It can combine with FSDP and mixed precision, but benchmark the full training workload.
Real-World Implementation
Training a deep transformer with a larger batch size on a single GPU by discarding and recomputing layer activations.
Fine-tuning vision models on high-resolution images where activation maps would otherwise overflow GPU memory.
Hugging Face Transformers enabling gradient_checkpointing=True to fit billion-parameter models during fine-tuning.
Combining checkpointing with FSDP so both parameters and activations are kept small, enabling training of very large language models.
The Future of Gradient Checkpointing
Gradient checkpointing is now standard in large-model training and is increasingly automated, with libraries selecting optimal checkpoint locations for you. It pairs naturally with FSDP, mixed precision, and offloading to push model sizes higher. Expect 'selective' checkpointing that recomputes only cheap operations while keeping expensive ones (like attention matrices) cached, plus compiler-driven approaches in tools like PyTorch's torch.compile that automatically decide what to save versus recompute for the best speed-memory balance.
Putting it into practice
Why it matters
- Cost and budget. Architecture decisions drive performance and operating cost for years.
- Clearer decisions. Technical education helps teams choose the right stack, not just the newest one.
- Quality control. Better engineering choices reduce reliability incidents in production.
Watch out for
- Optimizing one benchmark can hide broader system weaknesses.
- Infrastructure and maintenance costs are often underestimated.
- Security and observability gaps can grow as systems become more complex.
Next steps
- Define latency, quality, and cost targets before implementation.
- Benchmark under realistic load and data conditions.
- Instrument monitoring for errors, drift, and user impact.
- Prepare rollback and incident response paths before scaling.
Keep Exploring
Free newsletter
Keep up with AI in 3 minutes a day
One short email each weekday with the three AI stories that actually matter. Free forever, no ads.
One email each weekday. Unsubscribe in one click. We never sell or share your address.
Test yourself
Take the Gradient Checkpointing quiz
Instant feedback on every answer, and a shareable certificate with a verifiable ID once you pass a course.
Support free AI education. AI Understanding is a 501(c)(3) nonprofit — no ads, no paywall, ever. Make a donation
Frequently asked questions
What is Gradient Checkpointing?
Gradient checkpointing (also called activation checkpointing) is a memory-saving trick that throws away most intermediate activations during the forward pass and recomputes them on the fly during backpropagation. It lets you train deeper, larger networks by trading extra compute for much lower memory use.
What does gradient checkpointing primarily trade in order to save memory?
Gradient checkpointing recomputes discarded activations during the backward pass, spending extra compute in exchange for substantially reduced memory.
Why are activations normally stored during the forward pass?
Backprop computes gradients using the intermediate activations from the forward pass, so they must be available unless they are recomputed.
In a simplified chain of N similarly sized layers, what can checkpoint boundaries spaced about sqrt(N) layers apart achieve in the textbook schedule?
The square-root schedule is a theoretical result for a simplified sequential computation graph, not a universal implementation guarantee.
What most determines gradient-checkpointing runtime overhead in a practical model?
Recomputation overhead varies with checkpoint policy and the cost of operations selected for recomputation.
In PyTorch, which utility is commonly used to apply gradient checkpointing to a module?
torch.utils.checkpoint wraps a module so its internal activations are recomputed during backward instead of being stored.
Keep learning
Related guides
More guides picked for this topic