PyTorch Gradient Checkpointing — Core Concepts

The Memory Problem in Deep Learning

When a neural network processes data during training, it stores intermediate activations — the outputs of every layer during the forward pass. These activations are essential for computing gradients during backpropagation. For a model like ResNet-152 or a Transformer with hundreds of millions of parameters, these stored activations can consume tens of gigabytes of GPU memory.

This creates a hard ceiling: your model size is limited by your GPU’s RAM, not by your ideas.

How Gradient Checkpointing Works

Gradient checkpointing splits the model into segments. Only the activations at segment boundaries (the “checkpoints”) are kept in memory. Everything in between is discarded after the forward pass.

During backpropagation, when the optimizer needs those discarded activations, PyTorch re-runs the forward pass for that segment from the nearest checkpoint. The activations are recalculated just in time, used for gradient computation, and then discarded again.

The Tradeoff

Aspect	Without Checkpointing	With Checkpointing
Memory usage	O(n) — all layers	O(√n) — only checkpoints
Compute time	1× forward pass	~1.3× forward pass
Backward pass	Standard	Standard (after recomputation)

The √n memory reduction is the key insight. For a 100-layer network, you might store activations for 10 checkpointed layers instead of all 100.

When to Use It

Gradient checkpointing makes sense when:

• Your model nearly fits in GPU memory but batch size is painfully small
• You want to train a larger model without upgrading hardware
• Fine-tuning a large pretrained model on a single GPU
• Research experiments where memory is the bottleneck, not speed

It does not help when:

• Your model easily fits in memory (the overhead isn’t worth it)
• Inference only — checkpointing is purely a training technique
• Your bottleneck is data loading or I/O, not memory

How It Looks in Practice

PyTorch provides torch.utils.checkpoint.checkpoint to wrap specific model segments. You apply it to blocks of layers rather than individual layers — wrapping every single layer adds too much overhead for too little benefit.

Most practitioners checkpoint every Transformer block or every residual group in a CNN. The Hugging Face transformers library has a gradient_checkpointing_enable() method that handles this automatically for supported architectures.

Common Misconception

Many people think gradient checkpointing changes the mathematical result of training. It doesn’t. The gradients are identical — you’re recomputing activations, not approximating them. The model converges to the same solution; it just takes longer per step.

This distinguishes it from techniques like mixed-precision training or gradient accumulation, which change the numerical behavior (even if slightly).

Combining with Other Techniques

Gradient checkpointing stacks well with:

• Mixed-precision training (FP16/BF16): Reduces memory per activation, and checkpointing reduces the count
• Gradient accumulation: Simulates larger batch sizes without storing more activations simultaneously
• Model parallelism: Splits layers across GPUs; checkpointing reduces memory per GPU

Together, these techniques let researchers at universities train models that would otherwise require datacenter-scale infrastructure.

The one thing to remember: Gradient checkpointing is a pure time-for-memory swap — it recomputes activations instead of storing them, cutting memory by up to 70% with identical training results.