squall/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[misc] Add Torch profiler support for CPU-only devices (#7806)

Browse Source
This commit is contained in:
Jie Fu (傅杰) 2024-08-23 13:46:25 +08:00 committed by GitHub
parent fc5ebbd1d3
commit faeddb565d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 26 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
vllm/worker/cpu_worker.py
View File

@ -179,6 +179,32 @@ class CPUWorker(LoraNotSupportedWorkerBase, LocalOrDistributedWorkerBase):
self.cache_engine: List[CPUCacheEngine] self.cache_engine: List[CPUCacheEngine]
self.cpu_cache: List[List[torch.Tensor]] self.cpu_cache: List[List[torch.Tensor]]
# Torch profiler. Enabled and configured through env vars:
# VLLM_TORCH_PROFILER_DIR=/path/to/save/trace
if envs.VLLM_TORCH_PROFILER_DIR:
torch_profiler_trace_dir = envs.VLLM_TORCH_PROFILER_DIR
logger.info("Profiling enabled. Traces will be saved to: %s",
torch_profiler_trace_dir)
self.profiler = torch.profiler.profile(
activities=[
torch.profiler.ProfilerActivity.CPU,
],
with_stack=True,
on_trace_ready=torch.profiler.tensorboard_trace_handler(
torch_profiler_trace_dir, use_gzip=True))
else:
self.profiler = None
def start_profile(self):
if self.profiler is None:
raise RuntimeError("Profiler is not enabled.")
self.profiler.start()
def stop_profile(self):
if self.profiler is None:
raise RuntimeError("Profiler is not enabled.")
self.profiler.stop()
def init_device(self) -> None: def init_device(self) -> None:
if self.local_omp_cpuid != "all": if self.local_omp_cpuid != "all":
torch.ops._C_utils.init_cpu_threads_env(self.local_omp_cpuid) torch.ops._C_utils.init_cpu_threads_env(self.local_omp_cpuid)