Add `jit` compilation to the Torch connector with `thunder`

[edoaltamura]

What should we add?

Thunder could be integrated with upcoming versions of the Torch connector and speed up the PyTorch compilation. This jit compiler is newly released and has promising results. From their repo:

Thunder is a source-to-source compiler for PyTorch. It makes PyTorch programs faster by combining and using different hardware executors at once (ie: nvFuser, torch.compile, cuDNN, and TransformerEngine FP8).

Works on single accelerators and in multi-GPU settings. Thunder aims to be usable, understandable, and extensible. Thunder achieves a 40% speedup in training throughput compared to eager code on H100 using a combination of executors including nvFuser, torch.compile, cuDNN, and TransformerEngine FP8.

Link: https://github.com/Lightning-AI/lightning-thunder

Minimal example:

import torch
import thunder


def foo(a, b):
    return a + b


jfoo = thunder.jit(foo)

a = torch.full((2, 2), 1)
b = torch.full((2, 2), 3)

result = jfoo(a, b)

print(result)

# prints
# tensor(
#  [[4, 4]
#   [4, 4]])