diff --git a/vllm/v1/engine/utils.py b/vllm/v1/engine/utils.py index 28df7bbc718..cc8ce10e31e 100644 --- a/vllm/v1/engine/utils.py +++ b/vllm/v1/engine/utils.py @@ -155,15 +155,16 @@ class CoreEngineProcManager: try: for proc, local_dp_rank in zip(self.processes, local_dp_ranks): - # Adjust device control in DP for non-CUDA platforms - # as well as external and ray launchers - # For CUDA platforms, we use torch.accelerator.set_device_index()() + # Adjust device control in DP for platforms that cannot rely + # on torch.accelerator.set_device_index(), and for Ray launchers. device_control_context: contextlib.AbstractContextManager[None] = ( contextlib.nullcontext() ) + needs_device_env_isolation = not ( + current_platform.is_cuda_alike() or current_platform.is_xpu() + ) if is_dp and ( - not current_platform.is_cuda_alike() - or vllm_config.parallel_config.use_ray + needs_device_env_isolation or vllm_config.parallel_config.use_ray ): device_control_context = set_device_control_env_var( vllm_config, local_dp_rank diff --git a/vllm/v1/worker/xpu_worker.py b/vllm/v1/worker/xpu_worker.py index 102a04f5d88..0f01ee42a97 100644 --- a/vllm/v1/worker/xpu_worker.py +++ b/vllm/v1/worker/xpu_worker.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import gc import os -from typing import Any import torch @@ -40,19 +39,35 @@ class XPUWorker(Worker): assert device_config.device_type == "xpu" assert current_platform.is_xpu() - # Torch profiler. Enabled and configured through profiler_config. - self.profiler: Any | None = None - profiler_config = vllm_config.profiler_config - if profiler_config.profiler == "torch": - worker_name = f"{vllm_config.instance_id}-rank-{self.rank}" - self.profiler = TorchProfilerWrapper( - profiler_config, - worker_name=worker_name, - local_rank=self.local_rank, - activities=["CPU", "XPU"], + def init_device(self): + # In DP mode, XPU workers see all visible devices. + # Offset local_rank by the local DP shard. + parallel_config = self.parallel_config + if ( + parallel_config.distributed_executor_backend + not in ("ray", "external_launcher") + and parallel_config.data_parallel_backend != "ray" + and parallel_config.nnodes_within_dp == 1 + ): + dp_local_rank = parallel_config.data_parallel_rank_local + if dp_local_rank is None: + dp_local_rank = parallel_config.data_parallel_index + tp_pp_world_size = ( + parallel_config.pipeline_parallel_size + * parallel_config.tensor_parallel_size + ) + self.local_rank += dp_local_rank * tp_pp_world_size + + visible_device_count = torch.accelerator.device_count() + assert self.local_rank < visible_device_count, ( + f"DP adjusted local rank {self.local_rank} is out of bounds. " + ) + assert parallel_config.local_world_size <= visible_device_count, ( + f"local_world_size ({parallel_config.local_world_size}) must " + f"be less than or equal to the number of visible devices " + f"({visible_device_count})." ) - def init_device(self): device = self.device_config.device if ( isinstance(device, torch.device) @@ -117,3 +132,30 @@ class XPUWorker(Worker): if self.rank == 0: # If usage stat is enabled, collect relevant info. report_usage_stats(self.vllm_config) + + def profile(self, is_start: bool = True, profile_prefix: str | None = None): + if self.profiler_config is None or self.profiler_config.profiler is None: + raise RuntimeError( + "Profiling is not enabled. Please set --profiler-config to enable " + "profiling. Example: " + "'--profiler-config.profiler=torch --profiler-config.torch_profiler_dir" + "=YOUR_DIR_PATH_TO_DUMP_TRACE'" + ) + + if is_start and self.profiler is None: + from vllm.distributed.utils import get_worker_rank_suffix + + rank_suffix = get_worker_rank_suffix(global_rank=self.rank) + trace_name = ( + f"{profile_prefix}_{rank_suffix}" if profile_prefix else rank_suffix + ) + + self.profiler = TorchProfilerWrapper( + self.profiler_config, + worker_name=trace_name, + local_rank=self.local_rank, + activities=["CPU", "XPU"], + ) + logger.debug("Starting torch profiler with trace name: %s", trace_name) + + super().profile(is_start=is_start, profile_prefix=profile_prefix)