A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 21:00 UTC
[TPU][Profiler] Support start_profile/stop_profile in TPU worker (#13988)
Signed-off-by: Siyuan Liu <lsiyuan@google.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Siyuan Liu committed
Mar 4, 2025 at 11:40 UTC
beebf4742af80296d3c3a657c66d512615c550c1
3 files changed
+47
-5
requirements-tpu.txt
+6
-5
@@ -17,8 +17,9 @@ ray[default]
17
--find-links https://storage.googleapis.com/libtpu-releases/index.html
18
--find-links https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
19
--find-links https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html
20
-
21
-torch==2.7.0.dev20250226+cpu
22
-torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250226+cxx11-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"
23
-torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250226+cxx11-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"
24
-torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250226+cxx11-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"
20
+torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.7.0.dev20250227%2Bcxx11-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"
21
+torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.7.0.dev20250227%2Bcxx11-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"
22
+torch @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-2.7.0.dev20250227%2Bcxx11-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"
23
+torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250227%2Bcxx11-cp39-cp39-linux_x86_64.whl ; python_version == "3.9"
24
+torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250227%2Bcxx11-cp310-cp310-linux_x86_64.whl ; python_version == "3.10"
25
+torch_xla[tpu, pallas] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.7.0.dev20250227%2Bcxx11-cp311-cp311-linux_x86_64.whl ; python_version == "3.11"
vllm/v1/worker/tpu_worker.py
+19
@@ -7,6 +7,7 @@ import torch
7
import torch.distributed
8
import torch.nn as nn
9
import torch_xla.core.xla_model as xm
10
+import torch_xla.debug.profiler as xp
11
import torch_xla.runtime as xr
12
13
import vllm.envs as envs
@@ -65,6 +66,15 @@ class TPUWorker:
66
from vllm.utils import init_cached_hf_modules
67
init_cached_hf_modules()
68
69
+ self.profiler = None
70
+ if envs.VLLM_TORCH_PROFILER_DIR and self.rank < 1:
71
+ # For TPU, we can only have 1 active profiler session for 1 profiler
72
+ # server. So we only profile on rank0.
73
+ self.profile_dir = envs.VLLM_TORCH_PROFILER_DIR
74
+ logger.info("Profiling enabled. Traces will be saved to: %s",
75
+ self.profile_dir)
76
+ self.profiler = xp.start_server(9012)
77
+
78
def init_device(self):
79
os.environ["PJRT_DEVICE"] = "TPU"
80
torch.set_grad_enabled(False)
@@ -152,6 +162,15 @@ class TPUWorker:
162
output = self.model_runner.execute_model(scheduler_output)
163
return output if self.is_driver_worker else None
164
165
+ def profile(self, is_start: bool = True):
166
+ if self.rank < 1:
167
+ if self.profiler is None:
168
+ raise RuntimeError("Profiler is not enabled.")
169
+ if is_start:
170
+ xp.start_trace(self.profile_dir)
171
+ else:
172
+ xp.stop_trace()
173
+
174
def load_model(self) -> None:
175
self.model_runner.load_model()
176
vllm/worker/tpu_worker.py
+22
@@ -5,6 +5,7 @@ from typing import List, Optional, Tuple, Union
5
6
import torch
7
import torch_xla.core.xla_model as xm
8
+import torch_xla.debug.profiler as xp
9
import torch_xla.runtime as xr
10
11
import vllm.envs as envs
@@ -93,6 +94,27 @@ class TPUWorker(LoRANotSupportedWorkerBase, LocalOrDistributedWorkerBase):
94
f"tp{world_size}_rank{rank}")
95
xr.initialize_cache(per_rank_path, readonly=False)
96
97
+ self.profiler = None
98
+ if envs.VLLM_TORCH_PROFILER_DIR and self.rank < 1:
99
+ # For TPU, we can only have 1 active profiler session for 1 profiler
100
+ # server. So we only profile on rank0.
101
+ self.profile_dir = envs.VLLM_TORCH_PROFILER_DIR
102
+ logger.info("Profiling enabled. Traces will be saved to: %s",
103
+ self.profile_dir)
104
+ self.profiler = xp.start_server(9012)
105
+
106
+ def start_profile(self):
107
+ if self.rank < 1:
108
+ if self.profiler is None:
109
+ raise RuntimeError("Profiler is not enabled.")
110
+ xp.start_trace(self.profile_dir)
111
+
112
+ def stop_profile(self):
113
+ if self.rank < 1:
114
+ if self.profiler is None:
115
+ raise RuntimeError("Profiler is not enabled.")
116
+ xp.stop_trace()
117
+
118
def load_model(self):
119
self.model_runner.load_model()
120