Xla: GPU support in PyTorch XLA

[goswamig]

❓ Questions and Help

Do we have GPU support in PyTorch XLA? If so would you please point me to any resource to begin with ?

[jysohn23]

Yes, it does support GPUs, but we have limited testing with them tbh. We've just set up a release pipeline for torch_xla wheels built with GPU support, so we'll work on testing them more internally, adding nightly tests, and create a sample PyTorch/XLA GPU notebook from it once it's ready.

[goswamig]

I am happy to help to test it out, if you have any doc to setup and run basic MNIST ?

[dlibenzi]

Maybe check here: https://github.com/pytorch/xla/issues/2121

[goswamig]

Thank you @dlibenzi I tried those steps, container build was successful but then import torch_xla was failing. I think we can polish that instruction.
Also do you have wheel or docker container build which I can test or use ?

[stale[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[zcain117]

Sorry for the delay @gautamkmr , we are getting close to something that we can share but we're not quite there yet. We're planning on offering wheels and a docker image that you can use just like our normal wheels/docker image.

Once that is ready, it should be really easy to switch between TPU and GPU without having to change your code. We can let you know once it's ready if you're still interested in trying it

[gupta-abhay]

@zcain117 - I am interested in trying this as well. Please let me know so that I can give this a shot. Thank you!

[tyoc213]

Will I be able to install on a normal PC with CUDA support?

[stale[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[zcain117]

@goswamig @gupta-abhay @tyoc213

I tried this Friday and it worked for me:

1. Made a VM in Google Cloud using the boot disk: "Deep Learning Image: Pytorch 1.4 and fastai m55". Sub-description: "PyTorch 1.4.0 (and fastai) with CUDA 10.0 and Intel® MKL-DNN, Intel® MKL." I think any VM that has up-to-date CUDA installed would probably work.
2. Pulled the pytorch/xla docker image: docker run --gpus all -it --shm-size 16G gcr.io/tpu-pytorch/xla@sha256:efe47b7a3875ddfa3ea9c68a12ed5517c19cbb3cf272776fba64bec8a683299f. This was Friday's pytorch/xla CUDA image but the latest nightly should also work. The tag would be nightly_3.6_cuda or nightly_3.7_cuda or you can use a dated image like nightly_3.6_cuda_20200914.
3. Ran training with a new environment variable: GPU_NUM_DEVICES=4 python test/test_train_mp_mnist.py

I was able to train on GPUs with 0 code change. We don't have automated tests for XLA GPU so we won't call this officially supported yet, but I'm working on making those tests.

Let me know if you're able to get it working. If you run into any issues, please open a new issue that has the error you're seeing and details about VM software, TPU software, CUDA version, etc.