Activation Optimization: Checkpointing, Offloading & Recomputation

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Memory Breakdown in Training
5. Gradient/Activation Checkpointing
6. Selective Recomputation
7. CPU & NVMe Offloading
8. Flash Attention Memory Savings
9. Memory-Compute Tradeoff Analysis
10. Inference Memory Optimization
11. Common Pitfalls
12. Summary
13. Exercises

Learning Objectives

1. Decompose GPU memory usage into model, optimizer, gradient, and activation components.
2. Derive the memory savings of activation checkpointing (sqrt reduction).
3. Design selective recomputation strategies (which layers to checkpoint).
4. Analyze offloading bandwidth requirements and performance impact.
5. Explain Flash Attention's O(N) memory (vs O(N²) for standard attention).

Notation

• $L$ — number of layers
• $B$ — batch size, $S$ — sequence length, $d$ — hidden dimension
• $A_l$ — activations at layer $l$

Core Intuition

Training a transformer requires storing activations from the FORWARD pass to use during the BACKWARD pass (for computing gradients). For a 70B model with sequence length 4096 and batch 32: activations can consume 400GB+ — far more than model weights. Activation checkpointing trades COMPUTE for MEMORY: don't store activations; recompute them during backward. This enables training models that wouldn't fit in memory otherwise.

Memory Breakdown in Training

For a transformer model with $N$ parameters:

Total: $12N + \text{activations}$.

Example (7B model, B=8, S=4096, d=4096, L=32):

• Parameters + optimizer: $12 \times 7B = 84$ GB.
• Activations (no checkpointing): $\sim 120$ GB.
• Total: 204 GB (doesn't fit on single A100 80GB).

Gradient/Activation Checkpointing

Standard: Store all activations during forward; use them in backward.

• Memory: $O(L \times B \times S \times d)$.

Full checkpointing: Store only input activations at "checkpoint" boundaries. Recompute intermediate activations during backward.

Optimal strategy: Checkpoint every $\sqrt{L}$ layers:

• Memory: $O(\sqrt{L} \times B \times S \times d)$.
• Extra compute: one additional forward pass (33% overhead).

Per-layer checkpointing (most common): Checkpoint at every transformer layer boundary:

• Don't store intermediate activations within a layer (attention scores, FFN intermediates).
• Recompute them during backward pass of that layer.
• Memory savings: 5-10x reduction in activation memory.
• Compute overhead: 33% more forward compute.

Selective Recomputation

Not all activations are equal. Some are cheap to recompute; others are expensive.

What to checkpoint (cheap to recompute):

• Layer normalization outputs (very cheap).
• Dropout masks (just re-sample).
• Linear projections (moderate cost).

What NOT to checkpoint (expensive):

• Attention scores ($O(S^2)$ compute).
• FFN intermediate activations (if using Flash Attention, attention is already efficient).

Flash Attention + selective recompute: Flash Attention doesn't materialize the $S \times S$ attention matrix — saving $O(S^2)$ memory. Only need to store the layer input; everything else can be recomputed in the fused kernel.

CPU & NVMe Offloading

When GPU memory is insufficient even with checkpointing:

CPU offloading (ZeRO-Offload):

• Store optimizer states on CPU RAM (cheap, large capacity).
• Transfer gradients to CPU; update on CPU; transfer updated params back.
• Bandwidth: PCIe 4.0 = 32 GB/s. For 7B model: 14GB transfer per step ≈ 0.4s overhead.

NVMe offloading (ZeRO-Infinity):

• Store parameters AND optimizer on NVMe SSD.
• NVMe bandwidth: 6-12 GB/s per drive (RAID for more).
• Enables training models 10x larger than GPU memory.
• Significant slowdown (2-5x).

When to use:

• CPU offload: Model fits in GPU for forward/backward, but optimizer doesn't.
• NVMe offload: Nothing fits. Last resort (very slow).

Flash Attention Memory Savings

Standard attention memory: $O(B \times H \times S^2)$ — store full attention matrix.

Flash Attention memory: $O(B \times H \times S)$ — only store $O(S)$ running statistics.

Mechanism: Tile-based computation:

1. Load blocks of Q, K, V to SRAM (fast on-chip memory).
2. Compute partial attention within each tile.
3. Accumulate using online softmax (numerically stable running stats).
4. Never materialize full $S \times S$ matrix in HBM.

Memory savings example (S=4096, H=32, B=8, FP16):

• Standard: $8 \times 32 \times 4096^2 \times 2 = 8.6$ GB.
• Flash: $O(S) \approx$ negligible.
• Savings: 8.6GB of activation memory eliminated.

Combined with checkpointing: Flash Attention makes activation checkpointing nearly free — the expensive part (attention recomputation) is already fast in the fused kernel.

Memory-Compute Tradeoff Analysis

Without checkpointing: Maximum memory, minimum compute. Full checkpointing: Minimum memory ($\sqrt{L}$), maximum compute (+33%). Selective checkpointing: Sweet spot (checkpoint expensive layers only).

Decision framework:

• GPU memory sufficient? → No checkpointing (fastest).
• Barely fits? → Selective checkpointing (minimal overhead).
• Doesn't fit? → Full checkpointing (33% slower).
• Still doesn't fit? → Checkpointing + CPU offloading.
• Model too large for GPU? → ZeRO-3/FSDP + offloading.

Effective throughput:

Inference Memory Optimization

During inference (no gradients/optimizer):

• Model parameters: $2N$ bytes (BF16).
• KV-cache: dominant memory user (covered in KV-cache optimization).
• Activations: only current layer needed (no storing for backward).

Inference activation memory: $O(B \times S \times d)$ — just the current hidden state.

Optimization: Process long prompts in chunks (chunked prefill) to limit peak activation memory.

Common Pitfalls

Pitfall 1. Enabling activation checkpointing everywhere without measuring benefit. Some models (small ones) fit in memory without checkpointing. The 33% compute overhead is wasted.

Pitfall 2. CPU offloading without overlapping communication. If offloading is synchronous: each step waits for PCIe transfer. Overlap with next batch's forward pass for much better throughput.

Pitfall 3. Not using Flash Attention when sequence length is long. At S=4096+, the $O(S^2)$ attention memory dominates all other activations. Flash Attention is nearly mandatory.

Summary

• Activations often dominate training memory (more than weights + optimizer).
• Checkpointing: Trade 33% compute for 5-10x activation memory reduction.
• Selective recomputation: Checkpoint cheap-to-recompute ops; keep expensive ones.
• Flash Attention: $O(S)$ memory instead of $O(S^2)$; enables long sequences.
• CPU/NVMe offloading: When GPU memory is exhausted; significant slowdown.
• Decision: No checkpoint → selective → full → offloading (from fastest to most memory-efficient).

Exercises

Exercise 1. For LLaMA-13B (L=40, d=5120, B=4, S=4096): compute activation memory with and without checkpointing.

Exercise 2. Derive the optimal checkpointing interval: every $\sqrt{L}$ layers. Why is this the minimum memory solution?

Exercise 3. For a 70B model training with CPU offloading over PCIe 5.0 (64 GB/s): compute the communication overhead per step.

Exercise 4. Compare the total memory of: (a) standard attention + no checkpointing, (b) Flash Attention + full checkpointing, for S=8192 and a 7B model.

Exercise 5. Design a selective recomputation strategy for a GPT-style model: which operations to checkpoint and which to store. Estimate memory savings and compute overhead.