Training Deep Nets with Sublinear Memory Cost (Activation Checkpointing)

Editorial record

Plain-language summary

During the forward pass only a sparse set of checkpoint activations is kept and the rest are dropped; during backpropagation the missing activations are recomputed on demand from the nearest checkpoint. Placing checkpoints roughly every sqrt(n) layers reduces activation memory to O(sqrt(n)) at the cost of one extra forward computation. This gradient/activation checkpointing let much deeper networks and longer sequences be trained within a fixed GPU memory budget, and became a standard building block in later large-model training systems.