Delete XRT code paths

[carmocca]

Outline & Motivation

The PyTorch XLA project has dropped support of the XRT runtime in master (PR pytorch/xla#5240).
This will break our XLA integration.

We added support for PjRT recently with #17350. Which will be released with the next feature release.

Pitch

Remove everything that's not under if pjrt.using_pjrt().
Delete the XRT tests.
Update the docs (part of #17492)

Additional context

It's not clear yet if XRT support will be dropped with XLA 2.1 or 2.2

cc @justusschock @awaelchli @Borda @carmocca @JackCaoG @steventk-g @Liyang90

[will-cromar]

Responding to pytorch/xla#5240 (comment):

Am I understanding correctly that 2.1 will have 2 different wheels, one with XRT + PjRT and another with just PjRT?

Yes. We'll create two versions of the PyTorch/XLA release: one with XRT + PJRT that is built off of TensorFlow, and one with just PJRT that is built off of OpenXLA.

If so, will both be under the torch_xla package name? Will there be a recommended way to distinguish them other than calling hasattr on the removed code?

We'll keep the torch_xla module name. What works best from your perspective? My first thought was to distinguish the packages by version (e.g. 2.0.0+xrttpu). But this is harder to use to distinguish the packages at runtime. I'm happy to add an attribute like torch_xla.has_xrt_support if that makes it easier for you.

I would appreciate it if you could notify me when these wheels are in a release candidate stage.

We started building nightly versions of the XRT wheels last week for our end-to-end tests. Our install script and the wheels themselves are public if you want to test with them. We can give you a heads-up when we create RC wheels for the 2.1 release.

It's not clear yet if XRT support will be dropped with XLA 2.1 or 2.2

We'll be dropping XRT support in the 2.2 release.

[carmocca]

What works best from your perspective?

I just had a more thorough look at your removal PR and our implementation. Luckily you kept torch_xla.experimental.pjrt.using_pjrt which we always use before accessing XRT specific logic logic.
The only breaking change I found was the removal of this xenv: https://github.com/Lightning-AI/lightning/blob/master/src/lightning/fabric/plugins/environments/xla.py#L78. What is the replacement with PjRT?

We will need to avoid the deprecation warnings introduced in pytorch/xla#5155 from importing torch_xla.experimental.

We can give you a heads-up when we create RC wheels for the 2.1 release.

Yes please! I might also try before this happens if I get a chance.

We'll be dropping XRT support in the 2.2 release.

Thanks for the info. We'll do the same then.

[will-cromar]

The only breaking change I found was the removal of this xenv

Is this the only place you use xenv? I removed it from master because I didn't want anyone to confuse them for something valid in PJRT. If you use those env vars to detect whether XRT is being used, I can add them back so you can avoid an attribute error.

What is the replacement with PjRT?

Off the top of my head, I'm not sure we have a direct replacement for HOST_ORDINAL in PJRT. If you need it, I can add an way to get the host ID to torch_xla.runtime.

Let me know if you have any issues moving off of torch_xla.experimental. Almost all of the public APIs stayed identical and just moved to torch_xla.runtime, except for broadcast_master_param and rendezvous.

[carmocca]

I opened #17993 with the prospective changes required, luckily, I don't think you need to add any code back 😃

If you need it, I can add an way to get the host ID to torch_xla.runtime.

Yes, this would be useful. Thank you

Let me know if you have any issues moving off ...

We are in the process of renewing our TPU quota. I'll be able to test things when that happens.