Hi everyone,
I’m excited to share an early release of `ckptplan`, a Python library that helps with memory-aware activation checkpoint planning in PyTorch.
Rather than checkpointing every layer, ckptplan profiles model blocks that can be checkpointed and chooses where to place checkpoints based on a memory budget. It offers several planners: uniform, greedy, checkpoint-all, no-checkpoint, and dynamic programming.
The main goal is to make the tradeoff between memory and computation clear:
- PyTorch-based
- CUDA GPU memory profiling
- Dynamic-programming checkpoint placement
- Gradient correctness validation
- Benchmarking for peak memory and step latency
Here’s an example from a real A10G experiment with a 24-layer, 1.2 billion parameter Transformer and a sequence length of 2048:
- No checkpointing ran out of memory at batch size 4.
- ckptplan-DP sustained batch size 4 at approximately 11.8 GiB peak allocation.
- As expected, the tradeoff was extra recomputation time.
Install:
pip install ckptplan
GitHub:
PyPI:
I’d really appreciate any feedback on the API design, the planner assumptions, or ideas for integrating this with larger PyTorch training workflows.