Post
12544
#TPU #PyTorch #Jax
When You're trying to use PyTorch or Jax on TPU,
for v2/v3/v4:
use tpu-ubuntu2204-base
for v5p:
use v2-alpha-tpuv5
for v5e:
use v2-alpha-tpuv5-lite
You must use these base images for the system to 'boot'.
Previously used tpu-vm-v4-pt-1.13 images might seem to start the VM, but SSH connections do not work.
I thought it was a firewall issue and spent a lot of time on it before realizing it was a problem with the boot image 🥲
https://cloud.google.com/tpu/docs/runtimes#pytorch_and_jax
When You're trying to use PyTorch or Jax on TPU,
for v2/v3/v4:
use tpu-ubuntu2204-base
for v5p:
use v2-alpha-tpuv5
for v5e:
use v2-alpha-tpuv5-lite
You must use these base images for the system to 'boot'.
Previously used tpu-vm-v4-pt-1.13 images might seem to start the VM, but SSH connections do not work.
I thought it was a firewall issue and spent a lot of time on it before realizing it was a problem with the boot image 🥲
https://cloud.google.com/tpu/docs/runtimes#pytorch_and_jax