Xla: Unable to get any speedup from using TPUs

[dalmia]

❓ Questions and Help

Hi everyone. I am trying to use TPUs for the first time with PyTorch. I am following the official examples mentioned under contrib/colab to add support for this. However, the training is happening very slowly compared to a GPU. Can you please help me understand why that might be the case?

A Colab to reproduce my setup can be found here.

It first clones the Github repository that I am using for training, then it downloads the CIFAR-10 datasets and creates data version files that are going to be used for training and finally invokes the training script. Please let me know if you need more details. Also, simply loading the .yml files is taking too long.

[JackCaoG]

HI @dalmia , do you mind doing a debug run following the instruction in https://github.com/pytorch/xla/blob/master/contrib/colab/issue-report.ipynb? We can look at the debug run output and see where the slowness coming from.

[jysohn23]

Hi @dalmia looks like your workload might be CPU bound since Colab by default has very low # CPUs compared to the actual # cpus needed to fully saturate TPU tray.

Can you instead try cloning this notebook and pasting your contents onto it? We get special instances with larger CPUs: https://colab.research.google.com/github/pytorch/xla/blob/master/contrib/colab/getting-started.ipynb

[jysohn23]

Also note on the https://github.com/pytorch/xla/blob/master/contrib/colab/README.md:

Note: These colab notebooks typically run on small machines (the Compute VMs, which runs the input pipeline) and training is often bottlenecked on the small Compute VM machines. For optimal performance create a GCP VM and TPU pair following our GCP Tutorials:

[dalmia]

Thanks for getting back so promptly and confirming that low number of CPUs might be the main bottleneck. I'll follow the instruction in the debug run as well as try cloning the notebook linked by @jysohn23 and get back soon!

[zcain117]

Hey @dalmia did you have any more questions about this issue?

[dalmia]

Hi @zcain117. Sorry for not getting back earlier. We can close the issue for now. I switched to PyTorch Lightning which solved the issue for me!