About pytorch 2.0 and pytorch-lightning

[PoPdark]

Pytorch 2.0 will be released soon. I am wondering if you have the schedule to follow it? Thank you.

cc @Borda

[awaelchli]

Since 2.0 is completely backward compatible with 1.13, it will just work out of the box with Lightning. In fact, you can try it today from the nightlies. Installation instructions are here.

In time, the team will add 2.0 to the CI testing suite to run all tests against it.

Then there is the big new feature in 2.0 that allows you to compile the model with torch.compile. The expectation is that he users can do that directly on the LightningModule. Some tests will have to be added here.

We can't give you a concrete ETA. So far, Lightning has always released a version shortly after the stable release of each PyTorch version with updated compatibility. We intend to follow that same pattern for 2.0.

[Borda]

Pytorch 2.0 will be released soon. I am wondering if you have the schedule to follow it?

yes we will b testing it as son as they provide some public wheels, but until their official release (March 2023) we have about 4 months :)

[lantiga]

Just FYI @PoPdark, support for passing models produced by torch.compile to the Trainer that just landed in Lightning master. It will be included in the upcoming patch release this week.

In order to try it out you'll need to have a recent PyTorch nightly installed alongside lightning.

[lantiga]

This went out with 1.8.4 https://github.com/Lightning-AI/lightning/releases/tag/1.8.4, closing the issue

@PoPdark let me know if you have issues trying things out, remember you have to manually get the PyTorch nightlies to get torch.compile

[yangky11]

Thanks！ @lantiga

Does it also work out of the box iftorch.compile is not called manually but the model is created by Lightning CLI?

[lantiga]

For now we don't have a Trainer flag (say, compile=True), but we may add it in the future.

So for the time being you need to do things just like the PyTorch docs say:

model = MyModel()
model = torch.compile(model)

trainer = Trainer(...)
trainer.fit(model)

Once we add a trainer flag you'll be able to enable compilation from the Lightning CLI.

[yangky11]

Yep, I think a compile flag in the future would be really helpful!