I built a Python library for cost-aware activation checkpoint placement

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: - ...