class XPUWorker(Worker):
"""A XPU worker class."""
def __init__(
self,
vllm_config: VllmConfig,
local_rank: int,
rank: int,
distributed_init_method: str,
is_driver_worker: bool = False,
):
super().__init__(
vllm_config, local_rank, rank, distributed_init_method, is_driver_worker
)
device_config = self.device_config
assert device_config.device_type == "xpu"
assert current_platform.is_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.data_parallel_external_lb
or 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
replica_world_size = parallel_config.world_size
visible_device_count = torch.accelerator.device_count()
if parallel_config.data_parallel_external_lb:
if parallel_config.nnodes_within_dp > 1:
# The replica spans nodes, so use its node-local shard.
replica_world_size = parallel_config.local_world_size
if replica_world_size > visible_device_count:
raise ValueError(
f"Local TP/PP/PCP replica size ({replica_world_size}) "
"exceeds the number of visible XPU devices "
f"({visible_device_count})."
)
local_dp_capacity = visible_device_count // replica_world_size
elif replica_world_size < visible_device_count:
# A node can host multiple complete TP/PP/PCP replicas.
local_dp_capacity = visible_device_count // replica_world_size
if visible_device_count % replica_world_size != 0:
logger.warning_once(
"XPU external LB cannot evenly divide "
"%d visible devices into TP/PP/PCP replicas of "
"size %d. This node can host %d complete DP "
"replicas, leaving %d visible devices unused.",
visible_device_count,
replica_world_size,
local_dp_capacity,
visible_device_count % replica_world_size,
)
elif replica_world_size == visible_device_count:
# A node hosts exactly one complete TP/PP/PCP replica.
local_dp_capacity = 1
logger.warning_once(
"XPU external LB sees exactly enough devices for one "
"TP/PP/PCP replica. This may be the intended "
"configuration, but it may also indicate that device "
"visibility is misconfigured. Every DP rank must see "
"all XPU devices on its node; consider removing "
"ZE_AFFINITY_MASK or setting it to expose the complete "
"device set."
)
else:
# The topology says single-node, but the replica does not fit.
raise ValueError(
f"TP/PP/PCP replica size ({replica_world_size}) exceeds "
f"the number of visible XPU devices ({visible_device_count}), "
"but nnodes_within_dp is 1. Configure the multi-node "
"topology, or ensure every DP rank can see all devices "
"on its node."
)
# Strip the node component off the global DP index to get this
# engine's slot on its own node. Assumes the launcher assigns
# DP ranks to nodes in contiguous blocks (node 0 gets ranks
# 0..capacity-1, and so on), which is what the usual sequential
# and one-pod-per-rank deployments do. A round-robin or
# unbalanced assignment would silently map two engines onto the
# same device.
dp_local_rank = parallel_config.data_parallel_index % local_dp_capacity
self.local_rank += dp_local_rank * replica_world_size
device = self.device_config.device
if (
isinstance(device, torch.device)
and device.type == "xpu"
and current_platform.is_xpu()
):
assigned_physical_gpu_ids = self.parallel_config.assigned_physical_gpu_ids
if assigned_physical_gpu_ids is not None:
from vllm.platforms.interface import set_assigned_physical_gpu_ids
set_assigned_physical_gpu_ids(assigned_physical_gpu_ids)
assert self.local_rank < len(assigned_physical_gpu_ids), (
f"local_rank {self.local_rank} is out of bounds for "
f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}"
)
# NOTE: local_world_size is derived from parallel_config.nnodes,
# which is only set for the "mp" multi-node backend. With the
# "ray"/"external_launcher" backends nnodes stays 1, so
# local_world_size collapses to the full world_size and this
# check wrongly fires on cross-node deployments.
# assigned_physical_gpu_ids is already per-node and the
# local_rank bound above fully validates the mapping for
# these backends, so skip the check for them.
if parallel_config.distributed_executor_backend not in (
"ray",
"external_launcher",
):
assert parallel_config.local_world_size <= len(
assigned_physical_gpu_ids
), (
f"local_world_size ({parallel_config.local_world_size})"
" exceeds assigned_physical_gpu_ids count "
f"({len(assigned_physical_gpu_ids)})"
)
else:
assert self.local_rank < torch.accelerator.device_count(), (
f"DP adjusted local rank {self.local_rank} is out of "
f"bounds for {torch.accelerator.device_count()} devices."
)
visible_device_index = (
current_platform.logical_device_id_to_visible_device_id(self.local_rank)
)
self.device = torch.device(f"xpu:{visible_device_index}")
torch.accelerator.set_device_index(self.device)
current_platform.check_if_supports_dtype(self.model_config.dtype)
torch.accelerator.empty_cache()
self.init_gpu_memory = torch.xpu.get_device_properties(
visible_device_index
).total_memory
else:
raise RuntimeError(f"Unsupported device type: {self.device_config.device}")
ENV_CCL_ATL_TRANSPORT = os.getenv("CCL_ATL_TRANSPORT", "ofi")
ENV_LOCAL_WORLD_SIZE = os.getenv(
"LOCAL_WORLD_SIZE", str(self.parallel_config.world_size)
)
os.environ["CCL_ATL_TRANSPORT"] = ENV_CCL_ATL_TRANSPORT
os.environ["LOCAL_WORLD_SIZE"] = ENV_LOCAL_WORLD_SIZE
os.environ["LOCAL_RANK"] = str(self.local_rank)
init_worker_distributed_environment(
self.vllm_config,
self.rank,
self.distributed_init_method,
self.local_rank,
current_platform.dist_backend,
)
# oneCCL warm-up; only meaningful for multi-device runs. Requiring it
# with a single worker breaks platforms where oneCCL cannot enumerate
# device topology (e.g. paravirtualized GPUs).
if (
self.parallel_config.world_size > 1
and torch.distributed.is_xccl_available()
):
torch.distributed.all_reduce(torch.zeros(1).xpu())
if self.use_v2_model_runner:
logger.info_once("Using V2 Model Runner")
# Set random seed.
set_random_seed(self.model_config.seed)
# Now take memory snapshot after NCCL is initialized
gc.collect()
torch.accelerator.empty_cache()
# take current memory snapshot
self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
self.requested_memory = request_memory(init_snapshot, self.cache_config)
logger.debug("worker init memory snapshot: %r", self.init_snapshot)
logger.debug(
"worker requested memory: %sGiB", format_gib(self.requested_memory)
)
self._init_workspace_and_model_runner()
def _make_model_runner(self):
if self.use_v2_model_runner:
model_runner = (
XPUMMEncoderModelRunner
if self.vllm_config.is_mm_encoder_only
else XPUModelRunnerV2
)
else:
model_runner = XPUModelRunner # type: ignore[assignment]
return model_runner(self.vllm_config, self.device) # type: ignore
def shutdown(self) -> None:
logger.info(
"XPUWorker shutdown: cleaning up (rank=%d, local_rank=%d)",
self.rank,
self.local_rank,
)
super().shutdown()
from vllm.device_allocator.xpumem import XpuMemAllocator
if XpuMemAllocator.instance is not None:
XpuMemAllocator.instance.release_pools()
logger.info(
"XPUWorker shutdown: done (rank=%d, local_rank=%d)",
self.rank,
self.local_rank,
)