Does PyTorch/XLA nightly provide GPU support?

[titaiwangms]

❓ Questions and Help

In README.md, there is nightly support on TPU

pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly-cp310-cp310-linux_x86_64.whl' -f https://storage.googleapis.com/libtpu-releases/index.html

But there is no instructions of XLA nightly support on GPU plugin. Is there a way that I can download PyTorch/XLA compatible version with torch-nightly?

[JackCaoG]

@lsy323 can you take a look? I think we just need an instruction to download torch nightly with cuda

[lsy323]

Hi @titaiwangms, please install the GPU plug-in by following https://github.com/pytorch/xla?tab=readme-ov-file#gpu-plugin. And we have the GPU guide at https://github.com/pytorch/xla/blob/master/docs/gpu.md.

https://pytorch.org/xla/master/gpu.html

[titaiwangms]

Hi @lsy323,

Thanks for the quick reply! For the first link https://github.com/pytorch/xla?tab=readme-ov-file#gpu-plugin, Is there a way that I can install with pytorch-nightly? It's supporting 2.4 release version. I am looking for something like:

python -m pip install --upgrade --pre torch torchvision torchaudio torch_xla --index-url https://download.pytorch.org/whl/nightly/cu118

Or I can only build on xla-nightly docker image (or build from source)?

[lsy323]

Hi @titaiwangms,

We have the nightly GPU plugin at https://storage.googleapis.com/pytorch-xla-releases/wheels/cuda/12.1/torch_xla_cuda_plugin-2.5.0-py3-none-any.whl. The suffix is none for now and will be fixed after #7697 is resolved.