What is activation checkpointing and what is the trade-off?

Activation checkpointing is a memory optimization technique used during training of deep neural networks. Instead of saving all intermediate activations for the backward pass, only certain checkpoints are saved. The remaining activations are recomputed during the backward pass as needed.

Trade-off:

• Pros: Significant reduction in memory consumption, allowing for training larger models or using bigger batches.
• Cons: Increased training time due to additional computation of activations during the backward pass.

Example: in PyTorch, there is torch.utils.checkpoint, which allows wrapping parts of a model for checkpointing.