Gradient Checkpointing ♻️

Table of Contents

• 0. Intuition: The Hiker’s Dilemma
• 1. The Memory Problem
  • 1.1 Where Memory Goes During Training
  • 1.2 Activation Memory at GPT-2 Scale
• 2. The Technique
  • 2.1 Formal Setup
  • 2.2 Algorithm
• 3. Optimal Checkpoint Placement
  • 3.1 Cost Model
  • 3.2 The sqrt(n) Bound
  • 3.3 Pareto Trade-off and the Recursive Extension
• 4. Selective Recomputation
  • 4.1 Activation Cost Asymmetry in Transformers
  • 4.2 The Korthikanti et al. Strategy
  • 4.3 Memory Breakdown
  • 4.4 Results Comparison
• 5. PyTorch Implementation
  • 5.1 torch.utils.checkpoint API
  • 5.2 use_reentrant: True vs. False
  • 5.3 Wrapping Transformer Blocks
• 6. Interactions and Gotchas
  • 6.1 Dropout and RNG State
  • 6.2 Distributed Data Parallel
  • 6.3 Mixed Precision and Autocast
  • 6.4 Double Backward
• References

0. Intuition: The Hiker’s Dilemma 🧭

Imagine a hiker climbing a mountain who needs to retrace their path exactly on the way back down. They have two options: photograph every single step (expensive, heavy backpack), or photograph only a handful of landmarks and trust that they can re-walk each segment between landmarks when needed.

Gradient checkpointing is the second strategy applied to neural network training.

Backpropagation needs the activations from every layer of the network to compute gradients — the chain rule requires multiplying local Jacobians together, and each Jacobian depends on the activation at that layer. Naively, one stores all activations from the forward pass and reads them back during the backward pass.

The key observation: if you have the input to a segment of computation, you can always rerun that segment to regenerate its internal activations. You don’t need to store them — you just need to store enough checkpoints (snapshots at strategic positions) that no segment is too expensive to re-traverse. This trades some extra computation for a large reduction in peak memory.

1. The Memory Problem 💾

1.1 Where Memory Goes During Training

During training, GPU memory holds four distinct categories of tensors:

1. Parameters \(\theta\): the model weights, size \(\approx P\) bytes (where \(P\) counts parameters times bytes-per-parameter).
2. Gradients \(\nabla_\theta \mathcal{L}\): same shape as parameters, so another \(\approx P\) bytes.
3. Optimizer state: e.g. Adam maintains first and second moment estimates, adding \(2P\) more bytes for a total of \(3P\) bytes of optimizer state.
4. Activations: the intermediate tensors produced during the forward pass and needed for backpropagation.

For a transformer with \(L\) layers, batch size \(B\), sequence length \(T\), and hidden dimension \(d\), activation memory per layer consists of:

• Attention queries, keys, values: \(3 \times B \times T \times d\) elements
• Attention weight matrix: \(B \times H \times T \times T\) elements (where \(H\) = number of heads)
• MLP intermediate activations: \(\sim 4 \times B \times T \times d\) elements (for a \(4d\) expansion)
• LayerNorm inputs and outputs: \(\sim 2 \times B \times T \times d\) elements

Summing over \(L\) layers, total activation memory scales as \(O(BLTd)\), with an additional \(O(BLHT^2)\) term from attention weights.

1.2 Activation Memory at GPT-2 Scale

For concreteness, take GPT-2 Small (\(L = 12\), \(d = 768\), \(H = 12\), \(T = 1024\), \(B = 8\)) with fp32 (4 bytes/element).

Parameter and optimizer memory (independent of batch/sequence):

Activation memory per layer (fp32, \(B=8\), \(T=1024\)):

Total activation memory across all 12 layers: \(\approx 7.2\) GB.

Total parameter + optimizer memory: \(\approx 1.9\) GB.

Activations dominate by roughly 4:1 at training time. This ratio grows linearly with \(B\) and \(T\), making large-batch or long-sequence training infeasible without intervention. At inference time, no activations need to be stored across layers, so parameters dominate and the ratio inverts entirely.

2. The Technique 🔬

2.1 Formal Setup

Definition (Computation Graph). A forward pass defines a directed acyclic graph \(G = (V, E)\) where each node \(v \in V\) is an intermediate tensor and each edge \((u, v) \in E\) means \(v\) depends on \(u\). Backpropagation traverses \(G\) in reverse topological order, computing \(\partial \mathcal{L}/\partial v\) for each \(v\) using stored values of the tensors on which \(v\) depends.

Definition (Checkpoint). A checkpoint is a designated node \(c \in V\) whose value is retained in memory across the full forward-backward pass. All non-checkpoint nodes are discarded after they are produced in the forward pass.

Definition (Segment). Given a sequence of checkpoints \(c_0, c_1, \ldots, c_K\) (where \(c_0\) is the input and \(c_K\) is the loss), the \(i\)-th segment \(S_i\) is the subgraph of \(G\) between consecutive checkpoints \(c_{i-1}\) and \(c_i\). The segment \(S_i\) has a well-defined recomputation: starting from \(c_{i-1}\), one can re-execute the forward pass of \(S_i\) to regenerate all internal tensors.

2.2 Algorithm

The full gradient checkpointing procedure is:

import torch
from typing import Callable, List

def checkpointed_forward_backward(
    segments: List[Callable],
    x: torch.Tensor,
) -> torch.Tensor:
    """
    Forward pass with gradient checkpointing.
    
    Args:
        segments: list of K callables, each representing one segment S_i.
                  Segment i maps checkpoint c_{i-1} to checkpoint c_i.
        x:        initial input (checkpoint c_0).
    
    Returns:
        Output of the final segment (checkpoint c_K).
    """
    checkpoints = [x]  # store c_0

    # --- Forward pass: run each segment, store only the output checkpoint ---
    with torch.no_grad():
        for segment in segments:
            x = segment(x)
            checkpoints.append(x)  # store c_i; internal tensors of S_i are discarded

    # --- Backward pass: for each segment, recompute internals then differentiate ---
    loss = checkpoints[-1]
    loss.requires_grad_(True)

    for i in reversed(range(len(segments))):
        c_prev = checkpoints[i]
        c_prev.requires_grad_(True)

        # Rerun segment i under autograd to rebuild the computation graph
        with torch.enable_grad():
            c_next_recomputed = segments[i](c_prev)

        # Backpropagate through the recomputed segment
        c_next_recomputed.backward(checkpoints[i + 1].grad)

    return checkpoints[-1]

The key structural property: at any point during the backward pass, memory holds only the \(K\) checkpoint tensors plus the activations of the single segment currently being recomputed. Peak memory during backward is therefore:

\[M_{\text{backward}} = \underbrace{K \cdot a_c}_{\text{checkpoint tensors}} + \underbrace{\max_i M(S_i)}_{\text{largest segment recomputation}}\]

where \(a_c\) is the (uniform) size of a checkpoint tensor and \(M(S_i)\) is the peak activation memory needed to recompute segment \(S_i\).

3. Optimal Checkpoint Placement 📐

3.1 Cost Model

Consider a sequential network of \(n\) layers, each producing an activation of uniform size \(a\) (bytes or elements). We place \(K - 1\) checkpoints, dividing the network into \(K\) equal segments of \(n/K\) layers each.

Memory cost has two components:

1. Checkpoint storage: \(K\) checkpoints (including input and output), so \(K \cdot a\).
2. Segment recomputation: within any segment of \(n/K\) layers, activations are produced one by one; peak activation count within a segment is \(n/K\), so the recomputation buffer costs \((n/K) \cdot a\).

Total peak memory:

\[M(K) = K \cdot a + \frac{n}{K} \cdot a = a\left(K + \frac{n}{K}\right)\]

Recomputation cost: each of the \(K\) segments is re-executed once during backward, and each segment costs \(n/K\) layer evaluations, so total extra forward-pass work is:

\[C(K) = K \cdot \frac{n}{K} = n \quad \text{layer evaluations}\]

Recompute cost is always \(O(n)\) regardless of \(K\), i.e., one extra forward pass through the whole network. Only the memory cost depends on \(K\).

3.2 The sqrt(n) Bound

To minimize \(M(K) = a(K + n/K)\) over \(K > 0\), differentiate:

\[\frac{dM}{dK} = a\left(1 - \frac{n}{K^2}\right) = 0 \implies K^* = \sqrt{n}\]

The second derivative \(2an/K^3 > 0\) confirms this is a minimum. Substituting \(K^* = \sqrt{n}\):

\[M(K^*) = a\left(\sqrt{n} + \frac{n}{\sqrt{n}}\right) = a \cdot 2\sqrt{n} = O(\sqrt{n} \cdot a)\]

With \(\sqrt{n}\) uniformly spaced checkpoints, peak activation memory is \(O(\sqrt{n} \cdot a)\), reduced from \(O(n \cdot a)\), at a recompute cost of exactly one extra forward pass.

Optimality of Uniform Spacing

For sequential networks with equal activation sizes, uniform checkpoint spacing is optimal. The argument: given fixed \(K\), memory \(M(K)\) is minimized when the maximum segment length is minimized, which occurs when segments are equal. Any non-uniform partition increases the maximum segment length and hence increases the recomputation buffer term \(\max_i M(S_i)\).

3.3 Pareto Trade-off and the Recursive Extension

The \(K\)-checkpoint family traces a Pareto curve in (memory, recompute) space:

import numpy as np
import matplotlib.pyplot as plt

def memory_cost(K: int, n: int, a: float = 1.0) -> float:
    """Peak activation memory with K checkpoints (in units of a)."""
    return K + n / K  # = a*(K + n/K) / a

def recompute_cost(K: int, n: int) -> float:
    """Extra forward evaluations with K checkpoints."""
    return n  # always O(n) for this model

n = 96  # e.g. 96 transformer layers
K_values = np.arange(1, n + 1)
memories = [memory_cost(K, n) for K in K_values]
recomputes = [recompute_cost(K, n) for K in K_values]

# Minimum is at K = sqrt(n)
K_opt = int(np.sqrt(n))
print(f"Optimal K = {K_opt}, memory = {memory_cost(K_opt, n):.1f}a, extra compute = {n} layers")

Since recompute cost is flat at \(O(n)\) for all \(K\), the Pareto curve degenerates: the \(\sqrt{n}\) checkpoint placement is simultaneously memory-optimal and recompute-optimal among all \(K > 0\).

Recursive extension. Chen et al. (2016) also derive a recursive variant: apply checkpointing hierarchically within each segment. This reduces memory to \(O(\log n \cdot a)\) at the cost of \(O(n \log n)\) extra compute. Concretely, with \(k = 1\) checkpoint per segment and \(\lceil \log_2 n \rceil\) levels of recursion, memory is \(O(\log n)\).

4. Selective Recomputation 🎯

4.1 Activation Cost Asymmetry in Transformers

The uniform-cost assumption of Section 3 breaks down for transformers because different activation tensors have very different size-to-recompute-cost ratios.

Definition (Recomputation Efficiency). For an activation tensor of size \(s\) bytes, let \(f\) be the FLOPs required to recompute it from the nearest preceding checkpoint. The recomputation efficiency is \(s/f\): bytes saved per FLOP of recomputation. Higher values mean it is cheaper (per byte) to recompute.

In a transformer block with sequence length \(T\), hidden dimension \(d\), and \(H\) attention heads:

Attention weights (the \(T \times T\) softmax output) are large but cheap to recompute: the compute cost is \(O(T^2 d_h)\) while the memory cost is \(O(T^2)\), so efficiency scales as \(O(1/d_h)\). MLP inputs, by contrast, require rerunning the entire preceding attention block to regenerate.

4.2 The Korthikanti et al. Strategy

Korthikanti et al. (2022) introduce selective activation recomputation: store only the activations that are expensive to recompute; recompute the ones that are cheap (large but cheap).

Definition (Selective Recomputation Policy).

• Recompute (do not store): \(QK^T\) matrix product, softmax output, softmax dropout output, weighted value aggregation \(\text{softmax}(\cdot) V\).
• Store (retain in memory): LayerNorm inputs before self-attention, QKV projection inputs, MLP layer inputs, dropout masks for non-attention paths.

The rationale for recomputing attention activations: they account for the \(O(T^2)\) memory but can be regenerated from the stored Q, K, V projections with only \(O(T^2 d_h)\) work — which is the same work done during the original forward pass of the attention sublayer.

4.3 Memory Breakdown

For a single transformer layer without any parallelism (\(s =\) sequence length, \(b =\) batch size, \(h =\) hidden dimension, \(a =\) number of attention heads), activation memory (in bytes, bf16) is:

The \(5as^2b / h\) term — from attention weights — is the \(O(T^2)\) component that selective recomputation targets.

With selective recomputation (discarding attention weight tensors and related intermediates), the \(5as^2b\) term is eliminated from stored memory. The remaining stored activations cost \(sbh \cdot 34\) bytes per layer, independent of sequence length.

4.4 Results Comparison

The three regimes, at large scale (Megatron-LM 22B–530B parameter range):

Selective recomputation achieves comparable memory savings to full checkpointing at less than 10% compute overhead, compared to 30–40% for full checkpointing. This is the central quantitative result of Korthikanti et al. (2022).

5. PyTorch Implementation 💻

5.1 torch.utils.checkpoint API

The primary entry point is torch.utils.checkpoint.checkpoint:

torch.utils.checkpoint.checkpoint(
    function,           # callable: the segment to checkpoint
    *args,              # positional inputs to function
    use_reentrant=None, # bool: controls implementation variant (see below)
    context_fn=noop_context_fn,  # callable returning a pair of context managers
    determinism_check='default', # 'default' | 'none'
    debug=False,
    early_stop=True,    # non-reentrant only: stop recompute early when possible
    **kwargs,           # keyword arguments (non-reentrant only)
)

The function returns the same output as function(*args, **kwargs) but with lower peak activation memory.

5.2 use_reentrant: True vs. False

The use_reentrant flag selects between two fundamentally different implementations.

use_reentrant=True (legacy, reentrant variant):

The forward pass is run under torch.no_grad(), so no autograd graph is recorded for the checkpointed region. During backward, PyTorch re-executes the function in a nested (“reentrant”) backward context to reconstruct the computation graph on the fly.

Limitations: - Requires at least one input and one output to have requires_grad=True. - Cannot handle detach() or torch.no_grad() inside the checkpointed function — these confuse the reentrant backward. - Does not support torch.autograd.grad() (only .backward()). - Incompatible with double backward (second-order gradients).

use_reentrant=False (recommended, non-reentrant variant):

The forward pass records the autograd graph normally via saved-tensor hooks: instead of saving intermediate tensors, hooks are registered that trigger recomputation on demand when a saved tensor is accessed during backward.

Advantages: - No requires_grad constraint on inputs. - Handles detached tensors, nested checkpoints, keyword arguments. - Supports torch.autograd.grad(). - Compatible with double backward (with preserve_rng_state=True). - early_stop=True (default): recomputation halts as soon as all needed tensors have been regenerated — avoids recomputing the full segment when only early activations are needed.

Starting in PyTorch 2.9, omitting use_reentrant raises an error — always pass it explicitly.

5.3 Wrapping Transformer Blocks

Pattern 1: Checkpoint at the transformer block level (most common)

import torch
import torch.nn as nn
from torch.utils.checkpoint import checkpoint

class TransformerBlock(nn.Module):
    def __init__(self, d_model: int, n_heads: int, ffn_dim: int):
        super().__init__()
        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
        self.ln1 = nn.LayerNorm(d_model)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = nn.Sequential(
            nn.Linear(d_model, ffn_dim),
            nn.GELU(),
            nn.Linear(ffn_dim, d_model),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Attention sublayer
        attn_out, _ = self.attn(x, x, x, need_weights=False)
        x = self.ln1(x + attn_out)
        # FFN sublayer
        x = self.ln2(x + self.ffn(x))
        return x


class Transformer(nn.Module):
    def __init__(self, n_layers: int, d_model: int, n_heads: int, ffn_dim: int,
                 use_checkpointing: bool = True):
        super().__init__()
        self.layers = nn.ModuleList([
            TransformerBlock(d_model, n_heads, ffn_dim) for _ in range(n_layers)
        ])
        self.use_checkpointing = use_checkpointing

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            if self.use_checkpointing:
                # Checkpoint the entire transformer block.
                # use_reentrant=False is recommended for modern PyTorch.
                x = checkpoint(layer, x, use_reentrant=False)
            else:
                x = layer(x)
        return x

Pattern 2: Checkpoint every other block (partial checkpointing)

def forward(self, x: torch.Tensor) -> torch.Tensor:
    for i, layer in enumerate(self.layers):
        # Checkpoint even-indexed layers only: sqrt(n) memory with n/2 recomputes
        if self.use_checkpointing and i % 2 == 0:
            x = checkpoint(layer, x, use_reentrant=False)
        else:
            x = layer(x)
    return x

Pattern 3: torch.compile compatibility

torch.utils.checkpoint is compatible with torch.compile. The compiler understands recomputation semantics and can fuse operations across segment boundaries:

model = Transformer(n_layers=24, d_model=1024, n_heads=16, ffn_dim=4096,
                    use_checkpointing=True)

# torch.compile works with checkpointing; no special flags needed.
compiled_model = torch.compile(model)

6. Interactions and Gotchas ⚠️

6.1 Dropout and RNG State

Dropout is stochastic: on the original forward pass, a random mask is sampled. On the recomputed forward pass during backward, the same mask must be applied — otherwise gradients are computed with respect to a different function than was used to produce the output, silently corrupting gradient estimates.

PyTorch’s checkpointing handles this via RNG state saving:

# Pseudocode of what checkpoint() does internally (simplified):
import torch

def _checkpointed_segment(fn, *args):
    # Before original forward: save RNG state
    rng_state_cpu = torch.get_rng_state()
    rng_state_cuda = torch.cuda.get_rng_state()  # if on CUDA

    # Original forward (under no_grad for reentrant, or normal for non-reentrant)
    output = fn(*args)

    # ... store rng_state_cpu, rng_state_cuda alongside the checkpoint ...

    # During recompute (backward):
    # Restore RNG state so dropout samples the same mask
    torch.set_rng_state(rng_state_cpu)
    torch.cuda.set_rng_state(rng_state_cuda)
    output_recomputed = fn(*args)
    # Restore RNG state to current (post-recompute) state
    # so subsequent random ops are unaffected

This is controlled by the preserve_rng_state argument (default True). Setting preserve_rng_state=False disables RNG saving and restoring — only do this if the checkpointed segment contains no stochastic operations or if determinism is not required.

6.2 Distributed Data Parallel

Gradient checkpointing is safe with Distributed Data Parallel (DDP). Each rank independently performs its forward and backward pass, and checkpointing operates entirely within each rank’s local computation — there is no interaction with the AllReduce communication primitives that DDP uses for gradient synchronization.

model = Transformer(n_layers=24, d_model=1024, n_heads=16, ffn_dim=4096,
                    use_checkpointing=True)
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank])
# DDP + checkpointing: safe, no special configuration needed.

6.3 Mixed Precision and Autocast

PyTorch’s torch.autocast context manager casts operations to lower precision (e.g. bf16 or fp16) based on a dtype dispatch table. When a checkpointed segment is recomputed during backward, it must execute within the same autocast context as the original forward pass — otherwise operations that were bf16 in the forward become fp32 in the recompute, causing shape or dtype mismatches in the gradient computation.

torch.utils.checkpoint handles this automatically: it captures the autocast state (enabled/disabled, dtype) at the start of the original forward pass and restores it during recomputation.

# Correct: autocast wraps the entire forward, including checkpointed regions.
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
    output = model(x)
    loss = criterion(output, target)
loss.backward()
# The recomputed segments during backward() also run under bfloat16 autocast.

6.4 Double Backward

Double backward (computing gradients of gradients, as needed for MAML, higher-order optimization, or Hessian-vector products) requires that the backward pass itself be differentiable — i.e. that the computation graph produced during backward can itself be differentiated.

The reentrant variant (use_reentrant=True) is fundamentally incompatible with double backward: the nested backward pass runs under no_grad, so no graph is recorded for the recomputed activations.

The non-reentrant variant (use_reentrant=False) supports double backward when preserve_rng_state=True (the default), because it records the autograd graph during recomputation via saved-tensor hooks, which are themselves differentiable.

# Double backward: requires use_reentrant=False
x = checkpoint(layer, x, use_reentrant=False, preserve_rng_state=True)

# Compute first-order gradient
grads = torch.autograd.grad(loss, model.parameters(), create_graph=True)

# Compute second-order gradient (Hessian-vector product, etc.)
grad_norm = sum(g.pow(2).sum() for g in grads)
grad_norm.backward()  # This works with use_reentrant=False

References