lightning-thunder
lightning-thunder copied to clipboard
A converter for FXGraph with Torch calls -> FXGraph with Thunder calls
Dynamo represents the function argument of torch.utils.checkpoint.checkpoint(function, args...) as FX Graph, this FX graph has PyTorch operators in it. This PR creates a converter to replace the torch operators with the thunder equivalents, so it can be traced in the thunder checkpoint symbol. With this change native PyTorch activation checkpointing can be supported.