import torch import torch.distributed as dist from torch.distributed import ProcessGroup from vllm.platforms import current_platform if current_platform.is_tpu(): import torch_xla.core.xla_model as xm import torch_xla.runtime as xr from torch_xla._internal import pjrt from vllm.executor import ray_utils class TpuCommunicator: def init(self, group: ProcessGroup): if not current_platform.is_tpu(): self.disabled = True return self.disabled = False # NOTE(woosuk): When using TP > 1 on TPUs, every TPU on the same node # must be used together. Therefore, the local rank and world size can # be simply calculated as follows. global_rank = dist.get_rank(group) global_world_size = dist.get_world_size(group) # Calculate how many TPU nodes are in the current deployment. This # is the Ray placement group if it is deployed with Ray. Default # to the number of TPU nodes in the Ray cluster. The number of TPU # nodes is computed by the total number of TPUs divided by the # number of TPU accelerators per node, to account for clusters # with both CPUs and TPUs. num_nodes = ray_utils.get_num_tpu_nodes() num_nodes_in_pg = ray_utils.get_num_nodes_in_placement_group() if num_nodes_in_pg > 0: num_nodes = num_nodes_in_pg local_world_size = global_world_size // num_nodes local_rank = global_rank % local_world_size # Ensure environment variables are set for multihost deployments. # On GKE, this is needed for libtpu and TPU driver to know which TPU # chip is actually visible. Otherwise the TPU driver will fail to # initialize because the number of devices would be different from # the number of visible worker addresses. os.environ["CLOUD_TPU_TASK_ID"] = str(global_rank) os.environ["TPU_VISIBLE_CHIPS"] = str(local_rank) pjrt.initialize_multiprocess(local_rank, local_world_size) xr._init_world_size_ordinal() def all_reduce(self, x: torch.Tensor) -> torch.Tensor: return xm.all_reduce(xm.REDUCE_SUM, x) def all_gather(self, x: torch.Tensor, dim: int = -1) -> torch.Tensor: assert dim == -1, "TPUs only support dim=-1 for all-gather." return xm.all_gather(x, dim=dim) , I am using 4 tpus and i have setup the ray backend with ray start head and then connecting all nodes to it , but tpu initialization is still failing , help me debug it I will give you all the code you need for it [rank0]: File "/nfs_share/vllm/vllm/executor/ray_tpu_executor.py", line 39, in init [rank0]: super().init(*args, **kwargs) [rank0]: File "/nfs_share/vllm/vllm/executor/executor_base.py", line 47, in init [rank0]: self._init_executor() [rank0]: File "/nfs_share/vllm/vllm/executor/ray_tpu_executor.py", line 51, in _init_executor [rank0]: self._init_workers_ray(placement_group) [rank0]: File "/nfs_share/vllm/vllm/executor/ray_tpu_executor.py", line 197, in _init_workers_ray [rank0]: self._run_workers("init_device") [rank0]: File "/nfs_share/vllm/vllm/executor/ray_tpu_executor.py", line 264, in _run_workers [rank0]: driver_worker_output = self.driver_worker.execute_method( [rank0]: File "/nfs_share/vllm/vllm/worker/worker_base.py", line 465, in execute_method [rank0]: raise e [rank0]: File "/nfs_share/vllm/vllm/worker/worker_base.py", line 456, in execute_method [rank0]: return executor(*args, **kwargs) [rank0]: File "/nfs_share/vllm/vllm/worker/tpu_worker.py", line 83, in init_device [rank0]: ensure_model_parallel_initialized( [rank0]: File "/nfs_share/vllm/vllm/distributed/parallel_state.py", line 1054, in ensure_model_parallel_initialized [rank0]: initialize_model_parallel(tensor_model_parallel_size, [rank0]: File "/nfs_share/vllm/vllm/distributed/parallel_state.py", line 1035, in initialize_model_parallel [rank0]: _PP = init_model_parallel_group(group_ranks, [rank0]: File "/nfs_share/vllm/vllm/distributed/parallel_state.py", line 859, in init_model_parallel_group [rank0]: return GroupCoordinator( [rank0]: File "/nfs_share/vllm/vllm/distributed/parallel_state.py", line 229, in init [rank0]: self.tpu_communicator = TpuCommunicator(group=self.cpu_group) [rank0]: File "/nfs_share/vllm/vllm/distributed/device_communicators/tpu_communicator.py", line 53, in init [rank0]: pjrt.initialize_multiprocess(local_rank, local_world_size) [rank0]: File "/home/spars/.local/lib/python3.10/site-packages/torch_xla/_internal/pjrt.py", line 121, in initialize_multiprocess [rank0]: devices = xm.get_xla_supported_devices() [rank0]: File "/home/spars/.local/lib/python3.10/site-packages/torch_xla/core/xla_model.py", line 93, in get_xla_supported_devices [rank0]: devices = torch_xla._XLAC._xla_get_devices() [rank0]: RuntimeError: Bad StatusOr access: UNKNOWN: TPU initialization failed: Invalid --2a886c8_slice_builder_worker_addresses specified. Expected 16 worker addresses, got 4.