pytorch / xla

Enabling PyTorch on XLA Devices (e.g. Google TPU)
https://pytorch.org/xla
Other
2.49k stars 483 forks source link

Unable to run on Google Colab #6568

Open BrandonStudio opened 9 months ago

BrandonStudio commented 9 months ago

HTTP 404 for url http://metadata.google.internal/computeMetadata/v1/instance/attributes/tpu-env when executing

from torch_xla._internal import tpu
tpu.num_available_devices()

See https://github.com/Lightning-AI/pytorch-lightning/issues/19274

JackCaoG commented 9 months ago

which version did you use? Colab is still on the old TPU Node architure so the last version we support is pytorch 2.0 with the old XRT runtime. The api you are trying to use belongs to the new PJRT runtime I believe. Can you try to use kaggle? Colab is working on the new TPUVM support and we should see more updates in next few months.

BrandonStudio commented 9 months ago

I believe I had used Kaggle without problems.

Could you give a guide about how to use tpu on Colab? As Google Colab is a widely-used free TPU platform and there are a bunch of different and contradictory methods. Should I install libtpu instead of torch_xla?

JackCaoG commented 9 months ago

You can install pytorch 2.0 and torch_xla 2.0 XRT build at

https://storage.googleapis.com/tpu-pytorch/wheels/colab/torch_xla-2.0-cp310-cp310-linux_x86_64.whl

then it should work

BrandonStudio commented 9 months ago

No, it does not work. image

Installation output: image