Technical GUIDE
Activation Recomputation Tradeoffs
Activation recomputation (gradient or activation checkpointing) saves GPU memory during training by discarding intermediate activations in the forward pass and recomputing them during the backward pass.
Listen to this guideAI-narrated · ~2 min
On this page2 min read
Overview
It trades extra compute for the ability to train larger models or longer sequences on the same hardware.
Deep Dive
Backpropagation needs the forward-pass activations to compute gradients, so by default every layer's outputs are stored — a huge memory cost that grows with model size, batch size, and sequence length. Activation recomputation keeps only a few 'checkpoint' tensors (often just layer boundaries) and throws away the rest. During the backward pass, it re-runs the forward computation between checkpoints to regenerate the discarded activations on demand. For an idealized N-layer chain with comparable activation sizes, square-root checkpoint placement can reduce activation storage to O(sqrt(N)) at the cost of recomputation. This does not describe all parameter, optimizer, or workspace memory. Selective variants recompute only cheap-but-memory-heavy ops (like attention or dropout) while caching expensive ones, getting most of the memory savings for far less recompute overhead.
Real-World Implementation
Training a large transformer that wouldn't otherwise fit by checkpointing each layer block
Using PyTorch's torch.utils.checkpoint to wrap transformer blocks and cut activation memory
Selective recomputation of attention/softmax in Megatron-LM to save memory with minimal slowdown
Enabling longer sequence lengths on a fixed GPU budget by recomputing activations instead of storing them
The Future of Activation Recomputation Tradeoffs
Recomputation is increasingly automated and selective. Frameworks now profile each op's memory and FLOP cost to choose optimal checkpoints, and combine recomputation with activation offloading to CPU/NVMe and with parallelism strategies. As context lengths and model sizes keep growing, expect compiler-driven policies (in PyTorch, JAX/XLA) that pick per-op recompute decisions automatically, plus tighter overlap of recompute with communication so the extra FLOPs are partly hidden.
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 Activation Recomputation Tradeoffs 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 Activation Recomputation Tradeoffs?
Activation recomputation (gradient or activation checkpointing) saves GPU memory during training by discarding intermediate activations in the forward pass and recomputing them during the backward pass. It trades extra compute for the ability to train larger models or longer sequences on the same hardware.
What does activation recomputation trade away to save memory?
Recomputation discards stored activations and regenerates them in the backward pass, spending extra compute to reduce memory usage.
Why are forward-pass activations normally stored at all?
The backward pass uses the forward activations to compute gradients, so by default they are kept in memory until the backward pass runs.
Why should activation checkpointing overhead be measured on the actual model?
Recomputation spends additional compute, but the observed runtime depends on the selected operations and system.
For an idealized N-layer chain with comparable activations, what storage scaling can square-root checkpoint placement achieve?
The theoretical square-root result concerns activation checkpointing in an appropriate chain model, not every category of training memory.
What is the idea behind selective (not full) recomputation?
Selective recomputation targets ops that use lots of memory but little compute (like softmax or layernorm), while caching expensive GEMM results to minimize wasted FLOPs.
Keep learning
Related guides
More guides picked for this topic