[TPU][Bugfix] Use XLA rank for persistent cache path (vllm-project#8137)

axonn-ai · Sep 30, 2024 · 07ad14a · 07ad14a
1 parent c6e2066
commit 07ad14a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/docs/source/getting_started/tpu-installation.rst b/docs/source/getting_started/tpu-installation.rst
@@ -59,7 +59,7 @@ First, install the dependencies:
     $ export DATE="20240828"
     $ export TORCH_VERSION="2.5.0"
     $ pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-${TORCH_VERSION}.dev${DATE}-cp310-cp310-linux_x86_64.whl
-    $ pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-${TORCH_VERSION}.dev${DATE}-cp310-cp310-linux_x86_64.whl
+    $ pip install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-${TORCH_VERSION}.dev${DATE}-cp310-cp310-linux_x86_64.whl
 
     $ # Install JAX and Pallas.
     $ pip install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html

diff --git a/vllm/worker/tpu_worker.py b/vllm/worker/tpu_worker.py
@@ -102,8 +102,9 @@ def init_device(self) -> None:
         # NOTE(woosuk): Set per-rank cache path since different ranks
         # can have slightly different XLA graphs.
         world_size = self.parallel_config.world_size
+        rank = xr.global_ordinal()
         per_rank_path = os.path.join(envs.VLLM_XLA_CACHE_PATH,
-                                     f"tp{world_size}_rank{self.rank}")
+                                     f"tp{world_size}_rank{rank}")
         xr.initialize_cache(per_rank_path, readonly=False)
 
     def load_model(self):