Gradient Checkpointing

Gradient checkpointing is a memory-efficient training technique for deep neural networks that reduces peak activation memory by storing only a strategically chosen subset of intermediate activations (checkpoints) during the forward pass, then recomputing the discarded activations on demand during backpropagation. The technique was formalised by Chen et al. (2016) under the name ‘Training Deep Networks with Sublinear Memory Cost’, achieving O(sqrt(N)) memory in the number of layers N at the cost of one additional forward pass per training step. It is now a foundational primitive in large-model training, implemented natively in PyTorch, JAX, TensorFlow, and most major deep-learning frameworks.

Overview

• During standard Backpropagation, the chain rule requires that every intermediate activation produced in the Forward Pass be available when computing gradients in the backward pass. For a network with N layers trained on a mini-batch, this means keeping O(N) activation tensors in memory simultaneously, which can exhaust GPU Memory for very deep or wide models.
• Gradient checkpointing solves this by saving activations only at a subset of layer boundaries (the checkpoints), and discarding the activations in between. When the backward pass reaches a segment whose activations were discarded, it pauses and re-executes the forward computation for that segment alone, starting from the most recent checkpoint. The recomputed activations are used immediately and then released.
• The classical segmented scheme achieves O(sqrt(N)) memory growth in the number of layers, at the cost of roughly one extra forward pass per batch step (i.e. 1/3 more compute overall compared to full storage). More sophisticated policies — including learned or profile-guided checkpoint selection — can shift the memory–compute trade-off to better suit specific hardware budgets.
• The technique is architecture-agnostic: it applies to Convolutional Neural Network pipelines, Recurrent Neural Network unrolled through time (a particularly high-memory scenario), Transformer decoder stacks, and any other model expressible as a directed acyclic Computational Graph.

Key Mechanisms

• Checkpoint Placement
  • Uniform segmentation: divide the graph into sqrt(N) segments of sqrt(N) layers each, store only segment boundary activations.
  • Non-uniform placement: profile memory footprint per layer and place checkpoints at memory peaks, minimising recomputation cost for a given budget.
  • Selective checkpointing: mark only memory-heavy operations (e.g. attention softmax in Transformer blocks) rather than entire segments.
• Recomputation on Demand
  • During the backward pass, when a gradient requires a discarded activation, the framework pauses gradient flow, reruns the forward function for the enclosing segment, caches the result briefly, and resumes.
  • This relies on Automatic Differentiation engines (e.g. PyTorch autograd, JAX jit) that can separate forward and backward execution and inject custom checkpointing hooks.
• Segment Boundary Tensors
  • The boundary tensors (checkpoints) are kept in device memory across the entire backward pass, so checkpoint count directly determines minimum memory floor.
  • Offloading checkpoints to CPU RAM (combined with Offloading strategies) can further reduce GPU memory at the cost of PCIe transfer latency.
• Memory–Compute Trade-off
  • Full storage: O(N) memory, no recomputation overhead.
  • Gradient checkpointing (uniform): O(sqrt(N)) memory, ~33% extra compute.
  • Extreme: O(1) memory (store only input and output), recompute entire network per gradient — impractical but illustrates the continuum.
• Integration with Other Techniques
  • Composable with Mixed Precision Training (FP16/BF16 reduce checkpoint tensor sizes).
  • Composable with Gradient Accumulation (allows effective batch sizes larger than device memory independently of checkpointing).
  • Used within Pipeline Parallelism to reduce bubble memory at pipeline stages.

Applications and Use Cases

• Large Language Model Training
  • Training Transformer-based models such as GPT-class architectures with billions of parameters requires fitting long sequence activations in memory; gradient checkpointing is applied at every transformer block boundary as standard practice.
• Long Context Training
  • Training or fine-tuning with very long sequences (e.g. 32 K–1 M tokens) causes quadratic Attention Mechanism memory growth; per-block checkpointing is essential to keep memory tractable alongside techniques like Flash Attention.
• Fine-Tuning on Consumer Hardware
  • Parameter-efficient fine-tuning methods such as LoRA and QLoRA pair gradient checkpointing with quantisation to fit fine-tuning of 7B–70B parameter models on single consumer GPUs (24 GB VRAM class).
• Scientific Computing
  • Neural PDE solvers, physics-informed networks, and neural ODEs unrolled over thousands of time steps use gradient checkpointing (sometimes called rematerialisation in the JAX/XLA ecosystem) to make time-series backpropagation feasible.
• Video and 3D Models
  • Spatiotemporal models processing video frames or volumetric data accumulate activation memory proportional to the number of frames/voxels; checkpointing is applied per temporal or spatial segment.
• Recurrent Neural Network Training
  • Truncated Backpropagation Through Time combined with gradient checkpointing allows training over very long sequences without exploding memory.
• Edge AI and Resource-Constrained Training
  • On-device continual learning on embedded or mobile hardware leverages gradient checkpointing to keep memory within tight device constraints, bridging large-model methodology with edge deployment.

Semantic Classification

Standards and Context

• No formal standards body governs gradient checkpointing; it is documented as a first-class feature in major framework documentation.
• PyTorch: torch.utils.checkpoint.checkpoint() and checkpoint_sequential() APIs; semantics stabilised across PyTorch 1.x–2.x.
• JAX/XLA: jax.checkpoint (formerly jax.remat) provides rematerialisation as a function transformation, composable with jit, vmap, and pmap.
• TensorFlow: tf.recompute_grad decorator offers equivalent functionality.
• DeepSpeed: ZeRO-Offload and ZeRO-Infinity stages incorporate activation checkpointing as a configurable sub-option alongside parameter and optimiser-state offloading.
• Megatron-LM: applies selective activation recomputation per transformer layer as default for large-scale Distributed Training runs.
• Hugging Face Transformers: model.gradient_checkpointing_enable() wrapper activates per-block checkpointing for all supported architectures, used routinely in LoRA/QLoRA fine-tuning recipes.
• The underlying algorithm is described in Chen et al. 2016 (“Training Deep Networks with Sublinear Memory Cost”) and Kumar et al. 2019 (optimal checkpoint placement for general DAGs).

Provenance