Gradient Checkpointing (Activation Checkpointing) is the memory optimization technique that trades compute for memory during neural network training by selectively storing only a subset of intermediate activations and recomputing the rest during the backward pass — reducing memory consumption from O(N) to O(√N) for N layers, enabling training of models that would otherwise exceed GPU memory, at the cost of approximately 30-33% additional computation, making it essential infrastructure for training large transformers and deep networks on memory-constrained hardware.

The Memory Problem

Forward pass: Compute and STORE activations for backward pass
  Layer 1: a₁ = f₁(x)       → store a₁ (needed for grad computation)
  Layer 2: a₂ = f₂(a₁)      → store a₂
  ...
  Layer N: aₙ = fₙ(aₙ₋₁)   → store aₙ
  
Memory: O(N) activations stored simultaneously
For Llama-2-7B (32 layers, batch=4, seq=4096): ~60 GB activation memory

How Gradient Checkpointing Works

Without checkpointing (standard):
  Forward: Store ALL activations [a₁, a₂, a₃, ..., a₃₂]
  Backward: Use stored activations to compute gradients
  Memory: 32 × activation_size

With checkpointing (every 4 layers):
  Forward: Store only checkpoints [a₁, a₅, a₉, a₁₃, a₁₇, a₂₁, a₂₅, a₂₉]
  Backward at layer 12:
    Need a₁₂ but it wasn't stored!
    Recompute: a₁₀ = f₁₀(a₉), a₁₁ = f₁₁(a₁₀), a₁₂ = f₁₂(a₁₁)
    Use a₁₂ to compute gradient, then free it
  Memory: 8 checkpoints + 4 recomputed activations = 12 (vs. 32)

Memory-Compute Trade-off

Implementation

import torch
from torch.utils.checkpoint import checkpoint

class TransformerBlock(nn.Module):
    def forward(self, x):
        x = x + self.attention(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x

class Model(nn.Module):
    def forward(self, x):
        for block in self.blocks:
            # Without checkpointing: stores all activations
            # x = block(x)
            
            # With checkpointing: recomputes during backward
            x = checkpoint(block, x, use_reentrant=False)
        return x

# Memory savings for 32-layer model:
# Without: 32 layers of activations
# With: ~6 layers (√32 ≈ 6 checkpoints + recompute buffer)

Selective Checkpointing

• Not all layers consume equal memory.
• Attention: O(N²) memory for attention matrices — checkpoint these!
• FFN: O(N×d) memory — less benefit from checkpointing.
• Strategy: Checkpoint attention (high memory), skip FFN (low memory) → better ratio.

In Practice

Combined with Other Optimizations

Baseline:     Model weights (14 GB) + Activations (60 GB) + Gradients (14 GB) + Optimizer (56 GB)
              = 144 GB → doesn't fit on 80GB GPU

+ Checkpointing: Activations → 20 GB → Total 104 GB → still doesn't fit
+ Mixed precision: Activations in BF16 → 10 GB → Total 94 GB → close
+ DeepSpeed ZeRO-2: Optimizer → 28 GB → Total 66 GB → fits on 80GB!

Gradient checkpointing is the essential memory optimization that makes training large models possible on limited hardware — by accepting a modest ~33% compute overhead in exchange for dramatically reduced activation memory, checkpointing enables researchers and engineers to train models that would otherwise require 2-4× more GPUs, directly reducing the hardware cost and barrier to entry for training state-of-the-art deep learning models.