Gradient Checkpointing
Resources
- Arsh’s summer project: https://blog.janestreet.com/trading-off-compute-for-memory-with-activation-checkpointing/
- https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=activation_recomputation
Gradient checkpointing trades compute for memory by recomputing activations on the backward pass instead of storing them from the forward pass.
Why is this needed?
The backward pass needs every layer’s activation to compute its gradient. Storing them all is what makes activation memory blow up as models and sequence lengths grow, often more than the weights themselves.
Mechanics
- Mark certain layers as checkpoints, save their activations
- Discard activations between checkpoints
- On backward, re-run the forward pass between checkpoints to reconstruct what was discarded, then compute the gradient
Net effect: roughly halves activation memory, costs ~20% more time (one extra forward pass over the non-checkpointed segments).
- I have low confidence