TPU VM (v4-32, tpu-ubuntu2204-base) hangs when verifying TPU with `xm.get_xla_supported_devices("TPU")`

[renxiangdai]

Environment

• TPU hardware: v4-32 (TPU VM)
• OS image: tpu-ubuntu2204-base

Steps to reproduce

On a fresh TPU VM:

bash
sudo apt-get update
sudo apt-get install libopenblas-dev -y
pip install numpy
pip install torch torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html

Then run the verification command:

PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"

Expected behavior

The command should return the list of available TPU devices, e.g.:
['TPU:0', 'TPU:1', ...]

Actual behavior

The command hangs indefinitely and never prints any output.
This makes it impossible to confirm TPU accessibility from PyTorch/XLA on the TPU VM.

Additional context

• This was run directly on a TPU VM (not Colab).

• No error messages are shown — just hangs after import.

[ysiraichi]

I tried on a v5litepod-4 with software version v2-alpha-tpuv5-lite , and wasn't able to reproduce the issue.
Instead, I got:

Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "torch_xla/core/xla_model.py", line 82, in get_xla_supported_devices
    for i, _ in enumerate(devices[:max_devices] if max_devices else devices)
TypeError: slice indices must be integers or None or have an __index__ method

Which is (kind of) expected.
Note that get_xla_supported_devices() devkind parameter was deprecated at #9284.

Removing the "TPU" argument, I get:

['xla:0', 'xla:1', 'xla:2', 'xla:3']

[ysiraichi]

If you stop the execution, what do you get?

[Muinez]

I'm experiencing same issue with v4-32 (4 hosts).

The command PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices())" completes successfully only when run on all hosts simultaneously. If I run it on a single host, the process hangs indefinitely.
When running the official PyTorch TPU Pod examples from the documentation (https://cloud.google.com/tpu/docs/pytorch-pods), the training process appears to be desynchronized. It seems each host trains its own model independently.
The key indicators of this are:
The training steps are out of sync across hosts (e.g., one host is at step 80 while another is at 120).
It produces four separate logs instead of one unified log.
I've observed the exact same behavior with my own custom training script as well.

I have also tried using both .spawn and the .launch script to initiate the training. In both cases, when I launch the script across all hosts, the processes hang indefinitely and never start the actual training

@ysiraichi

[Muinez]

I'm experiencing same issue with v4-32 (4 hosts).

The command PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices())" completes successfully only when run on all hosts simultaneously. If I run it on a single host, the process hangs indefinitely. When running the official PyTorch TPU Pod examples from the documentation (https://cloud.google.com/tpu/docs/pytorch-pods), the training process appears to be desynchronized. It seems each host trains its own model independently. The key indicators of this are: The training steps are out of sync across hosts (e.g., one host is at step 80 while another is at 120). It produces four separate logs instead of one unified log. I've observed the exact same behavior with my own custom training script as well.

I have also tried using both .spawn and the .launch script to initiate the training. In both cases, when I launch the script across all hosts, the processes hang indefinitely and never start the actual training

@ysiraichi

I've recreated the instance, and that has fixed .launch and desync issues. I guess I messed up the network settings during the creation process

[SVISHNUVARDHAN3610]

Hi @renxiangdai I'm currently working on TPU V5e-8 and it works perfectly fine xm.get_xla_supported_devices() try to use this

import os
os.environ.pop('TPU_PROCESS_ADDRESSES')

[renxiangdai]

Thanks for all the input, especially the key observation from @Muinez.

After further investigation and testing, I believe this is not a bug in torch_xla but rather expected behavior due to the multi-host architecture of TPU v4 Pods.

Here's my reasoning:

• The command works perfectly on a tpu-v4-4, which is a single-host setup.

• The command hangs on a tpu-v4-32, which is a multi-host setup (4 hosts in this case).

• As @Muinez confirmed, the command only completes successfully on a v4-32 when it is executed on all hosts simultaneously.

This strongly suggests that initializing torch_xla or discovering devices on a TPU Pod slice requires a collective communication step across all hosts. When the script is run on a single host, it hangs because it's waiting for the other hosts in the pod to participate, which they never do. Given this, the correct way to run any code (even a simple verification script) on a TPU Pod is to launch it across all hosts.

Based on this understanding, I'm closing this issue.