Gradient Checkpointing

Overview

Gradient Checkpointing (aka Activation Checkpointing or Recomputation) reduces memory usage by selectively discarding intermediate activations during forward pass and recomputing them during backward pass. Trade-off: Save ~50% memory at cost of ~30% computation overhead.

• Foundational Idea: "Training Deep Nets with Sublinear Memory" (Chen et al., 2016)
• Trade-off: Memory: O(N) → O(√N), Compute: ~1.3x slower training
• Adoption: Standard in transformers (LLaMA, BERT, GPT all use it)
• Key Insight: Recompute is cheaper than storing large activations in GPU memory
• Modern Implementation: PyTorch provides native support

-

The Memory Problem During Training

Where Memory Goes

GPU Memory During Training:

Model Weights: 70B parameters = 140GB (float32)
Optimizer States: 2 × 70B = 140GB (Adam: m and v buffers)
Activations (stored): Depends on depth and batch size
 - Layer 1 output: (batch, seq_len, hidden) = 2GB
 - Layer 2 output: (batch, seq_len, hidden) = 2GB
 -...
 - Layer 80 output: (batch, seq_len, hidden) = 2GB
 - Total: 80 layers × 2GB = 160GB!

Total during training:
 - Weights + Optimizer: 280GB
 - Activations: 160GB
 - Gradients: 140GB
 - Total: 580GB!

Only A100 (80GB) can't fit even one example per GPU!
Problem: Activation storage is huge!

Storage breakdown:
 - Weights/Optimizer: 70% of memory
 - Activations: 30% of memory
 - Most problematic: Activations aren't needed immediately!

Why Store Activations?

Backward pass needs activations:

Forward pass produces:
 - Layer 1: a₁ (activation)
 - Layer 2: a₂ (activation)
 -...
 - Layer 80: a₈₀ (activation)
 - Loss: L

Backward pass computes:
 - ∂L/∂a₈₀ = dL/da₈₀ (computed from forward)
 - ∂L/∂W₈₀ = dL/da₈₀ × ∂a₈₀/∂W₈₀ (needs a₈₀!)
 - ∂L/∂a₇₉ = dL/da₈₀ × ∂a₈₀/∂a₇₉ (needs both a₈₀ and a₇₉!)
 -...

To compute ∂L/∂W at layer i:
 - Need activation a_i from forward pass!

Current approach:
 - Store all activations during forward: 160GB
 - Backward pass uses stored activations
 - Problem: Huge memory for storing!

Question: Can we compute ∂L/∂W without storing activations?
Answer: Yes! Recompute them!

-

Gradient Checkpointing: The Solution

Core Idea

Instead of:
- ┌─────────────────┐ ┌──────────────────┐
 - Forward Pass │ │ Backward Pass │
 - (store a₁..a₈₀) │ → │ (use a₁..a₈₀) │
 - ┘ └──────────────────┘
Memory: 160GB

Do:
- ┌─────────────────┐ ┌──────────────────┐
 - Forward Pass │ │ Recompute │ ┌──────────────────┐
 - (store a₁..a₁₀) │ → │ Forward (a₁..a₈₀)│ → │ Backward Pass │
 - ┘ └──────────────────┘ │ (use recomputed) │
 - Memory: ~16GB Compute: extra └──────────────────┘

Savings:
 - Forward: Store only ~10 activations (1/8 of total)
 - Backward: Recompute ~70 activations (extra compute)
 - Trade: 90% memory saving for 30% compute overhead
 - Net: Better than 90GB + waiting for OOM error!

Checkpointing Strategy

Transformer Block:
- ┌──────────────────────────────┐
 - Input x │
 - ┤
 - Self-Attention │ ← Save output here (checkpoint)
 - ┤
 - Feed-Forward Network (FFN) │
 - ┤
 - Output │ ← Save output here (checkpoint)
 - ┘

Checkpoint placement strategy:
 - Option 1: Every layer (expensive backward recompute)
 - Option 2: Every 2-3 layers (balanced)
 - Option 3: Only attention outputs (reasonable)
 - Typical: Save ~20-25% of activations, recompute 75-80%

---

Implementation

PyTorch Native Support

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

class TransformerBlock(nn.Module):
 def __init__(self, hidden_dim):
 super().__init__()
 self.self_attn = MultiHeadAttention(hidden_dim)
 self.ffn = FeedForward(hidden_dim)
 self.norm1 = nn.LayerNorm(hidden_dim)
 self.norm2 = nn.LayerNorm(hidden_dim)

 def forward(self, x):
 # Without checkpointing: stores all intermediate activations
 # x_attn = self.self_attn(self.norm1(x))
 # out = self.ffn(self.norm2(x + x_attn))
 # return x + out

 # With checkpointing: recompute during backward
 def create_attention_layer(x):
 return self.self_attn(self.norm1(x))

 def create_ffn_layer(x_attn, x):
 return self.ffn(self.norm2(x + x_attn))

 # checkpoint() recomputes during backward, doesn't store activations
 x_attn = checkpoint(create_attention_layer, x, use_reentrant=False)
 out = checkpoint(create_ffn_layer, x_attn, x, use_reentrant=False)
 return x + out

# Usage:
model = nn.Sequential(
 TransformerBlock(2048),
 TransformerBlock(2048),
 #... more blocks
)

# Training loop
x = torch.randn(32, 4096, 2048) # batch, seq_len, hidden_dim
y = model(x)
loss = criterion(y, target)
loss.backward() # Automatically recomputes activations!
optimizer.step()

Manual Checkpointing Control

class FullyShardedTransformer(nn.Module):
 def __init__(self, num_layers, hidden_dim):
 super().__init__()
 self.layers = nn.ModuleList([
 TransformerBlock(hidden_dim)
 for _ in range(num_layers)
])
 self.use_checkpoint = True

 def forward(self, x):
 for i, layer in enumerate(self.layers):
 if self.use_checkpoint:
 # Create wrapper that doesn't require self
 def layer_forward(x, layer=layer):
 return layer(x)

 x = checkpoint(layer_forward, x, use_reentrant=False)
 else:
 x = layer(x)
 return x

# Advanced
class SelectiveCheckpointTransformer(nn.Module):
 def __init__(self, num_layers, hidden_dim, checkpoint_interval=1):
 super().__init__()
 self.layers = nn.ModuleList([
 TransformerBlock(hidden_dim)
 for _ in range(num_layers)
])
 self.checkpoint_interval = checkpoint_interval

 def forward(self, x):
 for i, layer in enumerate(self.layers):
 # Checkpoint every N layers
 if (i % self.checkpoint_interval) == 0:
 def layer_forward(x, layer=layer):
 return layer(x)
 x = checkpoint(layer_forward, x, use_reentrant=False)
 else:
 x = layer(x)
 return x

-

The Trade-off Analysis

Memory Savings

Scenario: 70B model, batch_size=1, seq_len=4096

Without checkpointing:
 - Layer activations: 80 layers × 2GB = 160GB
 - Peak memory: 280GB (weights + optimizer) + 160GB = 440GB
 - Doesn't fit on A100 (80GB)!

With checkpointing (every layer):
 - Stored activations: ~1 layer (2GB)
 - Peak memory: 280GB + 2GB = 282GB (still doesn't fit with optimizer!)

With checkpointing + FSDP (8 GPUs):
 - Per GPU parameters: 140GB / 8 = 17.5GB
 - Per GPU optimizer: 140GB / 8 = 17.5GB
 - Per GPU activations: 2GB (checkpoint)
 - Total per GPU: ~37GB (fits on A100!)
 - Solution: FSDP + checkpointing enables large model training!

Compute Overhead

Computation breakdown (per layer):

Forward pass (one-time): F
Backward pass without checkpoint: B (uses stored activation)
Backward pass with checkpoint: F + B (recompute + backward)

Total compute:
 - Without checkpoint: F + B per layer = (F + B) × 80 layers
 - With checkpoint: (F + B + F) × 80 = (2F + B) × 80 layers
 - Overhead: F / (F + B) = typically 30-40%

Typical F/B ratio for transformers: F:B ≈ 1:2 (backward heavier)
 - Overhead = 1/3 ≈ 33% extra compute

Actual wall-clock time:
 - Without checkpoint: Faster per batch, but OOMs on large models
 - With checkpoint: 30-40% slower per batch, fits large models
 - Net-net: Can train 10-100x larger models with checkpoint
 - Trade is worthwhile!

Real example (GPT-3 70B):
 - Without checkpoint: ~120 hrs training, OOMs
 - With checkpoint: ~150 hrs training (1.25x), fits!
 - Can add more data, more GPUs, get to production

---

Advanced Techniques

Partial Checkpointing

Instead of checkpointing all layers, only checkpoint expensive ones:

Cost of operations (relative):
 - Self-attention: 1x (attention is relatively fast)
 - FFN: 2-3x (FFN is compute-intensive)
 - LayerNorm: 0.1x (super cheap)

Strategy:
 - Checkpoint FFN layers
 - Don't checkpoint attention layers
 - Save 50% activations, but keep cheap operations stored
 - Balances memory and compute better

Selective Checkpointing

class SmartCheckpointTransformer(nn.Module):
 def forward(self, x):
 for i, layer in enumerate(self.layers):
 # Only checkpoint expensive FFN layer
 if isinstance(layer.ffn, ExpensiveFFN):
 def ffn_checkpoint(x):
 return layer.ffn(x)
 ffn_out = checkpoint(ffn_checkpoint, layer.norm1(x))
 x = x + ffn_out
 else:
 # Cheap layers: no checkpoint
 x = layer(x)
 return x

Reentrant vs Non-reentrant Checkpointing

# Traditional (reentrant=True):
checkpoint(func, x, use_reentrant=True)

Pros:
Lower memory overhead
Compatible with older PyTorch versions

Cons:
Can cause issues with certain operations
Requires careful implementation

# Modern (reentrant=False):
checkpoint(func, x, use_reentrant=False)

Pros:
More stable with complex models
Better compatibility
PyTorch 1.11+

Cons:
Slightly higher memory overhead
Requires newer PyTorch

---

When to Use Gradient Checkpointing

Use checkpointing when

Model is large (>1B parameters)
Batch size is small (can't use more GPUs easily)
Memory is bottleneck (OOM errors)
Sequence length is long (more activations to store)
Can afford 30-40% compute overhead

Don't use checkpointing when

Memory is plentiful (small model + large GPUs)
Training speed is critical
Operations don't compose well (RNGs, special functions)

---

Key Takeaways

Trade 30% compute for 50-80% memory savings Essential for training LLMs on GPUs with limited memory Checkpointing enables 10-100x larger models to train Modern PyTorch makes it easy: just wrap function in checkpoint() Widely adopted: LLaMA, BERT, GPT all use it

-

Related Notes

• Distributed Training - Use together for large-scale training
• Mixed Precision Training (Amp) - Further memory reduction
• Lora - Alternative to full model training
• Llm Inference Optimization - Memory during inference