layer-wise checkpointing
**Layer-wise Activation Checkpointing** is a **memory optimization technique that treats each transformer block as a checkpoint boundary, saving activations at layer boundaries and recomputing within-layer activations during the backward pass** — providing a simple, tunable knob where adjusting the checkpoint frequency (every 1, 2, or 4 layers) directly controls the tradeoff between memory savings and recomputation overhead, making it the most widely used memory reduction technique for training large transformer models.
**What Is Layer-wise Checkpointing?**
- **Definition**: A gradient checkpointing strategy that saves the input activations at transformer layer boundaries and discards all intermediate activations within each layer — during the backward pass, the forward computation within each checkpointed layer is re-executed to regenerate the needed activations for gradient computation.
- **The Tradeoff**: Without checkpointing, all activations are saved (maximum memory, zero recomputation). With checkpointing every layer, only layer inputs are saved (minimum memory, maximum recomputation ~33% overhead). Checkpointing every N layers provides intermediate tradeoffs.
- **Natural Boundaries**: Transformer layers are ideal checkpoint units — each layer has clean input/output interfaces, self-contained forward computation, and well-defined gradient flow, making them natural points to save and restore state.
- **Tunable Frequency**: The checkpoint interval is the primary tuning parameter — checkpoint every 1 layer for maximum memory savings, every 2 layers for balanced performance, or every 4 layers for minimal speed impact.
**Checkpoint Frequency Tradeoffs**
| Frequency | Memory Usage | Speed Overhead | Best For |
|-----------|-------------|---------------|----------|
| No checkpointing | 100% (baseline) | 0% | Small models that fit in memory |
| Every 4 layers | ~70% | ~10% | Moderate memory pressure |
| Every 2 layers | ~50% | ~20% | Balanced speed/memory |
| Every 1 layer | ~30% | ~30% | Maximum memory savings |
| Selective (per-op) | ~50% | ~10-15% | Optimal but complex |
**Implementation**
- **PyTorch**: `torch.utils.checkpoint.checkpoint(layer, input)` wraps each transformer layer — the forward pass runs normally but activations are not saved; during backward, the forward is re-executed within a no-grad context to regenerate activations.
- **Hugging Face Transformers**: `model.gradient_checkpointing_enable()` activates layer-wise checkpointing for any supported model — a single method call that reduces memory by ~50% with ~20% training slowdown.
- **DeepSpeed**: Integrates checkpointing with ZeRO stages — combining activation checkpointing with optimizer state partitioning for maximum memory efficiency.
- **Megatron-LM**: Uses layer-wise checkpointing as the baseline, with selective recomputation as an advanced option for further optimization.
**Why Layer Boundaries Work**
- **Clean Interfaces**: Each transformer layer takes a hidden state tensor and returns a hidden state tensor — the checkpoint only needs to save this single tensor per layer boundary.
- **Efficient Recomputation**: Within-layer operations (attention, FFN, normalization) are computationally cheap relative to the memory they consume — recomputing them is fast.
- **Composable with Other Techniques**: Layer-wise checkpointing combines with tensor parallelism, pipeline parallelism, and ZeRO optimizer sharding — each technique addresses a different memory bottleneck.
**Layer-wise activation checkpointing is the standard memory optimization for large model training** — providing a simple, tunable checkpoint frequency that directly controls the speed-memory tradeoff at natural transformer layer boundaries, enabling training of models 2-3× larger than available GPU memory would otherwise allow.