| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | |
| 5 | import torch |
| 6 | import torch.distributed as dist |
| 7 | from torch.distributed import ProcessGroup |
| 8 | |
| 9 | import vllm.envs as envs |
| 10 | from vllm.logger import init_logger |
| 11 | |
| 12 | from .base_device_communicator import DeviceCommunicatorBase |
| 13 | |
| 14 | logger = init_logger(__name__) |
| 15 | |
| 16 | |
| 17 | class XpuCommunicator(DeviceCommunicatorBase): |
| 18 | def __init__( |
| 19 | self, |
| 20 | cpu_group: ProcessGroup, |
| 21 | device: torch.device | None = None, |
| 22 | device_group: ProcessGroup | None = None, |
| 23 | unique_name: str = "", |
| 24 | use_all2all: bool = False, |
| 25 | ): |
| 26 | super().__init__( |
| 27 | cpu_group, device, device_group, unique_name, use_all2all=use_all2all |
| 28 | ) |
| 29 | self.ca_comm: None = None |
| 30 | if self.use_all2all: |
| 31 | if self.all2all_backend in ("naive", "allgather_reducescatter"): |
| 32 | from .all2all import AgRsAll2AllManager |
| 33 | |
| 34 | self.all2all_manager = AgRsAll2AllManager(self.cpu_group) |
| 35 | logger.info("Using AgRs manager on XPU device.") |
| 36 | |
| 37 | else: # type: ignore[has-type] |
| 38 | logger.warning( |
| 39 | "`%s` all2all manager is not supported on XPU. " |
| 40 | "Falling back to AgRs manager for XPU, " |
| 41 | "which is the Default backend", |
| 42 | self.all2all_backend, # type: ignore[has-type] |
| 43 | ) |
| 44 | from .all2all import AgRsAll2AllManager |
| 45 | |
| 46 | self.all2all_manager = AgRsAll2AllManager(self.cpu_group) |
| 47 | logger.info("Using AgRs manager on XPU device.") |
| 48 | |
| 49 | def _fixed_rank_sum(self, input_: torch.Tensor) -> torch.Tensor: |
| 50 | flat_input = input_.reshape(-1) |
| 51 | gathered = torch.empty( |
| 52 | (self.world_size, flat_input.numel()), |
| 53 | dtype=input_.dtype, |
| 54 | device=input_.device, |
| 55 | ) |
| 56 | dist.all_gather_single(gathered.view(-1), flat_input, group=self.device_group) |
| 57 | output = gathered[0].clone() |
| 58 | for rank in range(1, self.world_size): |
| 59 | output.add_(gathered[rank]) |
| 60 | return output.view(input_.shape) |
| 61 | |
| 62 | def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: |
| 63 | if envs.VLLM_BATCH_INVARIANT: |
| 64 | # XCCL changes its reduction order with message size. Gather first |
| 65 | # and accumulate in rank order, independent of the token batch. |
| 66 | return self._fixed_rank_sum(input_) |
| 67 | |
| 68 | output = input_.clone() |
| 69 | dist.all_reduce(output, group=self.device_group) |
| 70 | return output |
| 71 | |
| 72 | def reduce_scatter(self, input_: torch.Tensor, dim: int = -1): |
| 73 | world_size = self.world_size |
| 74 | |
| 75 | if dim < 0: |
| 76 | # Convert negative dim to positive. |
| 77 | dim += input_.dim() |
| 78 | |
| 79 | # Note: This will produce an incorrect answer if we don't make |
| 80 | # the input_tensor contiguous. Possible bug in reduce_scatter_tensor? |
| 81 | input_tensor = input_.movedim(0, dim).contiguous() |
| 82 | |
| 83 | assert input_tensor.shape[0] % world_size == 0 |
| 84 | chunk_size = input_tensor.shape[0] // world_size |
| 85 | output_shape = (chunk_size,) + input_tensor.shape[1:] |
| 86 | |
| 87 | output = torch.empty( |
| 88 | output_shape, dtype=input_tensor.dtype, device=input_tensor.device |
| 89 | ) |
| 90 | |
| 91 | if envs.VLLM_BATCH_INVARIANT: |
| 92 | reduced = self._fixed_rank_sum(input_tensor) |
| 93 | start = self.rank_in_group * chunk_size |
| 94 | output.copy_(reduced.narrow(0, start, chunk_size)) |
| 95 | else: |
| 96 | dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group) |
| 97 | |
| 98 | # Reshape before returning |
| 99 | return output.movedim(0, dim).contiguous() |
| 100 | |
| 101 | def reduce_scatterv( |
| 102 | self, input_: torch.Tensor, dim: int = -1, sizes: list[int] | None = None |
| 103 | ): |
| 104 | world_size = self.world_size |
| 105 | |
| 106 | if dim < 0: |
| 107 | # Convert negative dim to positive. |
| 108 | dim += input_.dim() |
| 109 | |
| 110 | # Note: This will produce an incorrect answer if we don't make |
| 111 | # the input_tensor contiguous. Possible bug in reduce_scatter_tensor? |
| 112 | input_tensor = input_.movedim(0, dim).contiguous() |
| 113 | |
| 114 | if sizes is not None: |
| 115 | assert len(sizes) == world_size |
| 116 | assert input_tensor.shape[0] == sum(sizes) |
| 117 | chunk_size = sizes[self.rank_in_group] |
| 118 | else: |
| 119 | assert input_tensor.shape[0] % world_size == 0 |
| 120 | chunk_size = input_tensor.shape[0] // world_size |
| 121 | output_shape = (chunk_size,) + input_tensor.shape[1:] |
| 122 | |
| 123 | output = torch.empty( |
| 124 | output_shape, dtype=input_tensor.dtype, device=input_tensor.device |
| 125 | ) |
| 126 | if envs.VLLM_BATCH_INVARIANT: |
| 127 | reduced = self._fixed_rank_sum(input_tensor) |
| 128 | start = ( |
| 129 | self.rank_in_group * chunk_size |
| 130 | if sizes is None |
| 131 | else sum(sizes[: self.rank_in_group]) |
| 132 | ) |
| 133 | output.copy_(reduced.narrow(0, start, chunk_size)) |
| 134 | elif sizes is None or sizes.count(sizes[0]) == len(sizes): |
| 135 | dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group) |
| 136 | elif torch.xpu.device_count() < self.world_size: |
| 137 | # uneven reduce_scatter is not supported when the |
| 138 | # process cannot see every device in the group ( |
| 139 | # e.g. external-LBDP). |
| 140 | max_size = max(sizes) |
| 141 | padded = input_tensor.new_zeros( |
| 142 | (world_size * max_size,) + input_tensor.shape[1:] |
| 143 | ) |
| 144 | for rank, split in enumerate(input_tensor.split(sizes, dim=0)): |
| 145 | padded[rank * max_size : rank * max_size + sizes[rank]] = split |
| 146 | padded_output = input_tensor.new_empty((max_size,) + input_tensor.shape[1:]) |
| 147 | dist.reduce_scatter_tensor(padded_output, padded, group=self.device_group) |
| 148 | output = padded_output[:chunk_size] |
| 149 | else: |
| 150 | input_splits = list(input_tensor.split(sizes, dim=0)) |
| 151 | dist.reduce_scatter(output, input_splits, group=self.device_group) |
| 152 | # Reshape before returning |
| 153 | return output.movedim(0, dim).contiguous() |
| 154 | |
| 155 | def all_gatherv( |
| 156 | self, |
| 157 | input_: torch.Tensor | list[torch.Tensor], |
| 158 | dim: int = 0, |
| 159 | sizes: list[int] | None = None, |
| 160 | ): |
| 161 | if dim != 0: |
| 162 | raise NotImplementedError("only dim 0 all-gatherv is supported") |
| 163 | world_size = self.world_size |
| 164 | |
| 165 | # 'sizes' is not needed if all inputs in the same group have the same |
| 166 | # shape |
| 167 | if sizes is not None and all(s == sizes[0] for s in sizes): |
| 168 | sizes = None |
| 169 | |
| 170 | def _all_gather_single(input_: torch.Tensor, sizes: list[int] | None = None): |
| 171 | input_size = input_.size() |
| 172 | if sizes is not None: |
| 173 | assert len(sizes) == world_size |
| 174 | assert input_.shape[dim] == sizes[self.rank_in_group], ( |
| 175 | f"{input_.shape[dim]} != {sizes[self.rank_in_group]}" |
| 176 | ) |
| 177 | output_size = (sum(sizes),) + input_size[1:] |
| 178 | else: |
| 179 | output_size = (input_size[0] * world_size,) + input_size[1:] |
| 180 | # Allocate output tensor. |
| 181 | output_tensor = torch.empty( |
| 182 | output_size, dtype=input_.dtype, device=input_.device |
| 183 | ) |
| 184 | |
| 185 | if sizes is not None: |
| 186 | all_gather_list = [] |
| 187 | for size in sizes: |
| 188 | all_gather_list.append( |
| 189 | torch.empty( |
| 190 | (size,) + input_.shape[1:], |
| 191 | dtype=input_.dtype, |
| 192 | device=input_.device, |
| 193 | ) |
| 194 | ) |
| 195 | dist.all_gather(all_gather_list, input_, group=self.device_group) |
| 196 | output_tensor = torch.cat(all_gather_list, dim=0) |
| 197 | else: |
| 198 | dist.all_gather([output_tensor], input_, group=self.device_group) |
| 199 | return output_tensor |
| 200 | |
| 201 | if isinstance(input_, torch.Tensor): |
| 202 | return _all_gather_single(input_, sizes) |
| 203 | |
| 204 | output_list = [] |
| 205 | for inp in input_: |
| 206 | output_list.append(_all_gather_single(inp, sizes=sizes)) |
| 207 | return output_list |
| 208 | |
| 209 | def gather( |
| 210 | self, input_: torch.Tensor, dst: int = 0, dim: int = -1 |
| 211 | ) -> torch.Tensor | None: |
| 212 | assert -input_.dim() <= dim < input_.dim(), ( |
| 213 | f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" |
| 214 | ) |
| 215 | if dim < 0: |
| 216 | # Convert negative dim to positive. |
| 217 | dim += input_.dim() |
| 218 | # For xpu path, gather doesn't work properly together with ray |
| 219 | # cluster so we use all_gather instead for now. |
| 220 | input_size = input_.size() |
| 221 | # Allocate output tensor. |
| 222 | output_tensor = torch.empty( |
| 223 | (self.world_size,) + input_size, dtype=input_.dtype, device=input_.device |
| 224 | ) |
| 225 | # All-gather. |
| 226 | dist.all_gather_single(output_tensor, input_, group=self.device_group) |
| 227 | if self.rank_in_group == dst: |
| 228 | # Reshape |
| 229 | output_tensor = output_tensor.movedim(0, dim) |
| 230 | output_tensor = output_tensor.reshape( |
| 231 | input_size[:dim] |
| 232 | + (self.world_size * input_size[dim],) |
| 233 | + input_size[dim + 1 :] |
| 234 | ) |
| 235 | else: |
| 236 | output_tensor = None |
| 237 | return output_tensor |
| 238 | |
| 239 | def broadcast(self, input_: torch.Tensor, src: int = 0) -> None: |
| 240 | dist.broadcast(input_, src=src, group=self.device_group) |
| 241 | |
| 242 | def dispatch_router_logits( |
| 243 | self, |
| 244 | hidden_states: torch.Tensor, |
| 245 | router_logits: torch.Tensor, |
| 246 | is_sequence_parallel: bool = False, |
| 247 | extra_tensors: list[torch.Tensor] | None = None, |
| 248 | ) -> ( |
| 249 | tuple[torch.Tensor, torch.Tensor] |
| 250 | | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 251 | ): |
| 252 | """Dispatch the hidden states and router logits to the appropriate device. |
| 253 | This is a no-op in the base class. |
| 254 | """ |
| 255 | assert self.all2all_manager is not None |
| 256 | return self.all2all_manager.dispatch_router_logits( |
| 257 | hidden_states, |
| 258 | router_logits, |
| 259 | is_sequence_parallel, |
| 260 | extra_tensors, |
| 261 | ) |
| 262 | |
| 263 | def dispatch( |
| 264 | self, |
| 265 | hidden_states: torch.Tensor, |
| 266 | topk_weights: torch.Tensor, |
| 267 | topk_ids: torch.Tensor, |
| 268 | is_sequence_parallel: bool = False, |
| 269 | extra_tensors: list[torch.Tensor] | None = None, |
| 270 | ) -> ( |
| 271 | tuple[torch.Tensor, torch.Tensor, torch.Tensor] |
| 272 | | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 273 | ): |
| 274 | """Dispatch the hidden states and topk weights/ids to the appropriate device. |
| 275 | This is a no-op in the base class. |
| 276 | """ |
| 277 | assert self.all2all_manager is not None |
| 278 | return self.all2all_manager.dispatch( |
| 279 | hidden_states, |
| 280 | topk_weights, |
| 281 | topk_ids, |
| 282 | is_sequence_parallel, |
| 283 | extra_tensors=extra_tensors, |
| 284 | ) |
| 285 | |
| 286 | def combine( |
| 287 | self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False |
| 288 | ) -> torch.Tensor: |
| 289 | """Combine the hidden states and router logits from the appropriate device. |
| 290 | This is a no-op in the base class. |
| 291 | """ |
| 292 | assert self.all2all_manager is not None |
| 293 | return self.all2all_manager.combine( |
| 294 | hidden_states, |
| 295 | is_sequence_parallel, |
| 296 | ) |