Lightning-AI / lightning-thunder

Make PyTorch models up to 40% faster! Thunder is a source to source compiler for PyTorch. It enables using different hardware executors at once; across one or thousands of GPUs.
Apache License 2.0
1.08k stars 62 forks source link

backwards of autograd functions #666

Open t-vi opened 1 week ago

t-vi commented 1 week ago

661 allows us to trace through torch.autograd.Function .apply using the forward, but so far we ignore the backward (and emit a warning). It would be neat to fix this and also trace the backward / use it in a grad_transform or so.

tfogal commented 1 week ago

triage review: