| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import os |
| 5 | from typing import Any |
| 6 | |
| 7 | import torch |
| 8 | from torch.distributed import ProcessGroup |
| 9 | |
| 10 | from vllm.distributed.utils import pickle |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.platforms.interface import CpuArchEnum |
| 14 | |
| 15 | from .base_device_communicator import DeviceCommunicatorBase |
| 16 | |
| 17 | logger = init_logger(__name__) |
| 18 | |
| 19 | |
| 20 | class CpuCommunicator(DeviceCommunicatorBase): |
| 21 | def __init__( |
| 22 | self, |
| 23 | cpu_group: ProcessGroup, |
| 24 | device: torch.device | None = None, |
| 25 | device_group: ProcessGroup | None = None, |
| 26 | unique_name: str = "", |
| 27 | use_all2all: bool = False, |
| 28 | ): |
| 29 | super().__init__( |
| 30 | cpu_group, device, device_group, unique_name, use_all2all=use_all2all |
| 31 | ) |
| 32 | self.dist_module = torch.distributed |
| 33 | |
| 34 | if ( |
| 35 | ( |
| 36 | current_platform.get_cpu_architecture() == CpuArchEnum.X86 |
| 37 | or current_platform.get_cpu_architecture() == CpuArchEnum.ARM |
| 38 | or current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC |
| 39 | or current_platform.get_cpu_architecture() == CpuArchEnum.S390X |
| 40 | ) |
| 41 | and hasattr(torch.ops._C, "init_shm_manager") |
| 42 | and (unique_name.startswith("tp") or unique_name.startswith("pp")) |
| 43 | and self._all_group_ranks_share_shm_group_name() |
| 44 | ): |
| 45 | self.dist_module = _CPUSHMDistributed(self) |
| 46 | elif unique_name.startswith("tp") or unique_name.startswith("pp"): |
| 47 | logger.info( |
| 48 | "CPU SHM communicator disabled for group %s: ranks do not share " |
| 49 | "the same SHM group name, falling back to torch.distributed.", |
| 50 | unique_name, |
| 51 | ) |
| 52 | |
| 53 | # send/recv tensor_dict is only supported through the SHM communicator backend |
| 54 | self.supports_tensor_dict = isinstance(self.dist_module, _CPUSHMDistributed) |
| 55 | |
| 56 | if self.use_all2all: |
| 57 | if self.all2all_backend not in ( |
| 58 | "naive", |
| 59 | "allgather_reducescatter", |
| 60 | ): # type: ignore[has-type] |
| 61 | logger.warning( |
| 62 | "`%s` all2all manager is not supported on CPU. " |
| 63 | "Falling back to `allgather_reducescatter` manager.", |
| 64 | self.all2all_backend, # type: ignore[has-type] |
| 65 | ) |
| 66 | from .all2all import AgRsAll2AllManager |
| 67 | |
| 68 | self.all2all_manager = AgRsAll2AllManager(self.cpu_group) |
| 69 | logger.info("Using allgather_reducescatter all2all manager.") |
| 70 | |
| 71 | def _all_group_ranks_share_shm_group_name(self) -> bool: |
| 72 | """CPUSHM requires all ranks in this group to agree on one SHM group name. |
| 73 | This is a lightweight consistency check for VLLM_DIST_IDENT/name inputs. |
| 74 | """ |
| 75 | local_name = _CPUSHMDistributed.make_group_name(self) |
| 76 | names: list[str] = [""] * self.world_size |
| 77 | torch.distributed.all_gather_object( |
| 78 | names, |
| 79 | local_name, |
| 80 | group=self.device_group, |
| 81 | ) |
| 82 | return len(set(names)) == 1 |
| 83 | |
| 84 | def all_reduce(self, input_): |
| 85 | self.dist_module.all_reduce(input_, group=self.device_group) |
| 86 | return input_ |
| 87 | |
| 88 | def gather( |
| 89 | self, input_: torch.Tensor, dst: int = 0, dim: int = -1 |
| 90 | ) -> torch.Tensor | None: |
| 91 | """NOTE: We assume that the input tensor is on the same device across |
| 92 | all the ranks. |
| 93 | NOTE: `dst` is the local rank of the destination rank. |
| 94 | """ |
| 95 | world_size = self.world_size |
| 96 | assert -input_.dim() <= dim < input_.dim(), ( |
| 97 | f"Invalid dim ({dim}) for input tensor with shape {input_.size()}" |
| 98 | ) |
| 99 | if dim < 0: |
| 100 | # Convert negative dim to positive. |
| 101 | dim += input_.dim() |
| 102 | |
| 103 | # Allocate output tensor. |
| 104 | if self.rank_in_group == dst: |
| 105 | gather_list = [torch.empty_like(input_) for _ in range(world_size)] |
| 106 | else: |
| 107 | gather_list = None |
| 108 | |
| 109 | # Gather. |
| 110 | self.dist_module.gather( |
| 111 | input_, gather_list, dst=self.ranks[dst], group=self.device_group |
| 112 | ) |
| 113 | |
| 114 | if self.rank_in_group == dst: |
| 115 | output_tensor = torch.cat(gather_list, dim=dim) |
| 116 | else: |
| 117 | output_tensor = None |
| 118 | return output_tensor |
| 119 | |
| 120 | def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor: |
| 121 | if dim < 0: |
| 122 | # Convert negative dim to positive. |
| 123 | dim += input_.dim() |
| 124 | input_size = input_.size() |
| 125 | # NOTE: we have to use concat-style all-gather here, |
| 126 | # stack-style all-gather has compatibility issues with |
| 127 | # torch.compile . see https://github.com/pytorch/pytorch/issues/138795 |
| 128 | output_size = (input_size[0] * self.world_size,) + input_size[1:] |
| 129 | # Allocate output tensor. |
| 130 | output_tensor = torch.empty( |
| 131 | output_size, dtype=input_.dtype, device=input_.device |
| 132 | ) |
| 133 | # All-gather. |
| 134 | self.dist_module.all_gather_into_tensor( |
| 135 | output_tensor, input_, group=self.device_group |
| 136 | ) |
| 137 | |
| 138 | # Reshape |
| 139 | output_tensor = output_tensor.reshape((self.world_size,) + input_size) |
| 140 | output_tensor = output_tensor.movedim(0, dim) |
| 141 | output_tensor = output_tensor.reshape( |
| 142 | input_size[:dim] |
| 143 | + (self.world_size * input_size[dim],) |
| 144 | + input_size[dim + 1 :] |
| 145 | ) |
| 146 | return output_tensor |
| 147 | |
| 148 | def send_tensor_dict( |
| 149 | self, |
| 150 | tensor_dict: dict[str, torch.Tensor | Any], |
| 151 | dst: int, |
| 152 | ) -> None: |
| 153 | if not self.supports_tensor_dict: |
| 154 | raise NotImplementedError( |
| 155 | "CpuCommunicator does not support tensor dict fastpath with " |
| 156 | "torch.distributed backend." |
| 157 | ) |
| 158 | return self.dist_module.send_tensor_dict(tensor_dict, dst) |
| 159 | |
| 160 | def recv_tensor_dict( |
| 161 | self, |
| 162 | src: int, |
| 163 | ) -> dict[str, torch.Tensor | Any]: |
| 164 | if not self.supports_tensor_dict: |
| 165 | raise NotImplementedError( |
| 166 | "CpuCommunicator does not support tensor dict fastpath with " |
| 167 | "torch.distributed backend." |
| 168 | ) |
| 169 | return self.dist_module.recv_tensor_dict(src) |
| 170 | |
| 171 | def dispatch_router_logits( |
| 172 | self, |
| 173 | hidden_states: torch.Tensor, |
| 174 | router_logits: torch.Tensor, |
| 175 | is_sequence_parallel: bool = False, |
| 176 | extra_tensors: list[torch.Tensor] | None = None, |
| 177 | ) -> ( |
| 178 | tuple[torch.Tensor, torch.Tensor] |
| 179 | | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 180 | ): |
| 181 | """Dispatch the hidden states and router logits to the appropriate device. |
| 182 | This is a no-op in the base class. |
| 183 | """ |
| 184 | assert self.all2all_manager is not None |
| 185 | return self.all2all_manager.dispatch_router_logits( |
| 186 | hidden_states, |
| 187 | router_logits, |
| 188 | is_sequence_parallel, |
| 189 | extra_tensors, |
| 190 | ) |
| 191 | |
| 192 | def dispatch( |
| 193 | self, |
| 194 | hidden_states: torch.Tensor, |
| 195 | topk_weights: torch.Tensor, |
| 196 | topk_ids: torch.Tensor, |
| 197 | is_sequence_parallel: bool = False, |
| 198 | extra_tensors: list[torch.Tensor] | None = None, |
| 199 | ) -> ( |
| 200 | tuple[torch.Tensor, torch.Tensor, torch.Tensor] |
| 201 | | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]] |
| 202 | ): |
| 203 | """Dispatch the hidden states and topk weights/ids to the appropriate device. |
| 204 | This is a no-op in the base class. |
| 205 | """ |
| 206 | assert self.all2all_manager is not None |
| 207 | return self.all2all_manager.dispatch( |
| 208 | hidden_states, |
| 209 | topk_weights, |
| 210 | topk_ids, |
| 211 | is_sequence_parallel, |
| 212 | extra_tensors=extra_tensors, |
| 213 | ) |
| 214 | |
| 215 | def combine( |
| 216 | self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False |
| 217 | ) -> torch.Tensor: |
| 218 | """Combine the hidden states and router logits from the appropriate device. |
| 219 | This is a no-op in the base class. |
| 220 | """ |
| 221 | assert self.all2all_manager is not None |
| 222 | return self.all2all_manager.combine( |
| 223 | hidden_states, |
| 224 | is_sequence_parallel, |
| 225 | ) |
| 226 | |
| 227 | |
| 228 | class _CPUSHMDistributed: |
| 229 | def __init__(self, communicator: CpuCommunicator): |
| 230 | self.communicator = communicator |
| 231 | |
| 232 | self.group_name = self.make_group_name(communicator) |
| 233 | |
| 234 | self.handle = self._init_cpu_shm() |
| 235 | |
| 236 | @staticmethod |
| 237 | def make_group_name(communicator: CpuCommunicator) -> str: |
| 238 | instance_identifier = os.environ["VLLM_DIST_IDENT"] |
| 239 | unique_name = communicator.unique_name |
| 240 | instance_identifier = f"{instance_identifier}-{unique_name}" |
| 241 | group_ranks = [str(rank) for rank in communicator.ranks] |
| 242 | shm_group_identifier = f"[{'-'.join(group_ranks)}]" |
| 243 | return f"{instance_identifier}-{shm_group_identifier}-cpushm" |
| 244 | |
| 245 | def _init_cpu_shm(self) -> int: |
| 246 | thread_num_tensor = torch.tensor( |
| 247 | [torch.get_num_threads()], |
| 248 | dtype=torch.int64, |
| 249 | ) |
| 250 | torch.distributed.all_reduce( |
| 251 | thread_num_tensor, |
| 252 | op=torch.distributed.ReduceOp.MIN, |
| 253 | group=self.communicator.device_group, |
| 254 | ) |
| 255 | thread_num = thread_num_tensor.item() |
| 256 | |
| 257 | handle = torch.ops._C.init_shm_manager( |
| 258 | self.group_name, |
| 259 | self.communicator.world_size, |
| 260 | self.communicator.rank, |
| 261 | thread_num, |
| 262 | ) |
| 263 | torch.distributed.barrier(self.communicator.device_group) |
| 264 | torch.ops._C.join_shm_manager( |
| 265 | handle, |
| 266 | self.group_name, |
| 267 | ) |
| 268 | torch.distributed.barrier(self.communicator.device_group) |
| 269 | |
| 270 | return handle |
| 271 | |
| 272 | def all_reduce( |
| 273 | self, input: torch.Tensor, group: ProcessGroup | None = None |
| 274 | ) -> None: |
| 275 | torch.ops._C.shm_allreduce(self.handle, input) |
| 276 | |
| 277 | def gather( |
| 278 | self, |
| 279 | input: torch.Tensor, |
| 280 | gather_list: list[torch.Tensor] | None, |
| 281 | dst: int = -1, |
| 282 | group: ProcessGroup | None = None, |
| 283 | ) -> None: |
| 284 | # Note: different from the torch gather, here we use local dst rank. |
| 285 | torch.ops._C.shm_gather( |
| 286 | self.handle, |
| 287 | input, |
| 288 | gather_list, |
| 289 | torch.distributed.get_group_rank(group, dst), |
| 290 | ) |
| 291 | |
| 292 | def all_gather_into_tensor( |
| 293 | self, |
| 294 | output: torch.Tensor, |
| 295 | input: torch.Tensor, |
| 296 | group: ProcessGroup | None = None, |
| 297 | ) -> None: |
| 298 | torch.ops._C.shm_all_gather(self.handle, input, output) |
| 299 | |
| 300 | def send_tensor_dict( |
| 301 | self, |
| 302 | tensor_dict: dict[str, torch.Tensor | Any], |
| 303 | dst: int, |
| 304 | ) -> None: |
| 305 | key_list = list(tensor_dict.keys()) |
| 306 | value_list = list(tensor_dict.values()) |
| 307 | size_list = [] |
| 308 | for v in value_list: |
| 309 | if not isinstance(v, torch.Tensor): |
| 310 | raise RuntimeError("CpuCommunicator only supports sending tensors.") |
| 311 | size_list.append(v.size()) |
| 312 | key_size_tensor = torch.frombuffer( |
| 313 | pickle.dumps([key_list, size_list]), dtype=torch.uint8 |
| 314 | ) |
| 315 | value_list.append(key_size_tensor) |
| 316 | |
| 317 | torch.ops._C.shm_send_tensor_list(self.handle, value_list, dst) |
| 318 | |
| 319 | return None |
| 320 | |
| 321 | def recv_tensor_dict( |
| 322 | self, |
| 323 | src: int, |
| 324 | ) -> dict[str, torch.Tensor | Any]: |
| 325 | tensor_list = torch.ops._C.shm_recv_tensor_list(self.handle, src) |
| 326 | |
| 327 | value_list: list[torch.Tensor] = tensor_list[:-1] |
| 328 | key_size_tensor = tensor_list[-1] |
| 329 | |
| 330 | key_size = pickle.loads(key_size_tensor.numpy().tobytes()) |
| 331 | key_list = key_size[0] |
| 332 | size_list = key_size[1] |
| 333 | assert len(key_list) == len(size_list) |
| 334 | assert len(key_list) == len(value_list) |
| 335 | |
| 336 | tensor_dict: dict[str, torch.Tensor] = {} |
| 337 | for key, size, t in zip(key_list, size_list, value_list): |
| 338 | tensor_dict[key] = t.view(size) |
| 339 | return tensor_dict |