Gradient Checkpointing: How Does It Save GPU Memory?
Gradient checkpointing saves fewer intermediate activations during the forward pass and recomputes missing ones during backward computation. Use it when activations prevent the training batch or sequence length from fitting in GPU memory, and the extra computation is acceptable.
It reduces saved activation memory. It does not remove model weights, optimizer state, or every temporary buffer.
What gets recomputed?
Backward computation needs intermediate results from the forward pass to calculate gradients. Without checkpointing, the framework retains the required results. With checkpointing, it saves selected inputs or intermediate results and reruns the affected operations when their values are needed.
For example, a framework can save the input to a group of transformer layers and discard their internal activations. During backward computation it reruns that group from the saved input. The saved input is an activation checkpoint, which differs from a model checkpoint written to disk for recovery.
Chen et al.’s original analysis shows an evenly segmented strategy with O(sqrt(n)) activation memory for an n-layer network and approximately one extra forward pass per training step. Real transformers have unequal activation sizes, attention implementations, and framework policies. Measure the resulting memory and time rather than copying the asymptotic estimate as a deployment forecast.
Enable the integration your stack expects
For supported Hugging Face Trainer models, gradient_checkpointing=True enables the general mechanism. With FSDP, check the framework’s activation-checkpointing integration instead. The current Transformers Trainer documentation recommends activation_checkpointing for FSDP because generic gradient checkpointing can cause an extra all-gather during backward computation.
Compare the same training configuration with and without checkpointing. Record peak GPU allocation, examples or tokens per second, and whether the intended sequence length fits. Also check that gradients remain valid when your model uses custom operations or adapter training.
If optimizer state dominates memory, consider state sharding or parameter-efficient tuning instead; recomputing activations will not remove that allocation.
The gradient-checkpointing section of the LLM Engineering Guide explains the segmented calculation.