Torch compile support for distributed operations

[AugustDev]

🚀 Feature

Documentation says that torch compile is not supported over distributed training right now. Since torch compile can speed up training as much as 2x using Lightning Trainer is without compile is no longer cost efficient and it would be great to support it.

It's bit unclear to me what happens if I compile the model before passing to Lightning module, will it be used as compiled model over DDP or not?

[t-vi]

@AugustDev Thank you, did you want to file this here or with https://github.com/Lightning-AI/pytorch-lightning/issues ?

[awaelchli]

The approach we took in Fabric should be transferrable to Trainer as well:
Lightning-AI/pytorch-lightning#19280
Lightning-AI/pytorch-lightning#19382
Essentially, it is just ensuring that torch.compile is applied over the FSDP/DDP wrapped model.