ConnectionRefusedError: [Errno 111] Connection refused TPU

[huzama]

🐛 Bug Report

Description

When attempting to initialize a TPU pod, a ConnectionRefusedError: [Errno 111] Connection refused TPU error occurs.

To Reproduce

Steps to reproduce the behavior:

1. Setup TPU Pod: Ensure that your TPU pod is properly configured and active.
2. Install Dependencies: Ensure torch >= v2.4.1 and torch_xla >= v2.4.0 are installed.
3. Run the Following Script:

import torch.distributed as dist
import torch_xla.runtime as xr
import torch_xla.distributed.xla_backend  # Import to register the `xla://` init_method

xr.use_spmd()
dist.init_process_group("gloo", init_method="xla://")

Expected Behavior

The script should run without errors, initializing the process group correctly on the TPU pod.

Environment

• XLA Backend: TPU
• Torch XLA Version: 2.4.0
• Torch Version: 2.4.1 or greater

Additional Context

• The issue does not occur with torch and torch_xla versions <= 2.4.0.

Error Message

ConnectionRefusedError: [Errno 111] Connection refused TPU