Graph
Gradient Checkpointing

🧮Gradient Checkpointing

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.

Source

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.

Source

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

KNOWLEDGE AREAS

AI MLOps Infrastructure

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