NEW
Also known as: Activation Checkpointing, Recomputation
A memory-saving technique that recomputes intermediate activations during backpropagation instead of storing them all.
DESCRIPTION
Gradient checkpointing trades compute for memory by saving only selected activations (checkpoints) during the forward pass. During backward propagation, missing activations are recomputed to obtain gradients, reducing peak memory use significantly. It enables training larger models or batches on limited hardware. Frameworks like PyTorch and DeepSpeed provide built-in APIs for this technique. The trade-off is additional compute overhead and potential wall-clock slowdowns. Properly placing checkpoints balances memory savings and recomputation cost.
ALTERNATIVE DEFINITIONS
PyTorch
Checkpointing allows trading compute for memory by storing only a subset of activations and recomputing the rest during the backward pass.
Captures the compute–memory trade-off and recomputation mechanism.
DeepSpeed
Activation checkpointing reduces memory consumption by re-computing activations during backpropagation instead of storing them all in memory.
Emphasizes memory reduction via activation recomputation in DeepSpeed.
RELATED TERMS
Activation Memory ZeRO Tensor Parallelism Micro-batching Mixed Precision Offloading FlashAttention Gradient Accumulation Memory Optimization Sequence Parallelism
COMPETING TERMS
Full-Activation Training Activation Offloading
RELATIONS
- Broader: Checkpointing
KNOWLEDGE AREAS
TARGET AUDIENCE
AI Engineer MLOps Engineer Hardware Engineer
EXTERNAL SOURCES
- Wikipedia - Checkpoint restart
- PyTorch - Official checkpointing API
- DeepSpeed - Activation checkpointing tutorial
- arXiv - Reversible residual networks paper
TOP QUESTIONS
When should I enable Gradient Checkpointing?
Enable it when GPU memory limits prevent desired batch sizes or model sizes, accepting additional compute overhead.
Category: Relevancy
How much memory can Gradient Checkpointing save?
Savings vary by placement; 30–70% reductions are common in large transformers with careful checkpointing.
Category: Example
What is the performance cost of Gradient Checkpointing?
It adds recomputation, increasing training time; exact penalty depends on checkpoint frequency and model architecture.
Category: Risks
How do I use Gradient Checkpointing in PyTorch?
Wrap submodules with torch.utils.checkpoint or use model-wide flags in libraries that expose it.
Category: Requirements
Does Gradient Checkpointing affect numerical stability?
Not directly; it recomputes activations. Combine with mixed precision carefully to avoid accumulation errors.
Category: Testing
Can Gradient Checkpointing be combined with ZeRO?
Yes, it complements optimizer state partitioning and offloading to further extend memory limits.
Category: Management
Where should I place checkpoints?
Place them around memory-heavy regions (e.g., attention blocks), balancing recomputation cost and savings.
Category: Best practices
Does Gradient Checkpointing help during inference?
No, it targets training memory; inference benefits come from KV caching and quantization instead.
Category: Misunderstanding
What are tooling options for Gradient Checkpointing?
PyTorch, DeepSpeed, and JAX offer APIs; some libraries expose automatic strategies.
Category: Overview