Gradient Checkpointing: Trading Compute for Memory

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: Selective checkpointing
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