Xla: Running on Google Colab

[nmilosev]

Google Colab example (https://github.com/pytorch/xla/blob/master/contrib/colab/PyTorch_TPU_XRT_1_13.ipynb) is broken.

Here is the stack trace:

ImportError                               Traceback (most recent call last)
<ipython-input-3-16adb22a54bc> in <module>()
      4 import torch.nn.functional as F
      5 import torch.optim as optim
----> 6 from torchvision import datasets, transforms
      7 import torch_xla
      8 import torch_xla_py.utils as xu

5 frames
/usr/local/lib/python3.6/dist-packages/torchvision/__init__.py in <module>()
----> 1 from torchvision import models
      2 from torchvision import datasets
      3 from torchvision import ops
      4 from torchvision import transforms
      5 from torchvision import utils

/usr/local/lib/python3.6/dist-packages/torchvision/models/__init__.py in <module>()
      9 from .shufflenetv2 import *
     10 from . import segmentation
---> 11 from . import detection

/usr/local/lib/python3.6/dist-packages/torchvision/models/detection/__init__.py in <module>()
----> 1 from .faster_rcnn import *
      2 from .mask_rcnn import *
      3 from .keypoint_rcnn import *

/usr/local/lib/python3.6/dist-packages/torchvision/models/detection/faster_rcnn.py in <module>()
      5 import torch.nn.functional as F
      6 
----> 7 from torchvision.ops import misc as misc_nn_ops
      8 from torchvision.ops import MultiScaleRoIAlign
      9 

/usr/local/lib/python3.6/dist-packages/torchvision/ops/__init__.py in <module>()
----> 1 from .boxes import nms, box_iou
      2 from .roi_align import roi_align, RoIAlign
      3 from .roi_pool import roi_pool, RoIPool
      4 from .poolers import MultiScaleRoIAlign
      5 from .feature_pyramid_network import FeaturePyramidNetwork

/usr/local/lib/python3.6/dist-packages/torchvision/ops/boxes.py in <module>()
      1 import torch
----> 2 from torchvision import _C
      3 
      4 
      5 def nms(boxes, scores, iou_threshold):

ImportError: /usr/local/lib/python3.6/dist-packages/torchvision/_C.cpython-36m-x86_64-linux-gnu.so: undefined symbol: _ZN3c1019ComplexCUDATensorIdEv

Any advice on what is the issue? I tried downgrading TF to 1.13 but the error persists. Can someone explain why we need TF for this example? Is it because we are using it as a bridge since we are on Colab? This code looks outdated to me compared to the MNIST train test code.

On a related note I have written here:

https://nmilosev.svbtle.com/testing-pytorch-xla-with-google-colab-tpus

how to test xla on Google Colab as described in the project readme. It uses a SSH tunnel to expose provided Colab TPUs which can then be used locally. If you think this would be useful to someone, feel free to link it somewhere or I can send you a readme PR.

Thanks!

[ailzhang]

Hi @nmilosev I believe this error is because torchvision has a hard dependency on CUDA now which is missing in your environment. To fix it, try installing torchvision-cpu, that should fix it.
Btw torch_xla doesn't depend on TF but only XLA and XRT. ;)

[nmilosev]

@ailzhang Thank you for reply.

I replaced the contents of the first cell with:

!pip uninstall torch torchvision fastai -y
!pip install \
  http://storage.googleapis.com/pytorch-tpu-releases/tf-1.13/torch-1.0.0a0+1d94a2b-cp36-cp36m-linux_x86_64.whl  \
  http://storage.googleapis.com/pytorch-tpu-releases/tf-1.13/torch_xla-0.1+5622d42-cp36-cp36m-linux_x86_64.whl
!pip install torchvision==0.2.0 # workaround for https://github.com/pytorch/xla/issues/957

The cell after that (comparison) executes fine, but the main MNIST cell crashes the runtime. The log is attached, but I am not seeing anything unusual there. If you could take a look, that would be great. Thanks!

EDIT: I think the relevant part is:

Unknown opcode: cross-replica-sum

colab-jupyter.log

[jysohn23]

The GCS links to the wheels are obsolete and we have moved way past them. And yes that colab example is also broken/obsolete as it uses our old deprecated API (jit).

Are you using Colab's TPUs or starting a runtime on a compute VM and connecting to that local runtime from colab? If using Colab TPUs we won't necessarily work as colab has a fixed XRT version which we don't have wheels that have been released/tested against.

If you're using just as a frontend with a local runtime you can use these wheels (make sure the versions of the TPUs you create match the wheels):
pytorch-0.1: torch, torch_xla
pytorch-nightly: torch, torch_xla

[nmilosev]

The GCS links to the wheels are obsolete and we have moved way past them. And yes that colab example is also broken/obsolete as it uses our old deprecated API (jit).

Are you using Colab's TPUs or starting a runtime on a compute VM and connecting to that local runtime from colab? If using Colab TPUs we won't necessarily work as colab has a fixed XRT version which we don't have wheels that have been released/tested against.

I tried using the Colab directly, and got the error described above. I suspected that the example is outdated. I am not using Cloud, just Colab.

If you're using just as a frontend with a local runtime you can use these wheels (make sure the versions of the TPUs you create match the wheels):
pytorch-0.1: torch, torch_xla
pytorch-nightly: torch, torch_xla

SSH-ing to Colab, exposing TPUs via SSH tunnel, and running the nightly Docker image locally works perfectly so far.

Maybe it would be a good idea then to remove Colab example? Since it is broken.

Cheers

[dlibenzi]

The Colab example we have is still pinned to an old TPU VM release (TF 1.13) and uses the old JIT interface.
We will be cutting a new release based on TF 1.15 in 2..3 weeks, and will update the Colab example to use the current lazy tensor interface.
Sorry about that.

[nmilosev]

No worries! Thanks for your help.

If I can help with updating the Colab example, I would love to contribute. :)

[jysohn23]

Hi @nmilosev, I'll be cutting a release on our side (torch and torch_xla) and then create a quick demo colab based on those APIs. Once we have that I'll let you know and please do contribute to our colab examples 😄. Thanks!

[tmabraham]

Hello! If I understand correctly, using torch-xla will not work on the Colab TPU (connecting to cloud runtime). If this is the case, will there be an update to fix this? It would really make TPUs more accessible for PyTorch users. I am also interested in using fastai with TPUs so having the torch-xla work with Colab would be amazing.

[jysohn23]

HI @tmabraham! We're on making Colab TPUs to work with pytorch/xla as we speak. This should hopefully be available in the next coming weeks. We also plan to release some sample colab notebooks. I'll post on this thread once they're available.

Note that the VMs backing a Colab runtime only has a few CPU cores (and thus only a few CPU cores to run the input pipeline on), which will be far from enough to drive a TPU to its full performance. So we recommend creating a TPU on GCP eventually with a VM if you want to use the TPU to its full performance.

Thanks!

[tmabraham]

Thanks for the information!

[tmabraham]

@jysohn23 Just following up on this. I see there's an updated google colab in the repository, but it doesn't seem to work. Is this usable yet in Google Colab?

[jysohn23]

@jysohn23 Just following up on this. I see there's an updated google colab in the repository, but it doesn't seem to work. Is this usable yet in Google Colab?

As we speak, Colab is currently in the process of updating the runtimes to TF/XRT 1.15. Once they've updated their Colab runtimes, I will be adding a few notebooks that are TF/XRT 1.15 compatible in our repository in the next few coming days. Stay tuned!

[jysohn23]

https://github.com/pytorch/xla/pull/1153 Adds three sample Colab notebooks that should work with TF/XRT 1.15. Currently they're still upgrading to TF 1.15 (you can check on colab with a simple import tensorflow as tf; tf.__version__). But once they are done upgrading you should be able to use these notebooks. These notebooks may or may not work with 1.14 and there are bugs we've fixed since then so we recommend you use them once Colab is at 1.15.

Do note that Colab instances (the ones that are running the input pipeline, not the TPUs) are pretty small as is and so to get the full performance we recommend you create VM<>TPU pair on GCP.

[tmabraham]

@jysohn23 Thanks! I will try as soon as the Colab instances are upgraded to 1.15

I didn't understand your point about the CPUs. Do you mean that CPU-limited scripts will not achieve much gain using the TPU instances? Do the TPU instances get a better CPU and the GPU instances?

[tmabraham]

@jysohn23 Also what is the state of torchvision compatibility? It seems that in the colab document torchvision 0.3.0 version is installed?

[jysohn23]

I didn't understand your point about the CPUs. Do you mean that CPU-limited scripts will not achieve much gain using the TPU instances? Do the TPU instances get a better CPU and the GPU instances?

So the input pipelines get run on the CPUs and get converted into tensors. Those tensors then get sent to the TPU for computations on the model. So if your model is heavy on the input pipeline (ex. image processing) then you'd need many CPU cores to so that the TPU doesn't stay idle. Colab CPUs may not fit your needs so if you notice slow performance I was suggesting you create your own GCP VM with more CPU cores and go from there.

@jysohn23 Also what is the state of torchvision compatibility? It seems that in the colab document torchvision 0.3.0 version is installed?

The only reason we do that is due to https://github.com/pytorch/vision/issues/967 since we build pytorch from source. We are compatible with torchvision. In fact, one of the colab examples just downloads a pretrained torchvision model and uses that for inference on the TPU.

[tmabraham]

@jysohn23 Thanks for your response! I will note that the code worked with TF 1.14 as well.

[tmabraham]

@jysohn23 Will there be versions running on just a single TPU? Just to start with? The current ones seem to be for parallel training

[jysohn23]

By single TPU do you mean not using DataParallel and instead just using model.to(xm.xla_device()) directly?

Also you can use the Colab notebooks as of today. We're on XRT 1.15 on colab prod.

[tmabraham]

@jysohn23 I guess... Will it work just to do that? Don't we still have use the xm optimizer step functionality?

And again, I was able to run both CIFAR10 and MNIST training notebooks since the release 2 days ago even though it was using XRT 1.14

[jysohn23]

Yes you'd need to do something like xm.optimizer_step(optimizer, barrier=True) as suggested in our API_GUIDE.

And yes, but there were bugs/improvements since then.