Unit 8.6 · about 4 h
226 / 324Activation checkpointing and recomputation
Checkpointing converts the activation term of the memory model into extra compute. Done blindly it slows training for no reason; done deliberately it is what lets larger batches fit.
- implement
- measure
You can now
- Wrap a segment with
torch.utils.checkpointand state exactly what stays alive and what is recomputed. - Choose
use_reentrantFalsewith reasons, and pass it explicitly. - Measure the memory saving against the recompute cost and the RNG effects.
Covers
- Activation checkpointing and recomputation
Study, in this order
2 steps
- `torch.utils.checkpoint` APIPyTorch - checkpointing, numerical accuracy and CUDA semantics
checkpointsignature and semantics: onlyargsstay alive, region outputs are recomputed in backward;use_reentrantFalseversusTruedifferences;preserve_rng_statewith the multi-device caveat;determinism_checkanddebug - Which layers to checkpointPyTorch - reliability and performance recipes
General optimizations, Checkpoint intermediate buffers: large layer outputs with small recompute cost as targets
Focus on
- Only the checkpoint inputs stay alive; everything the region computes is rerun in backward. Memory saved is the activations the region would have kept, and the price is one extra forward of the region during backward.
- Non-reentrant is the recommended path: it supports grad, keyword args and nested structures, and it stops recomputation early once the needed activations exist.
- RNG state is stashed and restored by default so dropout masks match. Turning that off changes gradients, not forward outputs.
- Recompute reruns the region code. Globals read inside the region, device moves, and multi-device RNG all break the equivalence silently or loudly.
Reproduce closed-book
use_reentrant=False and show its gradients equal the plain block at the same seed (torch.equal held on CPU with PyTorch 2.14). Rerun with preserve_rng_state=False and show the gradients differ while the forward outputs still match.Traps
- Omitting
use_reentrantand assuming a default. Current PyTorch warns that it must be passed explicitly; passFalsedeliberately. - Wrapping tiny cheap regions. Recompute then costs more than the memory it saves; target large outputs with small recompute cost.
- Changing code or globals between forward and backward. The recompute must see the identical function, or gradients are silently wrong.
Practice
TrenTorch: Gradient Checkpointing
Implement forward_full, backward_full, forward_checkpointed and backward_checkpointed for a stack of layers, plus count_stored_activations
The checkpointed gradients must match the full ones, since recomputation repeats the same operations. The section mini-project then measures the memory saved in PyTorch with equivalence checks.
Self-test
-
What stays alive across a checkpointed region, and what is recomputed?
Only the input args stay alive. The region reruns in backward to rebuild every forward tensor it needs. -
With
preserve_rng_stateFalse, forward outputs still match but gradients differ. Why?Forward runs once either way. The flag only controls the RNG state used by the backward recompute, so only dropout-dependent gradients change. -
What happens when
use_reentrantis omitted, and which value is recommended?PyTorch warns that it must be passed explicitly.Falseis recommended: it supportsgrad,kwargsand nested structures with early stop. -
With
use_reentrant=Trueand no input requiring grad, the region's parameters get no gradients. Why, and which variant avoids it?The reentrant variant runs the region underno_gradinside an autogradFunctionwhose output requires grad only if an input does. With none, backward never enters the region and PyTorch warns that gradients will beNone. The non-reentrant variant has no such requirement.
Completion criterion
Stop when
Your checkpointed block matches plain gradients at a fixed seed within a stated tolerance, and you can price any block in memory saved versus recompute added.
Mark passed, then open 8.7 Only mark this when you can do it without notes.Leads to
- 8.9 Profiling training phases and verifying optimizations Attribute the recompute cost to the backward phase in a profile.
Comes after
- 8.3 GPU memory: allocated, reserved, peak and memory modelsTrade the activation term for recompute once the model is measured.
Resources for this unit
- PyTorch - reliability and performance recipes Use the Profiler, Automatic Mixed Precision, model loading and Benchmark recipes, then the Performance Tuning Guide. Apply each to one compact model from an earlier section and keep a CPU fallback when CUDA is unavailable.
- PyTorch - checkpointing, numerical accuracy and CUDA semantics Use the official user-guide entries for torch.utils.checkpoint, numerical accuracy, CUDA semantics, reproducibility and serialization. Predict memory categories before measuring and define resume equivalence before running.
Also in this section
- PyTorch - compiler troubleshooting Read graph breaks, guards, recompilations and logging. Compile one stable fragment, intentionally trigger a graph break or recompilation, inspect it, and compare eager and compiled outputs and gradients.
- Deep Learning Training Debugger Use the failure-injection, measurement and evidence rules as the deeper standard. The full Tier 2 project is a later route and is not added to this stage milestone.
- TensorTonic - mixed-precision memory and transformer FLOPs Solve only Mixed-Precision Training Memory Accountant and Transformer Training FLOP Estimator. Reconcile both with the compact model's profiler and document every included and excluded term.
- Learn PyTorch - profiler walkthrough Use notebook 10 only if the official profiler route remains unclear. Transfer the relevant annotations into the compact model instead of keeping the tutorial code.
- TrenTorch - precision, checkpointing and resume Solve Loss Scaling, FP16/BF16 Representable Range, Gradient Checkpointing and Resume-from-checkpoint in NumPy, then reproduce each behavior with the PyTorch API in the compact model.
- Python os - `os.replace` Read
os.replacefor the atomic-rename step of checkpoint writes: save to a temp file in the same directory,flushandfsync, then replace the final name. - Hands-On Machine Learning with Scikit-Learn and PyTorch Read Appendix B, Reduced Precision Models and Mixed-Precision Training, for why
float16training underflows and whatfp32primary weights and loss scaling fix.