Support latest version of Lightning

[awaelchli]

Lightning 2.0 is rolling out soon (mid march) and comes with PyTorch 2.0 support and other significant simplifications and UX improvements. I have been working on converting the OpenFold code to the latest version: https://github.com/awaelchli/openfold/pull/1. I currently investigate torch.compile on the model as well. As you can see in the linked draft PR, the user-facing changes are very minimal.

Are you interested in a PR that updates the code to Lightning 2.0?

The TODOs left are:

• Get flash attention compiling and working with torch 2.0
• Run experiments to demonstrate convergence to same numbers as previous code (guidance on this part would be appreciated :))
• Speed benchmarks with torch 2.0 and deepspeed
• Backward compatibility with older Lightning versions (if desired?)

Please let me know if there is interest, thanks!

cc @gahdritz @NZ99

[a-r-j]

@awaelchli I think this would be great. Optimizations and new features aside (I think they're great btw!), I'm currently having to run openfold with both lightning and pytorch_lightning installed which is obviously not ideal.

[amorehead]

Agreed. A PyTorch/Lighting 2.0 version of openfold would be very welcome.