lightning-thunder icon indicating copy to clipboard operation
lightning-thunder copied to clipboard

A converter for FXGraph with Torch calls -> FXGraph with Thunder calls

Open kiya00 opened this issue 1 year ago • 0 comments

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.

kiya00 avatar Oct 04 '24 14:10 kiya00