Neural networks, deep learning, NLP · insight-0470
Trade compute for memory with activation checkpointing before shrinking the model
- Evidence type
- documented fact
- Difficulty
- medium
- Roles
- ML engineer
- Industries
- cross-industry
Summary: situation, action, reason, boundary
Situation
Training runs out of GPU memory as you increase sequence length or batch size, and the plan is to switch to a smaller model.
Action
Measure peak memory, then enable activation checkpointing on the largest blocks and use gradient accumulation to keep the effective batch size. Compare throughput and validation metrics with the original setup.
Reason
Checkpointing stores only some activations and recomputes the rest during backward, cutting activation memory for roughly one extra forward pass.
Boundary
Deep CNNs and Transformers limited by activation memory.
Symptom or risk
A smaller model may lose accuracy you need, when the actual constraint is stored activations for the backward pass.
Exceptions
When model states, not activations, dominate memory, checkpointing helps little; see the 16-bytes-per-parameter budget.
Trade-offs
Training slows because of recomputation; tune which blocks to checkpoint.
Pitfalls
Disabling RNG-state preservation (preserve_rng_state=False) or hand-rolling recomputation in blocks that contain dropout, so the recomputed forward pass differs from the original one.
Lessons and concepts
Claims and their sources
Chen et al. reduce training memory to O(sqrt(n)) for an n-layer network at the cost of about one extra forward pass per minibatch.
PyTorch describes activation checkpointing as trading compute for memory by not keeping forward tensors alive until backward and recomputing them.