Activation checkpointing saves only selected intermediate activations during the forward pass. During backward, it recomputes the discarded activations when they are needed instead of storing every layer's output.

This reduces memory at the cost of extra forward computation. More checkpoints use more memory but require less recomputation. Fewer checkpoints save more memory but make backward slower.

Selective activation checkpointing

Selective activation checkpointing applies checkpointing only to chosen operations or regions. Expensive-to-store activations can be saved, while cheap-to-recompute activations are discarded and recomputed during backward. This gives a tunable middle ground between keeping every activation and checkpointing every layer, balancing memory savings against extra computation.

A useful rule is to checkpoint activations that occupy a lot of memory but come from operations that are quick to recompute. Avoid checkpointing expensive operations when their recomputation cost would dominate the saved memory.

A practical approach is to profile the graph during a forward pass, recording memory use and FLOPs for each operation. This makes it easier to find checkpoints with a favorable memory-to-recomputation trade-off.

References

• Current and New Activation Checkpointing Techniques in PyTorch
• PyTorch Memory Tuning