What is gradient checkpointing?
A training trick that saves GPU memory by throwing away most in-between results on the forward pass and recomputing them when the backward pass needs them.
intermediate · Also called Activation Checkpointing, Activation Recomputation
You load a model for fine-tuning, it fits on the GPU, and the first training step still fails with "CUDA out of memory". Often the weights are not the problem. The in-between results of the forward pass are. Gradient checkpointing is the standard fix: keep fewer of them, and work the rest out again when needed.
What eats the memory
Training has 2 passes. The forward pass runs the input through every layer. The backward pass (BackpropagationBackpropagationAn algorithm for computing gradients in deep networks via the chain rule.) then goes back through the layers to compute the gradientsGradientThe vector of partial derivatives of the loss function with respect to model parameters, used to update weights during training., and at each layer it needs that layer's input from the forward pass. So by default every layer's output, its "activations", is kept until the backward pass reaches it.
Activations grow with batch size and with sequence length. With long sequences they can take more memory than the model itself, and they are why a model that runs fine for answering questions runs out of memory when you try to train it.
The trade
Gradient checkpointing keeps only some activations, at checkpoints, say the input to every few layers, and drops the rest. When the backward pass needs a dropped one, it reruns the forward pass from the nearest checkpoint to get it back.
A 2016 paper by Tianqi Chen and others showed that by keeping checkpoints about every square root of the number of layers, memory for activations drops from growing with the number of layers to growing with its square root. For 100 layers, that is roughly 10 layers' worth instead of 100. The cost is about one extra forward pass per step, so each step is slower, but the memory saved often allows a larger batch or a longer context.
How to switch it on
In Hugging Face Transformers it is one line, model.gradient_checkpointing_enable(), or gradient_checkpointing=True in the training arguments. In plain PyTorch, torch.utils.checkpoint wraps any part of a model. Most fine-tuning recipes for large models, including QLoRAQLoRALoRA trained on a copy of the base model stored in 4 bits, so fine-tuning a large model needs a fraction of the GPU memory, with 3 tricks that keep the quality of 16-bit fine-tuning., turn it on by default.
When to use it
Turn it on when training runs out of memory on activations: long sequences, large batches, or a big model on a small GPU. Leave it off when memory is not a problem, since it only costs time then. It matters only for training. Running a model has no backward pass, so there is nothing to recompute.
Sources