Lightning-AI / Lightning-AI/lightning-thunder
Raise an error when PyTorch's activation checkpointing is used with Thunder-jitted model
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🐛 Bug
Activation checkpointing is a technique to reduce memory usage by clearing activations of certain layers and recomputing them during the backward pass. Thunder doesn't know how to deal with PyTorch's annotations and wrappers implementing this feature that leads to silent or loud problems, like https://github.com/Lightning-AI/lightning-thunder/issues/582.
When nn.Module passed to `thunder.jit` includes an instance of `torch.distributed.algorithms._checkpoint.checkpoint_wrapper.CheckpointWrapper` in its submodules Thunder should raise an error.
Is there any other indicator that the user used `torch.distributed.algorithms._checkpoint.checkpoint_wrapper.apply_activation_checkpointing`?
cc @apaz-cli
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.