Gradient Checkpointing

Reduce Training Memory with Gradient Checkpointing

Training a neural network can run out of accelerator memory even when the model parameters fit comfortably. The missing piece is often activations: intermediate values produced during the forward pass and retained because backpropagation needs them later. Gradient checkpointing, also called activation checkpointing, trades extra computation for lower activation memory. Instead of keeping every intermediate activation until the backward pass, training keeps selected checkpoints and recomputes missing forward values when their gradients are needed.