Gradient checkpointing

Deep dive

How it works

In: a network split into segments, a batch · Out: the same gradients as plain backprop · State: each segment's input (and RNG state) kept until its backward · Cost: +1 forward pass (about +33% FLOPs), O(√n) activation memory.

1. Partition the network into segments (typically one transformer block each).
2. Forward: run each segment under no_grad, keep only its input (and RNG state), and discard the internals.
3. Backward: on reaching segment k, re-run its forward with grad enabled from the saved input, then backprop through the rebuilt local graph with the incoming gradient. Its activations die immediately.

forward:   x0 ─[seg1]─ x1 ─[seg2]─ x2 ─[seg3]─ x3 ─ loss    keep x0..x3 only
backward:                          recompute seg3 → bwd seg3 → free
                    recompute seg2 → bwd seg2 → free
   recompute seg1 → bwd seg1 → free

With n layers and segments of k layers, memory ≈ n/k boundaries + k internals. That is minimized at k = √n, giving 2√n. Recursive checkpointing reaches O(log n) memory at O(n log n) compute (Griewank's binomial "Revolve" schedule), which is rarely worth it on GPUs.

Transformer activation memory per layer in 16-bit precision, no tensor parallelism (Korthikanti et al.), with s = sequence length, b = micro-batch, h = hidden size, a = heads:

bytes ≈ s·b·h·(34 + 5·a·s/h)

For s = 2048, h = 12288, a = 96 (GPT-3 shape), b = 1: 34·sbh ≈ 0.86 GB plus the 5as/h term ≈ 2.0 GB, so 2.9 GB/layer and about 280 GB over 96 layers. The s² term comes from attention scores, softmax and dropout masks. FlashAttention never stores them, which is why "selective" recomputation targets exactly that part.

Code

A chain of n elementwise layers y = tanh(w ⊙ x) with loss sum(y_n). It computes the gradients once with all activations stored and once with √n checkpointing, then compares the results and the costs.

fn layer(w: &[f32], x: &[f32]) -> Vec<f32> { w.iter().zip(x).map(|(w, x)| (w * x).tanh()).collect() }

/// Backward through one layer: writes dL/dw, returns dL/dx.
fn layer_bwd(w: &[f32], x: &[f32], y: &[f32], g: &[f32], dw: &mut [f32]) -> Vec<f32> {
    (0..x.len()).map(|j| { let d = g[j] * (1.0 - y[j] * y[j]); dw[j] = d * x[j]; d * w[j] }).collect()
}

#[derive(Default, Debug)]
struct Cost { layer_forwards: usize, peak_vectors: usize }

fn grads_store_all(ws: &[Vec<f32>], x0: &[f32], c: &mut Cost) -> Vec<Vec<f32>> {
    let mut acts = vec![x0.to_vec()];
    for w in ws { let y = layer(w, acts.last().unwrap()); acts.push(y); c.layer_forwards += 1; }
    c.peak_vectors = acts.len();                                   // n + 1 activations alive
    let (mut g, mut dws) = (vec![1.0; x0.len()], vec![vec![0.0; x0.len()]; ws.len()]);
    for i in (0..ws.len()).rev() { g = layer_bwd(&ws[i], &acts[i], &acts[i + 1], &g, &mut dws[i]); }
    dws
}

fn grads_checkpointed(ws: &[Vec<f32>], x0: &[f32], k: usize, c: &mut Cost) -> Vec<Vec<f32>> {
    let n = ws.len();
    let (mut ckpt, mut x) = (vec![x0.to_vec()], x0.to_vec());     // forward: segment inputs only
    for i in 0..n {
        x = layer(&ws[i], &x);
        c.layer_forwards += 1;
        if (i + 1) % k == 0 && i + 1 < n { ckpt.push(x.clone()); }
    }
    let (mut g, mut dws) = (vec![1.0; x0.len()], vec![vec![0.0; x0.len()]; n]);
    for (s, start) in (0..n).step_by(k).enumerate().rev() {
        let end = (start + k).min(n);
        let mut acts = vec![ckpt[s].clone()];                      // recompute this segment
        for i in start..end { let y = layer(&ws[i], acts.last().unwrap()); acts.push(y); c.layer_forwards += 1; }
        c.peak_vectors = c.peak_vectors.max(ckpt.len() + acts.len());
        for i in (start..end).rev() {
            g = layer_bwd(&ws[i], &acts[i - start], &acts[i - start + 1], &g, &mut dws[i]);
        }                                                           // acts dropped here
    }
    dws
}

fn main() {
    let (n, d) = (64, 1024);
    let ws: Vec<Vec<f32>> = (0..n).map(|i| (0..d).map(|j| 0.5 + ((i * 31 + j * 17) % 97) as f32 / 97.0).collect()).collect();
    let x0: Vec<f32> = (0..d).map(|j| ((j * 13) % 29) as f32 / 29.0 - 0.5).collect();
    let (mut full, mut ck) = (Cost::default(), Cost::default());
    let g_full = grads_store_all(&ws, &x0, &mut full);
    let g_ck = grads_checkpointed(&ws, &x0, (n as f64).sqrt() as usize, &mut ck);
    assert_eq!(g_full, g_ck);                                       // same ops, same order: bitwise equal
    println!("store-all: {full:?}\ncheckpoint: {ck:?}");            // 65 vs 17 vectors; 64 vs 128 forwards
}

The second forward doubles forward work. Backward costs about 2 forwards, so the whole step grows from 3 to 4 forward-equivalents: +33%.

Real frameworks wrap the same idea around autograd: PyTorch torch.utils.checkpoint.checkpoint(block, h, use_reentrant=False) saves the RNG state so dropout masks replay identically, and JAX jax.checkpoint takes a policy such as dots_with_no_batch_dims_saveable.

Complexity and constants

• Full per-block recompute: +1 forward ≈ +33% FLOPs. Activation memory drops from the full per-layer amount to about s·b·h·2 bytes per layer (the block input), plus one block's working set during backward.
• Selective recompute (Megatron --recompute-granularity selective) recomputes only QKᵀ, softmax, dropout and attention-over-V. Megatron reports about 5x less activation memory and over 90% less recompute overhead than full recompute. FlashAttention makes this the default behavior.
• Offloading the checkpoints to CPU over PCIe 5 (about 50 GB/s) can be hidden only if per-layer compute exceeds the transfer time. That holds for large h, not for small models.

Efficiency

• Cost: one extra forward per checkpointed segment: +33% FLOPs, observed slowdown 20-40%. GPT-3 shape (s = 2048, h = 12288, 96 layers, b = 1) stores about 2.9 GB per layer, 280 GB total, without checkpointing (formula above); with it only the block inputs plus one block's working set. Bound: compute traded for memory.
• Levers:
  • checkpoint at block boundaries — O(√n) memory for n layers at +1 forward (Chen et al.) — any deep stack
  • selective recompute of attention internals — about 5x less activation memory, over 90% less recompute overhead than full recompute (Korthikanti et al.) — transformers
  • FlashAttention — never stores the s² tensors, so selective recompute comes built in — attention layers
  • recompute only cheap elementwise ops, save matmul outputs — most of the memory for few FLOPs — policy-based frameworks
  • CPU offload of checkpoints — frees GPU memory without recompute (order of 50 GB/s over PCIe 5) — large h
• Naive → standard → state of the art:
  • store everything · 0% overhead, memory O(n)
  • full block recompute · +33% FLOPs, block inputs only
  • selective recompute or FlashAttention · about 5x less activation memory, a few percent overhead

Implementation notes

• Checkpoint at block boundaries. The block input is small (sbh), and the block internals (QKV, MLP 4h intermediate, attention probabilities) are large.
• Policies beat all-or-nothing. Save matmul outputs and recompute elementwise ops (cheap FLOPs, many bytes). JAX exposes this as checkpoint policies; PyTorch has torch.utils.checkpoint with context_fn / selective activation checkpointing, and torch.compile min-cut rematerialization does it automatically.
• With pipeline parallelism, the first stage holds p micro-batches of activations, so checkpointing there pays off most.
• Recompute repeats any all-gather (FSDP) or all-reduce (tensor parallel) inside the segment. Account for the extra communication.

Pitfalls

• Reentrant checkpoint with an input that does not require grad (e.g. raw embeddings fed from frozen layers): the output does not require grad, and the segment's parameters silently get no gradients. Use use_reentrant=False.
• Dropout or other RNG ops without RNG replay: recomputed masks differ, and the gradients are subtly wrong (no error raised).
• BatchNorm in train mode updates running stats twice under recompute.
• Checkpointing every layer of a cheap model can make it 30% slower for a few hundred MB. Measure peak memory (torch.cuda.max_memory_allocated) before and after.
• Non-deterministic kernels (atomics) make the recomputed activations differ slightly from the originals. This is usually harmless but breaks bitwise reproducibility tests.

Variants and alternatives

In production

torch.utils.checkpoint, FSDP activation-checkpointing wrappers, Megatron-LM (full and selective), DeepSpeed activation checkpointing (partitioned and CPU-offloaded), JAX jax.checkpoint/remat (used throughout MaxText and T5X). Whisper (speech) training used activation checkpointing; Revolve provides checkpoint schedules for the adjoint (reverse) mode of program differentiation.

Implementations

Sources

• Chen, Xu, Zhang, Guestrin, Training Deep Nets with Sublinear Memory Cost, 2016. paper · code
• Griewank and Walther, Algorithm 799: Revolve, ACM TOMS, 2000. paper
• Korthikanti et al., Reducing Activation Recomputation in Large Transformer Models, 2022. paper · code
• Gomez, Ren, Urtasun, Grosse, The Reversible Residual Network, NeurIPS 2017. paper
• Radford et al., Robust Speech Recognition via Large-Scale Weak Supervision (Whisper), 2022 (trained with activation checkpointing). paper