Gradient Accumulation & Checkpointing

Overcoming the VRAM Memory Wall

Training large neural networks (like 70B LLMs or high-resolution vision models) hits a hard hardware barrier: GPU VRAM Memory Limits.

GPU VRAM is consumed by four main components:

  TOTAL VRAM = Model Parameters + Optimizer States + Gradients + ACTIVATION TENSORS!

For long sequence lengths (e.g. 8,000 tokens), Activation Tensors (intermediate outputs stored for backpropagation) consume up to $80%$ of total GPU VRAM!

Gradient Accumulation and Activation Checkpointing allow training massive models on limited hardware.

┌──────────────────────────┬──────────────────────────┐
│ 1. GRADIENT ACCUMULATION │ 2. ACTIVATION CHECKPOINT │
├──────────────────────────┼──────────────────────────┤
│ Simulates large global   │ Discards intermediate    │
│ batch sizes by running   │ activation tensors in    │
│ N micro-batches before   │ forward pass. Recomputes │
│ calling optimizer.step().│ them in backward pass!   │
│ ZERO extra compute cost. │ Trades +20% Compute Time │
│ Saves Batch Memory!      │ to Save -70% VRAM Memory!│
└──────────────────────────┴──────────────────────────┘

1. Gradient Accumulation

Suppose your model requires a Effective Global Batch Size of $256$ for stable training convergence.

However, your GPU crashes with an Out of Memory (OOM) error if mini-batch size exceeds $32$ samples.

How Gradient Accumulation Works

Instead of running a single batch of $256$, run $8$ micro-batches of $32$ samples each:

effective_batch_size = 256
micro_batch_size = 32
accumulation_steps = 8 # 256 / 32 = 8 steps

for i, (inputs, targets) in enumerate(dataloader):
    outputs = model(inputs)
    loss = criterion(outputs, targets) / accumulation_steps # Scale loss!
    loss.backward() # Accumulates gradients into param.grad!
    
    if (i + 1) % accumulation_steps == 0:
        optimizer.step() # Takes single optimizer step after 8 micro-batches!
        optimizer.zero_grad() # Reset accumulated gradients!

• Benefit: Trains with large effective batch sizes on small single GPUs.
• Cost: Zero added FLOP compute overhead!

2. Activation Checkpointing (Gradient Checkpointing)

During a standard forward pass, PyTorch stores activation tensors for every single layer in VRAM so backpropagation can compute $\frac{\partial L}{\partial a_l}$.

For a 100-layer Transformer, storing activations for all 100 layers consumes tens of gigabytes of VRAM.

How Activation Checkpointing Works (Chen et al., 2016)

Instead of saving activations for all 100 layers:

1. Save activations ONLY for a few Checkpoint Layers (e.g. every 10th layer). Discard intermediate activations!
2. During the backward pass, when backpropagation reaches layer 15, re-run a local forward pass from Checkpoint Layer 10 to 15 to recompute missing activations on the fly!

  STANDARD FORWARD:      [ L1 ] ──► [ L2 ] ──► [ L3 ] ──► [ L4 ]  (Store ALL Activations in VRAM!)
  CHECKPOINTED FORWARD:  [ L1* ] ──► [ L2 ] ──► [ L3 ] ──► [ L4* ] (Store ONLY Checkpoints L1* and L4*!)
                                              ▲
                                              └─ Recompute L2 & L3 on demand during Backward Pass!

• Benefit: Reduces activation VRAM memory footprint by up to $75%$.
• Cost: Adds $\sim 20%$ extra compute time due to re-running partial forward passes during backpropagation.

Say this out loud

Gradient Accumulation simulates large batch sizes by accumulating gradients over N micro batches before calling optimizer step, saving peak batch memory at zero compute cost. Activation Checkpointing trades compute for memory by discarding intermediate activations during the forward pass and recomputing them on demand during backpropagation, saving up to 75 percent activation VRAM.

Followups to expect

1. What is Selective Activation Checkpointing? Saving high-memory, low-compute activations (like Dropout masks and Attention Softmax probabilities) while recomputing low-memory, high-compute matrix multiplications.
2. How does FlashAttention complement Activation Checkpointing? FlashAttention avoids storing $O(N^2)$ intermediate attention matrix activations altogether by computing attention using tiled GPU SRAM blocks, reducing activation memory dramatically.