| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import gc |
| 4 | import os |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.config import VllmConfig |
| 9 | from vllm.logger import init_logger |
| 10 | from vllm.platforms import current_platform |
| 11 | from vllm.utils.mem_utils import MemorySnapshot, format_gib |
| 12 | from vllm.utils.torch_utils import set_random_seed |
| 13 | from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment |
| 14 | from vllm.v1.worker.xpu_model_runner import ( |
| 15 | XPUMMEncoderModelRunner, |
| 16 | XPUModelRunner, |
| 17 | XPUModelRunnerV2, |
| 18 | ) |
| 19 | |
| 20 | from .utils import request_memory |
| 21 | |
| 22 | logger = init_logger(__name__) |
| 23 | |
| 24 | |
| 25 | class XPUWorker(Worker): |
| 26 | """A XPU worker class.""" |
| 27 | |
| 28 | def __init__( |
| 29 | self, |
| 30 | vllm_config: VllmConfig, |
| 31 | local_rank: int, |
| 32 | rank: int, |
| 33 | distributed_init_method: str, |
| 34 | is_driver_worker: bool = False, |
| 35 | ): |
| 36 | super().__init__( |
| 37 | vllm_config, local_rank, rank, distributed_init_method, is_driver_worker |
| 38 | ) |
| 39 | device_config = self.device_config |
| 40 | assert device_config.device_type == "xpu" |
| 41 | assert current_platform.is_xpu() |
| 42 | |
| 43 | def init_device(self): |
| 44 | # In DP mode, XPU workers see all visible devices. |
| 45 | # Offset local_rank by the local DP shard. |
| 46 | parallel_config = self.parallel_config |
| 47 | if ( |
| 48 | parallel_config.distributed_executor_backend |
| 49 | not in ("ray", "external_launcher") |
| 50 | and parallel_config.data_parallel_backend != "ray" |
| 51 | and ( |
| 52 | parallel_config.data_parallel_external_lb |
| 53 | or parallel_config.nnodes_within_dp == 1 |
| 54 | ) |
| 55 | ): |
| 56 | dp_local_rank = parallel_config.data_parallel_rank_local |
| 57 | if dp_local_rank is None: |
| 58 | dp_local_rank = parallel_config.data_parallel_index |
| 59 | replica_world_size = parallel_config.world_size |
| 60 | visible_device_count = torch.accelerator.device_count() |
| 61 | |
| 62 | if parallel_config.data_parallel_external_lb: |
| 63 | if parallel_config.nnodes_within_dp > 1: |
| 64 | # The replica spans nodes, so use its node-local shard. |
| 65 | replica_world_size = parallel_config.local_world_size |
| 66 | if replica_world_size > visible_device_count: |
| 67 | raise ValueError( |
| 68 | f"Local TP/PP/PCP replica size ({replica_world_size}) " |
| 69 | "exceeds the number of visible XPU devices " |
| 70 | f"({visible_device_count})." |
| 71 | ) |
| 72 | local_dp_capacity = visible_device_count // replica_world_size |
| 73 | elif replica_world_size < visible_device_count: |
| 74 | # A node can host multiple complete TP/PP/PCP replicas. |
| 75 | local_dp_capacity = visible_device_count // replica_world_size |
| 76 | if visible_device_count % replica_world_size != 0: |
| 77 | logger.warning_once( |
| 78 | "XPU external LB cannot evenly divide " |
| 79 | "%d visible devices into TP/PP/PCP replicas of " |
| 80 | "size %d. This node can host %d complete DP " |
| 81 | "replicas, leaving %d visible devices unused.", |
| 82 | visible_device_count, |
| 83 | replica_world_size, |
| 84 | local_dp_capacity, |
| 85 | visible_device_count % replica_world_size, |
| 86 | ) |
| 87 | elif replica_world_size == visible_device_count: |
| 88 | # A node hosts exactly one complete TP/PP/PCP replica. |
| 89 | local_dp_capacity = 1 |
| 90 | logger.warning_once( |
| 91 | "XPU external LB sees exactly enough devices for one " |
| 92 | "TP/PP/PCP replica. This may be the intended " |
| 93 | "configuration, but it may also indicate that device " |
| 94 | "visibility is misconfigured. Every DP rank must see " |
| 95 | "all XPU devices on its node; consider removing " |
| 96 | "ZE_AFFINITY_MASK or setting it to expose the complete " |
| 97 | "device set." |
| 98 | ) |
| 99 | else: |
| 100 | # The topology says single-node, but the replica does not fit. |
| 101 | raise ValueError( |
| 102 | f"TP/PP/PCP replica size ({replica_world_size}) exceeds " |
| 103 | f"the number of visible XPU devices ({visible_device_count}), " |
| 104 | "but nnodes_within_dp is 1. Configure the multi-node " |
| 105 | "topology, or ensure every DP rank can see all devices " |
| 106 | "on its node." |
| 107 | ) |
| 108 | # Strip the node component off the global DP index to get this |
| 109 | # engine's slot on its own node. Assumes the launcher assigns |
| 110 | # DP ranks to nodes in contiguous blocks (node 0 gets ranks |
| 111 | # 0..capacity-1, and so on), which is what the usual sequential |
| 112 | # and one-pod-per-rank deployments do. A round-robin or |
| 113 | # unbalanced assignment would silently map two engines onto the |
| 114 | # same device. |
| 115 | dp_local_rank = parallel_config.data_parallel_index % local_dp_capacity |
| 116 | |
| 117 | self.local_rank += dp_local_rank * replica_world_size |
| 118 | |
| 119 | device = self.device_config.device |
| 120 | if ( |
| 121 | isinstance(device, torch.device) |
| 122 | and device.type == "xpu" |
| 123 | and current_platform.is_xpu() |
| 124 | ): |
| 125 | assigned_physical_gpu_ids = self.parallel_config.assigned_physical_gpu_ids |
| 126 | if assigned_physical_gpu_ids is not None: |
| 127 | from vllm.platforms.interface import set_assigned_physical_gpu_ids |
| 128 | |
| 129 | set_assigned_physical_gpu_ids(assigned_physical_gpu_ids) |
| 130 | assert self.local_rank < len(assigned_physical_gpu_ids), ( |
| 131 | f"local_rank {self.local_rank} is out of bounds for " |
| 132 | f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}" |
| 133 | ) |
| 134 | # NOTE: local_world_size is derived from parallel_config.nnodes, |
| 135 | # which is only set for the "mp" multi-node backend. With the |
| 136 | # "ray"/"external_launcher" backends nnodes stays 1, so |
| 137 | # local_world_size collapses to the full world_size and this |
| 138 | # check wrongly fires on cross-node deployments. |
| 139 | # assigned_physical_gpu_ids is already per-node and the |
| 140 | # local_rank bound above fully validates the mapping for |
| 141 | # these backends, so skip the check for them. |
| 142 | if parallel_config.distributed_executor_backend not in ( |
| 143 | "ray", |
| 144 | "external_launcher", |
| 145 | ): |
| 146 | assert parallel_config.local_world_size <= len( |
| 147 | assigned_physical_gpu_ids |
| 148 | ), ( |
| 149 | f"local_world_size ({parallel_config.local_world_size})" |
| 150 | " exceeds assigned_physical_gpu_ids count " |
| 151 | f"({len(assigned_physical_gpu_ids)})" |
| 152 | ) |
| 153 | else: |
| 154 | assert self.local_rank < torch.accelerator.device_count(), ( |
| 155 | f"DP adjusted local rank {self.local_rank} is out of " |
| 156 | f"bounds for {torch.accelerator.device_count()} devices." |
| 157 | ) |
| 158 | |
| 159 | visible_device_index = ( |
| 160 | current_platform.logical_device_id_to_visible_device_id(self.local_rank) |
| 161 | ) |
| 162 | self.device = torch.device(f"xpu:{visible_device_index}") |
| 163 | torch.accelerator.set_device_index(self.device) |
| 164 | current_platform.check_if_supports_dtype(self.model_config.dtype) |
| 165 | torch.accelerator.empty_cache() |
| 166 | self.init_gpu_memory = torch.xpu.get_device_properties( |
| 167 | visible_device_index |
| 168 | ).total_memory |
| 169 | else: |
| 170 | raise RuntimeError(f"Unsupported device type: {self.device_config.device}") |
| 171 | |
| 172 | ENV_CCL_ATL_TRANSPORT = os.getenv("CCL_ATL_TRANSPORT", "ofi") |
| 173 | ENV_LOCAL_WORLD_SIZE = os.getenv( |
| 174 | "LOCAL_WORLD_SIZE", str(self.parallel_config.world_size) |
| 175 | ) |
| 176 | os.environ["CCL_ATL_TRANSPORT"] = ENV_CCL_ATL_TRANSPORT |
| 177 | os.environ["LOCAL_WORLD_SIZE"] = ENV_LOCAL_WORLD_SIZE |
| 178 | os.environ["LOCAL_RANK"] = str(self.local_rank) |
| 179 | |
| 180 | init_worker_distributed_environment( |
| 181 | self.vllm_config, |
| 182 | self.rank, |
| 183 | self.distributed_init_method, |
| 184 | self.local_rank, |
| 185 | current_platform.dist_backend, |
| 186 | ) |
| 187 | |
| 188 | # oneCCL warm-up; only meaningful for multi-device runs. Requiring it |
| 189 | # with a single worker breaks platforms where oneCCL cannot enumerate |
| 190 | # device topology (e.g. paravirtualized GPUs). |
| 191 | if ( |
| 192 | self.parallel_config.world_size > 1 |
| 193 | and torch.distributed.is_xccl_available() |
| 194 | ): |
| 195 | torch.distributed.all_reduce(torch.zeros(1).xpu()) |
| 196 | |
| 197 | if self.use_v2_model_runner: |
| 198 | logger.info_once("Using V2 Model Runner") |
| 199 | |
| 200 | # Set random seed. |
| 201 | set_random_seed(self.model_config.seed) |
| 202 | |
| 203 | # Now take memory snapshot after NCCL is initialized |
| 204 | gc.collect() |
| 205 | torch.accelerator.empty_cache() |
| 206 | |
| 207 | # take current memory snapshot |
| 208 | self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device) |
| 209 | self.requested_memory = request_memory(init_snapshot, self.cache_config) |
| 210 | logger.debug("worker init memory snapshot: %r", self.init_snapshot) |
| 211 | logger.debug( |
| 212 | "worker requested memory: %sGiB", format_gib(self.requested_memory) |
| 213 | ) |
| 214 | |
| 215 | self._init_workspace_and_model_runner() |
| 216 | |
| 217 | def _make_model_runner(self): |
| 218 | if self.use_v2_model_runner: |
| 219 | model_runner = ( |
| 220 | XPUMMEncoderModelRunner |
| 221 | if self.vllm_config.is_mm_encoder_only |
| 222 | else XPUModelRunnerV2 |
| 223 | ) |
| 224 | else: |
| 225 | model_runner = XPUModelRunner # type: ignore[assignment] |
| 226 | return model_runner(self.vllm_config, self.device) # type: ignore |
| 227 | |
| 228 | def shutdown(self) -> None: |
| 229 | logger.info( |
| 230 | "XPUWorker shutdown: cleaning up (rank=%d, local_rank=%d)", |
| 231 | self.rank, |
| 232 | self.local_rank, |
| 233 | ) |
| 234 | super().shutdown() |
| 235 | from vllm.device_allocator.xpumem import XpuMemAllocator |
| 236 | |
| 237 | if XpuMemAllocator.instance is not None: |
| 238 | XpuMemAllocator.instance.release_pools() |
| 239 | logger.info( |
| 240 | "XPUWorker shutdown: done (rank=%d, local_rank=%d)", |
| 241 | self.rank, |
| 242 | self.local_rank, |
| 243 | ) |