activation checkpoint
**Gradient Checkpointing (Activation Checkpointing)** is the **memory optimization technique that trades compute for memory during neural network training by selectively storing only a subset of intermediate activations and recomputing the rest during the backward pass** — reducing memory consumption from O(N) to O(√N) for N layers, enabling training of models that would otherwise exceed GPU memory, at the cost of approximately 30-33% additional computation, making it essential infrastructure for training large transformers and deep networks on memory-constrained hardware.
**The Memory Problem**
```
Forward pass: Compute and STORE activations for backward pass
Layer 1: a₁ = f₁(x) → store a₁ (needed for grad computation)
Layer 2: a₂ = f₂(a₁) → store a₂
...
Layer N: aₙ = fₙ(aₙ₋₁) → store aₙ
Memory: O(N) activations stored simultaneously
For Llama-2-7B (32 layers, batch=4, seq=4096): ~60 GB activation memory
```
**How Gradient Checkpointing Works**
```
Without checkpointing (standard):
Forward: Store ALL activations [a₁, a₂, a₃, ..., a₃₂]
Backward: Use stored activations to compute gradients
Memory: 32 × activation_size
With checkpointing (every 4 layers):
Forward: Store only checkpoints [a₁, a₅, a₉, a₁₃, a₁₇, a₂₁, a₂₅, a₂₉]
Backward at layer 12:
Need a₁₂ but it wasn't stored!
Recompute: a₁₀ = f₁₀(a₉), a₁₁ = f₁₁(a₁₀), a₁₂ = f₁₂(a₁₁)
Use a₁₂ to compute gradient, then free it
Memory: 8 checkpoints + 4 recomputed activations = 12 (vs. 32)
```
**Memory-Compute Trade-off**
| Strategy | Memory | Extra Compute | When to Use |
|----------|--------|-------------|-------------|
| No checkpointing | O(N) | 0% | Fits in memory |
| Checkpoint every √N layers | O(√N) | ~33% | Standard choice |
| Checkpoint every layer | O(1) per layer | ~100% | Extreme memory limit |
| Selective checkpointing | Variable | 10-30% | Target expensive layers |
**Implementation**
```python
import torch
from torch.utils.checkpoint import checkpoint
class TransformerBlock(nn.Module):
def forward(self, x):
x = x + self.attention(self.norm1(x))
x = x + self.ffn(self.norm2(x))
return x
class Model(nn.Module):
def forward(self, x):
for block in self.blocks:
# Without checkpointing: stores all activations
# x = block(x)
# With checkpointing: recomputes during backward
x = checkpoint(block, x, use_reentrant=False)
return x
# Memory savings for 32-layer model:
# Without: 32 layers of activations
# With: ~6 layers (√32 ≈ 6 checkpoints + recompute buffer)
```
**Selective Checkpointing**
- Not all layers consume equal memory.
- Attention: O(N²) memory for attention matrices — checkpoint these!
- FFN: O(N×d) memory — less benefit from checkpointing.
- Strategy: Checkpoint attention (high memory), skip FFN (low memory) → better ratio.
**In Practice**
| Framework | API | Default Behavior |
|-----------|-----|------------------|
| PyTorch | torch.utils.checkpoint | Manual per module |
| DeepSpeed | activation_checkpointing config | Automatic |
| Megatron-LM | --activations-checkpoint-method | Uniform or selective |
| FSDP | auto_wrap_policy + checkpoint | Integrated |
| HuggingFace | gradient_checkpointing=True | Simple flag |
**Combined with Other Optimizations**
```
Baseline: Model weights (14 GB) + Activations (60 GB) + Gradients (14 GB) + Optimizer (56 GB)
= 144 GB → doesn't fit on 80GB GPU
+ Checkpointing: Activations → 20 GB → Total 104 GB → still doesn't fit
+ Mixed precision: Activations in BF16 → 10 GB → Total 94 GB → close
+ DeepSpeed ZeRO-2: Optimizer → 28 GB → Total 66 GB → fits on 80GB!
```
Gradient checkpointing is **the essential memory optimization that makes training large models possible on limited hardware** — by accepting a modest ~33% compute overhead in exchange for dramatically reduced activation memory, checkpointing enables researchers and engineers to train models that would otherwise require 2-4× more GPUs, directly reducing the hardware cost and barrier to entry for training state-of-the-art deep learning models.