Gradient Checkpointing & Activation Recomputation

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. The Activation Memory Problem
5. Basic Checkpointing
6. Optimal Checkpoint Placement
7. Selective Recomputation
8. Memory Savings Analysis
9. Integration with Parallelism
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Compute activation memory for a transformer model.
2. Derive the memory-compute tradeoff for checkpointing.
3. Prove that $\sqrt{L}$ checkpoints minimize memory for a chain of $L$ layers.
4. Explain selective recomputation and which operations to recompute.
5. Analyze checkpointing interaction with pipeline parallelism.

Notation

• $L$ — number of layers
• $a$ — activation memory per layer
• $C$ — number of checkpoints
• $T_f$ — forward time per layer

Core Intuition

Training requires storing intermediate activations from the forward pass to compute gradients in the backward pass. For a 70B model processing 4K tokens: activation memory can reach 100+ GB — often more than the model weights. Checkpointing drops some activations and recomputes them when needed during backward, trading ~33% more compute for dramatically less memory.

The Activation Memory Problem

Standard training: Store all $L$ layers' activations:

For a transformer with $T=4096, d=4096, L=32$:

• Per-layer activations: $\sim T \times d \times 4$ (attention + FFN intermediate) $\approx 256$ MB.
• Total: $32 \times 256$ MB $= 8$ GB per sequence.
• With batch size 4: 32 GB just for activations.

Basic Checkpointing

Strategy: Only save activations at $C$ evenly-spaced checkpoints. During backward, recompute from the nearest checkpoint.

Memory: $C$ checkpointed activations + 1 segment of recomputed activations:

(Store $C$ checkpoints + at most $L/C$ layers between checkpoints during recomputation.)

Compute overhead: Re-run forward pass for each segment: ~33% more total compute (each layer's forward is computed twice: once in the original forward, once during backward recomputation — but the second time only for the segment being processed).

Optimal Checkpoint Placement

Minimize $f(C) = C + L/C$ (memory expression without constant $a$).

Minimum memory: $2\sqrt{L} \times a$.

For $L=32$: $C = 5$–$6$ checkpoints, memory $\approx 11a$ (vs $32a$ without checkpointing). 3x reduction.

Selective Recomputation

Not all operations are equal. Some activations are cheap to recompute and expensive to store:

Always recompute (cheap compute, large memory):

• Attention matrices ($O(T^2)$ memory but $O(T^2d)$ compute — actually, with FlashAttention, they're never stored anyway).
• Dropout masks (random, trivially regeneratable from seed).

Always store (expensive compute, small memory):

• LayerNorm statistics (mean, variance): tiny memory, moderate compute.
• Linear projection outputs: moderate memory, expensive to recompute.

Selective strategy: Checkpoint between transformer blocks (store input to each block); recompute attention + FFN within blocks. Memory: 1 activation per layer ($T \times d$) instead of all intermediates.

Memory Savings Analysis

Without checkpointing:

• Activations per layer: $\sim 10Td$ (attention QKV, attention output, FFN up, FFN down, residuals).
• Total: $10LTd$.

With checkpointing per block:

• Store: $LTd$ (one hidden state per layer).
• Recompute: intermediates within each block during backward.
• Savings: $\sim 10\times$ reduction.

With FlashAttention + checkpointing:

• Don't store $T \times T$ attention matrix (FlashAttention handles this).
• Don't store intermediate FFN activations (recomputed).
• Store only: layer inputs + LayerNorm stats.
• Near-optimal memory usage.

Integration with Parallelism

Pipeline parallelism: Each stage must store activations for $P$ micro-batches (1F1B schedule). Checkpointing reduces each micro-batch's activation footprint.

Tensor parallelism: Activations are split across TP ranks. Checkpointing saves the split activations and recomputes when needed (each rank recomputes independently).

Combined savings: For $L=32, P=4, M=8$:

• Without checkpointing: $10LTd \times P = 10 \times 32 \times T \times d \times 4$ (per micro-batch × stages in memory).
• With checkpointing: $LTd \times P = 32 \times T \times d \times 4$. 10x less.

Common Pitfalls

Pitfall 1. Checkpointing too frequently. Every checkpoint → no recomputation → wastes memory on the checkpoint storage itself.

Pitfall 2. Not using deterministic dropout. If dropout masks aren't reproducible during recomputation, gradients will be incorrect. Save the RNG state at each checkpoint.

Pitfall 3. Checkpointing with mixed precision incorrectly. FP16 activations saved at checkpoints must be used (not recomputed in FP32) to maintain gradient consistency.

Summary

• Activation memory grows linearly with layers and sequence length.
• Checkpointing trades ~33% extra compute for $O(\sqrt{L})$ memory.
• Optimal: $\sqrt{L}$ checkpoints for $O(\sqrt{L})$ memory.
• Selective recomputation: Recompute cheap operations (attention, dropout); store expensive ones.
• Essential for training large models — used universally with Flash Attention.

Exercises

Exercise 1. For $L=48$ layers: compute optimal $C$, memory with and without checkpointing.

Exercise 2. Derive the exact compute overhead of checkpointing (show why it's 33% for uniform segments).

Exercise 3. For a 70B model with $T=8192$: compute activation memory with and without checkpointing + FlashAttention.

Exercise 4. Design a selective recomputation strategy that minimizes memory while adding < 10% compute overhead.

Exercise 5. Explain why checkpointing is especially important for pipeline parallelism (hint: $P$ micro-batches in flight).