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, 17:15 UTC
[Hardware][TPU][Bugfix] Fix v1 mp profiler (#15409)
Signed-off-by: Siyuan Liu <lsiyuan@google.com>
Siyuan Liu committed
Mar 25, 2025 at 01:43 UTC
4157f563b4e82881309521c5aecabd4527db0a52
1 file changed
+8
-2
vllm/v1/worker/tpu_worker.py
+8
-2
@@ -66,14 +66,18 @@ class TPUWorker:
66
from vllm.utils import init_cached_hf_modules
67
init_cached_hf_modules()
68
69
+ # Delay profiler initialization to the start of the profiling.
70
+ # This is because in vLLM V1, MP runtime is initialized before the
71
+ # TPU Worker is initialized. The profiler server needs to start after
72
+ # MP runtime is initialized.
73
self.profiler = None
74
+ self.profile_dir = None
75
if envs.VLLM_TORCH_PROFILER_DIR and self.rank < 1:
76
# For TPU, we can only have 1 active profiler session for 1 profiler
77
# server. So we only profile on rank0.
78
self.profile_dir = envs.VLLM_TORCH_PROFILER_DIR
79
logger.info("Profiling enabled. Traces will be saved to: %s",
80
self.profile_dir)
76
- self.profiler = xp.start_server(9012)
81
82
if self.model_config.seed is None:
83
self.model_config.seed = 0
@@ -168,9 +172,11 @@ class TPUWorker:
172
173
def profile(self, is_start: bool = True):
174
if self.rank < 1:
171
- if self.profiler is None:
175
+ if self.profile_dir is None:
176
raise RuntimeError("Profiler is not enabled.")
177
if is_start:
178
+ if self.profiler is None:
179
+ self.profiler = xp.start_server(9012)
180
xp.start_trace(self.profile_dir)
181
else:
182
xp.stop_trace()