Running PyTorch on TPU: a bag of tricks | by Zahar Chikishev | Towards Data Science

At the time of writing these lines running PyTorch code on TPUs is not a well-trodden path. Naturally, TPUs have been optimized for and mainly used with TensorFlow. But Kaggle and Google distribute free TPU time on some of its competitions, and one doesn’t simply change his favorite framework, so this is a memo on my (mostly successful) experience of training PyTorch models with TPU on GCP.
PyTorch/XLA is the project that allows doing it. It is still in active development, issues get fixed. Hopefully in the near future the experience of running it will get smoother, with some bugs fixed and the best practices better communicated.
My setup includes:
GCP Compute Engine virtual machine with prebuilt PyTorch/XLA image. Follow section “Consume Prebuilt Compute VM Images” on PyTorch/XLA github page to setup.TPU node, use this instruction with “GCP Console” option to create one for yourself. If your GCP project has a free TPU quota you received an email describing versions and zones that you can use. Besides that, I did not find a direct way to verify that created TPU node is indeed free, besides looking at the bill.Jupyter notebook, follow this funny article to set it up. I just can’t work without jupyter anymore. You may notice that you breathe deeper when it works for you.I have also found it convenient to work with Google Cloud Storage to transfer files in some cases. For example you can use the following lines on your virtual machine to copy Kaggle API token and download competition data with it. Use gsutil cp also to copy your files back to GS bucket.
GCP Compute Engine virtual machine with prebuilt PyTorch/XLA image. Follow section “Consume Prebuilt Compute VM Images” on PyTorch/XLA github page to setup.
TPU node, use this instruction with “GCP Console” option to create one for yourself. If your GCP project has a free TPU quota you received an email describing versions and zones that you can use. Besides that, I did not find a direct way to verify that created TPU node is indeed free, besides looking at the bill.
Jupyter notebook, follow this funny article to set it up. I just can’t work without jupyter anymore. You may notice that you breathe deeper when it works for you.
I have also found it convenient to work with Google Cloud Storage to transfer files in some cases. For example you can use the following lines on your virtual machine to copy Kaggle API token and download competition data with it. Use gsutil cp also to copy your files back to GS bucket.
gcloud auth logingsutil cp gs://bucket-name/kaggle-keys/kaggle.json ~/.kagglechmod 600 ~/.kaggle/kaggle.jsonkaggle competitions download -c recursion-cellular-image-classification
Besides google storage I used github repository to transfer data and code from my local machine to the GCP virtual machine and back.
Note that there are software versions running on TPU node as well. It must match the conda environment that you use on your VM. As PyTorch/XLA is under active development at the moment I use nightly TPU version:
Let’s get to the code. PyTorch/XLA has its own way of running multi-core, and as TPUs are multi-core you want to exploit it. But before you do, you may want to replace device = ‘cuda’ in your model with
import torch_xla_py.xla_model as xm...device = xm.xla_device()...xm.optimizer_step(optimizer)xm.mark_step()...
to test your model on only one core of TPU. The last two lines in the code snippet above replace the regular optimizer.step() call.
For multi-core training PyTorch/XLA uses its own DataParallel class. An example of a training loop with DataParallel can be found here in the test directory, and I want to highlight the following three points related to it.
1) DataParallel holds copies of the model object (one per TPU device), which are kept synchronized with identical weights. One can save the weights by accessing one of the models:
torch.save(model_parallel._models[0].state_dict(), filepath)
2) DataParallel cores must run the same number of batches each, and only full batches are allowed. Therefore, each epoch runs with less than 100% of the samples, and the residual is left out. With a dataset shuffling, it is not a big issue for the training loop, but it becomes a problem for inference. I am using single-core run for inference, as described above.
3) DataParallel code running directly from a jupyter notebook is highly unstable for me. It can run for some time, but then it throws system errors, kernel dies or even jupyter crashes. Running it as a script seems to be stable, so I created a second notebook for running with the following lines:
!jupyter nbconvert --to script MyModel.ipynb!python MyModel.py
PyTorch/XLA design results in a list of limitations on PyTorch functionality. In fact, these limitations are general to TPU devices, and apparently apply to TensorFlow models as well, at least partially. Specifically
Tensor shapes are advised to be the same between iterations, which also limits usage of masks.Loops with a different number of iterations between steps are to be avoided.
Tensor shapes are advised to be the same between iterations, which also limits usage of masks.
Loops with a different number of iterations between steps are to be avoided.
Not following the guidelines results in (severe) performance degradation. Unfortunately, in my loss function I needed to use both masks and loops. In my case I moved everything to CPU and it is much faster now. Just do my_tensor.cpu().detach().numpy() on all your tensors. Of course, it does not work on tensors where one needs to track the gradients, and it also incurs a slowdown of its own because of moving to CPU.
My Kaggle competition teammate Yuval Reina has kindly agreed to share his machine configuration and training speed for comparison in this section. I have also added a column for my notebook (a physical machine, not a jupyter this time), but it is not a match to these heavy-weights and the code running on it was not well-optimized for performance.
The input to the networks is 512 by 512 images with 6 channels. We measured the images per second processed in the training loop, and on this metric the described TPU configuration outperforms Yuval’s Tesla V100 by 110 to 57.
The performance for single core TPU as described above (without DataParallel) is 26 images per second, approximately 4 times slower than all 8 cores together.
We do not disclose the architecture used by Yuval as the competition is still ongoing, but it is not significantly different in size from resnet50. But please note that as we did not run the same architectures the comparison is not fair.
Trying to switch to GCP SSD disk for the train images did not improve performance.
In summary, my experience with PyTorch/XLA is mixed. I have encountered several bugs/artifacts (not all mentioned here), the existing documentation and examples are limited and the TPU intrinsic limitations can be too prohibitive for more creative architectures. On the other hand, it mostly works, and when it works the performance is good.
Finally, and arguably most important, don’t forget to stop your GCP VM when you are finished!