Lightning-AI / Lightning-AI/lightning-thunder

Raise an error when PyTorch's activation checkpointing is used with Thunder-jitted model

Open
#770 2 comments 0 reactions 1 assignee Claimed by @IvanYashchuk View on GitHub
warnings & errors
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.