Gradient Checkpointing

Gradient checkpointing fundamentally trades increased computation for reduced memory. Instead of storing all intermediate activations from the forward pass for the backward pass, it selectively stores only a subset (checkpoints). The non-stored activations are recomputed on-demand during backpropagation from the nearest checkpoint. This reduces peak memory usage from O(n) to O(√n) for a chain of n layers, where n is the sequence length or depth.

The efficiency of gradient checkpointing depends heavily on where checkpoints are placed. Common strategies include:

• Uniform Placement: Placing checkpoints at regular intervals (e.g., every k layers).
• Optimal Dynamic Programming: Using an algorithm to find the placement that minimizes total recomputation cost for a given memory budget.
• Manual Placement for Model Architecture: Targeting memory-intensive operations (e.g., after large convolutional blocks or attention layers in transformers). Poor placement can lead to excessive recomputation with minimal memory savings.

Gradient checkpointing is a standard feature in deep learning frameworks, abstracting the complex recomputation logic.

• PyTorch: torch.utils.checkpoint.checkpoint function wraps a segment of the model. It runs the segment twice: once without storing intermediates (forward) and once with gradients enabled (for recomputation during backward).
• TensorFlow: tf.recompute_grad decorator or the GradientTape with persistent=False can be used to trigger similar behavior.
• JAX: The jax.checkpoint (formerly jax.remat) transformation is a first-class citizen, allowing flexible checkpointing policies ('none', 'max', 'min' memory).

This technique is critical for overcoming hardware memory limits when training massive models.

• Long Sequence Transformers: Enables training language models with context windows longer than what would fit in GPU memory if all activations were stored.
• Large-Scale Vision Models: Allows training of high-resolution image models (e.g., diffusion models, large CNNs) by checkpointing within the U-Net or residual blocks.
• Scientific Models: Facilitates training of deep physics-informed neural networks (PINNs) or other models with exceptionally deep or wide architectures.

The overhead is predictable but non-negligible.

• Memory Reduction: Can reduce activation memory by 50-90%, depending on the model and checkpointing strategy.
• Compute Increase: Typically increases total training computation by 20-40% due to the second forward pass (recomputation) within checkpointed segments.
• Wall-Clock Time: Training time increases, but the alternative (model parallelism or offloading) often has higher communication overhead. It enables training larger batches or models on existing hardware.

Gradient checkpointing is often combined with other optimizations in a memory reduction stack:

• With Activation Offloading (CPU RAM): Checkpoints can be stored on CPU, further reducing GPU memory pressure.
• With Model Parallelism: Used within individual model-parallel devices to handle layers that are still too large.
• Complementary to Quantization/Pruning: Checkpointing targets activation memory during training, while quantization and pruning primarily target parameter memory for inference. They address different parts of the memory problem.
• Contrast with KV Caching: KV caching optimizes inference memory for transformers. Gradient checkpointing optimizes training memory for any autograd-based model.