@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a0231b7c25
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 23:30 UTC

[platform] add base class for communicators (#13208)

Signed-off-by: youkaichao <youkaichao@gmail.com>

youkaichao committed Feb 16, 2025 at 22:14 UTC a0231b7c25d65d8552863671903356d4c431c6ad
13 files changed +364 -282
vllm/distributed/device_communicators/base_device_communicator.py new
+117
@@ -0,0 +1,117 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Optional
3 +
4 +import torch
5 +import torch.distributed as dist
6 +from torch.distributed import ProcessGroup
7 +
8 +
9 +class DeviceCommunicatorBase:
10 + """
11 + Base class for device-specific communicator.
12 + It can use the `cpu_group` to initialize the communicator.
13 + If the device has PyTorch integration (PyTorch can recognize its
14 + communication backend), the `device_group` will also be given.
15 + """
16 +
17 + def __init__(self,
18 + cpu_group: ProcessGroup,
19 + device: Optional[torch.device] = None,
20 + device_group: Optional[ProcessGroup] = None,
21 + unique_name: str = ""):
22 + self.device = device or torch.device("cpu")
23 + self.cpu_group = cpu_group
24 + self.device_group = device_group
25 + self.unique_name = unique_name
26 + self.rank = dist.get_rank(cpu_group)
27 + self.world_size = dist.get_world_size(cpu_group)
28 + self.ranks = dist.get_process_group_ranks(cpu_group)
29 + self.global_rank = dist.get_rank()
30 + self.global_world_size = dist.get_world_size()
31 + self.rank_in_group = dist.get_group_rank(self.cpu_group,
32 + self.global_rank)
33 +
34 + def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
35 + dist.all_reduce(input_, group=self.device_group)
36 + return input_
37 +
38 + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
39 + if dim < 0:
40 + # Convert negative dim to positive.
41 + dim += input_.dim()
42 + input_size = input_.size()
43 + # NOTE: we have to use concat-style all-gather here,
44 + # stack-style all-gather has compatibility issues with
45 + # torch.compile . see https://github.com/pytorch/pytorch/issues/138795
46 + output_size = (input_size[0] * self.world_size, ) + input_size[1:]
47 + # Allocate output tensor.
48 + output_tensor = torch.empty(output_size,
49 + dtype=input_.dtype,
50 + device=input_.device)
51 + # All-gather.
52 + dist.all_gather_into_tensor(output_tensor,
53 + input_,
54 + group=self.device_group)
55 + # Reshape
56 + output_tensor = output_tensor.reshape((self.world_size, ) + input_size)
57 + output_tensor = output_tensor.movedim(0, dim)
58 + output_tensor = output_tensor.reshape(input_size[:dim] +
59 + (self.world_size *
60 + input_size[dim], ) +
61 + input_size[dim + 1:])
62 + return output_tensor
63 +
64 + def gather(self,
65 + input_: torch.Tensor,
66 + dst: int = 0,
67 + dim: int = -1) -> Optional[torch.Tensor]:
68 + """
69 + NOTE: We assume that the input tensor is on the same device across
70 + all the ranks.
71 + NOTE: `dst` is the local rank of the destination rank.
72 + """
73 + world_size = self.world_size
74 + assert -input_.dim() <= dim < input_.dim(), (
75 + f"Invalid dim ({dim}) for input tensor with shape {input_.size()}")
76 + if dim < 0:
77 + # Convert negative dim to positive.
78 + dim += input_.dim()
79 +
80 + # Allocate output tensor.
81 + if self.rank_in_group == dst:
82 + gather_list = [torch.empty_like(input_) for _ in range(world_size)]
83 + else:
84 + gather_list = None
85 + # Gather.
86 + torch.distributed.gather(input_,
87 + gather_list,
88 + dst=self.ranks[dst],
89 + group=self.device_group)
90 + if self.rank_in_group == dst:
91 + output_tensor = torch.cat(gather_list, dim=dim)
92 + else:
93 + output_tensor = None
94 + return output_tensor
95 +
96 + def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
97 + """Sends a tensor to the destination rank in a non-blocking way"""
98 + """NOTE: `dst` is the local rank of the destination rank."""
99 + if dst is None:
100 + dst = (self.rank_in_group + 1) % self.world_size
101 + torch.distributed.send(tensor, self.ranks[dst], self.device_group)
102 +
103 + def recv(self,
104 + size: torch.Size,
105 + dtype: torch.dtype,
106 + src: Optional[int] = None) -> torch.Tensor:
107 + """Receives a tensor from the source rank."""
108 + """NOTE: `src` is the local rank of the source rank."""
109 + if src is None:
110 + src = (self.rank_in_group - 1) % self.world_size
111 +
112 + tensor = torch.empty(size, dtype=dtype, device=self.device)
113 + torch.distributed.recv(tensor, self.ranks[src], self.device_group)
114 + return tensor
115 +
116 + def destroy(self):
117 + pass
vllm/distributed/device_communicators/cpu_communicator.py new
+33
@@ -0,0 +1,33 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from typing import Optional
4 +
5 +import torch
6 +from torch.distributed import ProcessGroup
7 +
8 +from .base_device_communicator import DeviceCommunicatorBase
9 +
10 +
11 +class CpuCommunicator(DeviceCommunicatorBase):
12 +
13 + def __init__(self,
14 + cpu_group: ProcessGroup,
15 + device: Optional[torch.device] = None,
16 + device_group: Optional[ProcessGroup] = None,
17 + unique_name: str = ""):
18 + super().__init__(cpu_group, device, device_group, unique_name)
19 + self.ipex_available = False
20 + self.dist_module = torch.distributed
21 + try:
22 + import intel_extension_for_pytorch as ipex
23 + self.ipex_available = True
24 + self.dist_module = ipex.distributed
25 + except ImportError:
26 + """
27 + Intel IPEX not found. Falling back to PyTorch native
28 + all_reduce for CPU (e.g. MacOS)
29 + """
30 + pass
31 +
32 + def all_reduce(self, input_):
33 + return self.dist_module.all_reduce(input_, group=self.device_group)
vllm/distributed/device_communicators/cuda_communicator.py new
+106
@@ -0,0 +1,106 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +from typing import Optional
4 +
5 +import torch
6 +from torch.distributed import ProcessGroup
7 +
8 +from .base_device_communicator import DeviceCommunicatorBase
9 +
10 +
11 +class CudaCommunicator(DeviceCommunicatorBase):
12 +
13 + def __init__(self,
14 + cpu_group: ProcessGroup,
15 + device: Optional[torch.device] = None,
16 + device_group: Optional[ProcessGroup] = None,
17 + unique_name: str = ""):
18 + super().__init__(cpu_group, device, device_group, unique_name)
19 + if "pp" in unique_name:
20 + # pipeline parallel does not need custom allreduce
21 + use_custom_allreduce = False
22 + else:
23 + from vllm.distributed.parallel_state import (
24 + _ENABLE_CUSTOM_ALL_REDUCE)
25 + use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE
26 + use_pynccl = True
27 +
28 + self.use_pynccl = use_pynccl
29 + self.use_custom_allreduce = use_custom_allreduce
30 +
31 + # lazy import to avoid documentation build error
32 + from vllm.distributed.device_communicators.custom_all_reduce import (
33 + CustomAllreduce)
34 + from vllm.distributed.device_communicators.pynccl import (
35 + PyNcclCommunicator)
36 +
37 + self.pynccl_comm: Optional[PyNcclCommunicator] = None
38 + if use_pynccl and self.world_size > 1:
39 + self.pynccl_comm = PyNcclCommunicator(
40 + group=self.cpu_group,
41 + device=self.device,
42 + )
43 +
44 + self.ca_comm: Optional[CustomAllreduce] = None
45 + if use_custom_allreduce and self.world_size > 1:
46 + # Initialize a custom fast all-reduce implementation.
47 + self.ca_comm = CustomAllreduce(
48 + group=self.cpu_group,
49 + device=self.device,
50 + )
51 +
52 + def all_reduce(self, input_):
53 + # always try custom allreduce first,
54 + # and then pynccl.
55 + ca_comm = self.ca_comm
56 + if ca_comm is not None and not ca_comm.disabled and \
57 + ca_comm.should_custom_ar(input_):
58 + out = ca_comm.custom_all_reduce(input_)
59 + assert out is not None
60 + return out
61 + pynccl_comm = self.pynccl_comm
62 + assert pynccl_comm is not None
63 + out = pynccl_comm.all_reduce(input_)
64 + if out is None:
65 + # fall back to the default all-reduce using PyTorch.
66 + # this usually happens during testing.
67 + # when we run the model, allreduce only happens for the TP
68 + # group, where we always have either custom allreduce or pynccl.
69 + out = input_.clone()
70 + torch.distributed.all_reduce(out, group=self.device_group)
71 + return out
72 +
73 + def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
74 + """Sends a tensor to the destination rank in a non-blocking way"""
75 + """NOTE: `dst` is the local rank of the destination rank."""
76 + if dst is None:
77 + dst = (self.rank_in_group + 1) % self.world_size
78 +
79 + pynccl_comm = self.pynccl_comm
80 + if pynccl_comm is not None and not pynccl_comm.disabled:
81 + pynccl_comm.send(tensor, dst)
82 + else:
83 + torch.distributed.send(tensor, self.ranks[dst], self.device_group)
84 +
85 + def recv(self,
86 + size: torch.Size,
87 + dtype: torch.dtype,
88 + src: Optional[int] = None) -> torch.Tensor:
89 + """Receives a tensor from the source rank."""
90 + """NOTE: `src` is the local rank of the source rank."""
91 + if src is None:
92 + src = (self.rank_in_group - 1) % self.world_size
93 +
94 + tensor = torch.empty(size, dtype=dtype, device=self.device)
95 + pynccl_comm = self.pynccl_comm
96 + if pynccl_comm is not None and not pynccl_comm.disabled:
97 + pynccl_comm.recv(tensor, src)
98 + else:
99 + torch.distributed.recv(tensor, self.ranks[src], self.device_group)
100 + return tensor
101 +
102 + def destroy(self):
103 + if self.pynccl_comm is not None:
104 + self.pynccl_comm = None
105 + if self.ca_comm is not None:
106 + self.ca_comm = None
vllm/distributed/device_communicators/hpu_communicator.py
+14 -19
@@ -2,45 +2,40 @@
2
3 import torch
4 import torch.distributed as dist
5 -from torch.distributed import ProcessGroup
5
6 from vllm.platforms import current_platform
7
8 +from .base_device_communicator import DeviceCommunicatorBase
9 +
10 if current_platform.is_hpu():
11 import habana_frameworks.torch as htorch # noqa: F401
12
13
13 -class HpuCommunicator:
14 -
15 - def __init__(self, group: ProcessGroup):
16 - if not current_platform.is_hpu():
17 - self.disabled = True
18 - return
19 - self.disabled = False
20 - self.group = group
21 - self.world_size = dist.get_world_size(self.group)
14 +class HpuCommunicator(DeviceCommunicatorBase):
15
23 - def all_reduce(self, x: torch.Tensor) -> torch.Tensor:
16 + def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
17 # FIXME(kzawora): this is a workaround for a bug in Habana PT bridge
18 # occurring when PT_HPU_ENABLE_LAZY_COLLECTIVES=true env var is used
19 # (which is required for tensor parallel HPUGraph inference)
20 htorch.core.mark_step()
28 - dist.all_reduce(x, group=self.group)
29 - return x
21 + dist.all_reduce(input_, group=self.device_group)
22 + return input_
23
31 - def all_gather(self, x: torch.Tensor, dim: int = -1) -> torch.Tensor:
24 + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
25 world_size = self.world_size
26 if dim < 0:
27 # Convert negative dim to positive.
35 - dim += x.dim()
36 - input_size = x.size()
28 + dim += input_.dim()
29 + input_size = input_.size()
30 # Allocate output tensor.
31 output_tensor = torch.empty((world_size, ) + input_size,
39 - dtype=x.dtype,
40 - device=x.device)
32 + dtype=input_.dtype,
33 + device=input_.device)
34 # All-gather.
35 htorch.core.mark_step()
43 - dist.all_gather_into_tensor(output_tensor, x, group=self.group)
36 + dist.all_gather_into_tensor(output_tensor,
37 + input_,
38 + group=self.device_group)
39 # Reshape
40 output_tensor = output_tensor.movedim(0, dim)
41 output_tensor = output_tensor.reshape(input_size[:dim] +
vllm/distributed/device_communicators/tpu_communicator.py
+16 -13
@@ -1,13 +1,15 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 import os
4 +from typing import Optional
5
6 import torch
6 -import torch.distributed as dist
7 from torch.distributed import ProcessGroup
8
9 from vllm.platforms import current_platform
10
11 +from .base_device_communicator import DeviceCommunicatorBase
12 +
13 if current_platform.is_tpu():
14 import torch_xla.core.xla_model as xm
15 import torch_xla.runtime as xr
@@ -16,19 +18,20 @@ if current_platform.is_tpu():
18 from vllm.executor import ray_utils
19
20
19 -class TpuCommunicator:
21 +class TpuCommunicator(DeviceCommunicatorBase):
22
21 - def __init__(self, group: ProcessGroup):
22 - if not current_platform.is_tpu():
23 - self.disabled = True
24 - return
25 - self.disabled = False
23 + def __init__(self,
24 + cpu_group: ProcessGroup,
25 + device: Optional[torch.device] = None,
26 + device_group: Optional[ProcessGroup] = None,
27 + unique_name: str = ""):
28 + super().__init__(cpu_group, device, device_group, unique_name)
29
30 # NOTE(woosuk): When using TP > 1 on TPUs, every TPU on the same node
31 # must be used together. Therefore, the local rank and world size can
32 # be simply calculated as follows.
30 - global_rank = dist.get_rank(group)
31 - global_world_size = dist.get_world_size(group)
33 + global_rank = self.global_rank
34 + global_world_size = self.global_world_size
35
36 # Calculate how many TPU nodes are in the current deployment. This
37 # is the Ray placement group if it is deployed with Ray. Default
@@ -55,9 +58,9 @@ class TpuCommunicator:
58 pjrt.initialize_multiprocess(local_rank, local_world_size)
59 xr._init_world_size_ordinal()
60
58 - def all_reduce(self, x: torch.Tensor) -> torch.Tensor:
59 - return xm.all_reduce(xm.REDUCE_SUM, x)
61 + def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
62 + return xm.all_reduce(xm.REDUCE_SUM, input_)
63
61 - def all_gather(self, x: torch.Tensor, dim: int = -1) -> torch.Tensor:
64 + def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
65 assert dim == -1, "TPUs only support dim=-1 for all-gather."
63 - return xm.all_gather(x, dim=dim)
66 + return xm.all_gather(input_, dim=dim)
vllm/distributed/device_communicators/xpu_communicator.py deleted
-49
@@ -1,49 +0,0 @@
1 -# SPDX-License-Identifier: Apache-2.0
2 -
3 -import torch
4 -import torch.distributed as dist
5 -from torch.distributed import ProcessGroup
6 -
7 -from vllm.platforms import current_platform
8 -
9 -
10 -class XpuCommunicator:
11 -
12 - def __init__(self, group: ProcessGroup):
13 - if not current_platform.is_xpu():
14 - self.disabled = True
15 - return
16 - self.disabled = False
17 - self.group = group
18 - self.world_size = dist.get_world_size(self.group)
19 -
20 - def all_reduce(self, x: torch.Tensor) -> torch.Tensor:
21 - dist.all_reduce(x, group=self.group)
22 - return x
23 -
24 - def gather(self,
25 - input_: torch.Tensor,
26 - rank_in_group: int,
27 - dst: int = 0,
28 - dim: int = -1):
29 - # For xpu path, gather doesn't work properly together with ray
30 - # cluster so we use all_gather instead for now.
31 - input_size = input_.size()
32 - # Allocate output tensor.
33 - output_tensor = torch.empty((self.world_size, ) + input_size,
34 - dtype=input_.dtype,
35 - device=input_.device)
36 - # All-gather.
37 - torch.distributed.all_gather_into_tensor(output_tensor,
38 - input_,
39 - group=self.group)
40 - if rank_in_group == dst:
41 - # Reshape
42 - output_tensor = output_tensor.movedim(0, dim)
43 - output_tensor = output_tensor.reshape(input_size[:dim] +
44 - (self.world_size *
45 - input_size[dim], ) +
46 - input_size[dim + 1:])
47 - else:
48 - output_tensor = None
49 - return output_tensor
vllm/distributed/parallel_state.py
+48 -201
@@ -39,9 +39,12 @@ from torch.distributed import Backend, ProcessGroup
39
40 import vllm.distributed.kv_transfer.kv_transfer_agent as kv_transfer
41 import vllm.envs as envs
42 +from vllm.distributed.device_communicators.base_device_communicator import (
43 + DeviceCommunicatorBase)
44 from vllm.distributed.utils import StatelessProcessGroup
45 from vllm.logger import init_logger
44 -from vllm.utils import direct_register_custom_op, supports_custom_op
46 +from vllm.utils import (direct_register_custom_op, resolve_obj_by_qualname,
47 + supports_custom_op)
48
49 if TYPE_CHECKING:
50 from vllm.config import VllmConfig
@@ -130,9 +133,8 @@ class GroupCoordinator:
133 PyTorch ProcessGroup is bound to one specific communication backend,
134 e.g. NCCL, Gloo, MPI, etc.
135 GroupCoordinator takes charge of all the communication operations among
133 - the processes in the group. It can route the communication to
134 - a specific implementation (e.g. switch allreduce implementation
135 - based on the tensor size and cuda graph mode).
136 + the processes in the group. It manages both CPU and device
137 + communication.
138 """
139
140 # available attributes:
@@ -150,11 +152,8 @@ class GroupCoordinator:
152 rank_in_group: int # rank inside the group
153 cpu_group: ProcessGroup # group for CPU communication
154 device_group: ProcessGroup # group for device communication
153 - use_pynccl: bool # a hint of whether to use PyNccl
154 - use_custom_allreduce: bool # a hint of whether to use CustomAllreduce
155 - # communicators are only created for world size > 1
156 - pynccl_comm: Optional[Any] # PyNccl communicator
157 - ca_comm: Optional[Any] # Custom allreduce communicator
155 + use_device_communicator: bool # whether to use device communicator
156 + device_communicator: DeviceCommunicatorBase # device communicator
157 mq_broadcaster: Optional[Any] # shared memory broadcaster
158
159 def __init__(
@@ -162,11 +161,7 @@ class GroupCoordinator:
161 group_ranks: List[List[int]],
162 local_rank: int,
163 torch_distributed_backend: Union[str, Backend],
165 - use_pynccl: bool,
166 - use_custom_allreduce: bool,
167 - use_tpu_communicator: bool,
168 - use_hpu_communicator: bool,
169 - use_xpu_communicator: bool,
164 + use_device_communicator: bool,
165 use_message_queue_broadcaster: bool = False,
166 group_name: Optional[str] = None,
167 ):
@@ -196,56 +191,26 @@ class GroupCoordinator:
191 assert self.device_group is not None
192
193 from vllm.platforms import current_platform
194 +
195 + # TODO: fix it for other platforms
196 if current_platform.is_cuda_alike():
197 self.device = torch.device(f"cuda:{local_rank}")
198 else:
199 self.device = torch.device("cpu")
200
204 - self.use_pynccl = use_pynccl
205 - self.use_custom_allreduce = use_custom_allreduce
206 - self.use_tpu_communicator = use_tpu_communicator
207 - self.use_hpu_communicator = use_hpu_communicator
208 - self.use_xpu_communicator = use_xpu_communicator
209 -
210 - # lazy import to avoid documentation build error
211 - from vllm.distributed.device_communicators.custom_all_reduce import (
212 - CustomAllreduce)
213 - from vllm.distributed.device_communicators.pynccl import (
214 - PyNcclCommunicator)
215 -
216 - self.pynccl_comm: Optional[PyNcclCommunicator] = None
217 - if use_pynccl and self.world_size > 1:
218 - self.pynccl_comm = PyNcclCommunicator(
219 - group=self.cpu_group,
220 - device=self.device,
221 - )
201 + self.use_device_communicator = use_device_communicator
202
223 - self.ca_comm: Optional[CustomAllreduce] = None
224 - if use_custom_allreduce and self.world_size > 1:
225 - # Initialize a custom fast all-reduce implementation.
226 - self.ca_comm = CustomAllreduce(
227 - group=self.cpu_group,
203 + self.device_communicator: DeviceCommunicatorBase = None # type: ignore
204 + if use_device_communicator and self.world_size > 1:
205 + device_comm_cls = resolve_obj_by_qualname(
206 + current_platform.get_device_communicator_cls())
207 + self.device_communicator = device_comm_cls(
208 + cpu_group=self.cpu_group,
209 device=self.device,
210 + device_group=self.device_group,
211 + unique_name=self.unique_name,
212 )
213
231 - from vllm.distributed.device_communicators.tpu_communicator import (
232 - TpuCommunicator)
233 - self.tpu_communicator: Optional[TpuCommunicator] = None
234 - if use_tpu_communicator and self.world_size > 1:
235 - self.tpu_communicator = TpuCommunicator(group=self.cpu_group)
236 -
237 - from vllm.distributed.device_communicators.hpu_communicator import (
238 - HpuCommunicator)
239 - self.hpu_communicator: Optional[HpuCommunicator]
240 - if use_hpu_communicator and self.world_size > 1:
241 - self.hpu_communicator = HpuCommunicator(group=self.device_group)
242 -
243 - from vllm.distributed.device_communicators.xpu_communicator import (
244 - XpuCommunicator)
245 - self.xpu_communicator: Optional[XpuCommunicator]
246 - if use_xpu_communicator and self.world_size > 1:
247 - self.xpu_communicator = XpuCommunicator(group=self.device_group)
248 -
214 from vllm.distributed.device_communicators.shm_broadcast import (
215 MessageQueue)
216 self.mq_broadcaster: Optional[MessageQueue] = None
@@ -253,6 +218,9 @@ class GroupCoordinator:
218 self.mq_broadcaster = MessageQueue.create_from_process_group(
219 self.cpu_group, 1 << 22, 6)
220
221 + from vllm.platforms import current_platform
222 + self.use_custom_op_call = current_platform.is_cuda_alike()
223 +
224 @property
225 def first_rank(self):
226 """Return the global rank of the first process in the group"""
@@ -296,9 +264,16 @@ class GroupCoordinator:
264 else:
265 stream = graph_capture_context.stream
266
299 - ca_comm = self.ca_comm
300 - maybe_ca_context = nullcontext(
301 - ) if ca_comm is None else ca_comm.capture()
267 + # only cuda uses this function,
268 + # so we don't abstract it into the base class
269 + maybe_ca_context = nullcontext()
270 + from vllm.distributed.device_communicators.cuda_communicator import (
271 + CudaCommunicator)
272 + if self.device_communicator is not None:
273 + assert isinstance(self.device_communicator, CudaCommunicator)
274 + ca_comm = self.device_communicator.ca_comm
275 + if ca_comm is not None:
276 + maybe_ca_context = ca_comm.capture() # type: ignore
277
278 # ensure all initialization operations complete before attempting to
279 # capture the graph on another stream
@@ -328,54 +303,14 @@ class GroupCoordinator:
303 if self.world_size == 1:
304 return input_
305
331 - if input_.is_cpu:
332 - try:
333 - import intel_extension_for_pytorch as ipex
334 - ipex.distributed.all_reduce(input_, group=self.device_group)
335 - return input_
336 - except ImportError:
337 - """
338 - Intel IPEX not found. Falling back to PyTorch native
339 - all_reduce for CPU
340 - """
341 - torch.distributed.all_reduce(input_, group=self.device_group)
342 - return input_
343 -
344 - if self.tpu_communicator is not None and \
345 - not self.tpu_communicator.disabled:
346 - # TPU handles Dynamo with its own logic.
347 - return self.tpu_communicator.all_reduce(input_)
348 -
349 - if self.hpu_communicator is not None and \
350 - not self.hpu_communicator.disabled:
351 - return self.hpu_communicator.all_reduce(input_)
352 -
353 - if self.xpu_communicator is not None and \
354 - not self.xpu_communicator.disabled:
355 - return self.xpu_communicator.all_reduce(input_)
356 -
357 - return torch.ops.vllm.all_reduce(input_, group_name=self.unique_name)
306 + if self.use_custom_op_call:
307 + return torch.ops.vllm.all_reduce(input_,
308 + group_name=self.unique_name)
309 + else:
310 + return self._all_reduce_out_place(input_)
311
312 def _all_reduce_out_place(self, input_: torch.Tensor) -> torch.Tensor:
360 - # always try custom allreduce first,
361 - # and then pynccl.
362 - ca_comm = self.ca_comm
363 - if ca_comm is not None and not ca_comm.disabled and \
364 - ca_comm.should_custom_ar(input_):
365 - out = ca_comm.custom_all_reduce(input_)
366 - assert out is not None
367 - return out
368 - pynccl_comm = self.pynccl_comm
369 - assert pynccl_comm is not None
370 - out = pynccl_comm.all_reduce(input_)
371 - if out is None:
372 - # fall back to the default all-reduce using PyTorch.
373 - # this usually happens during testing.
374 - # when we run the model, allreduce only happens for the TP
375 - # group, where we always have either custom allreduce or pynccl.
376 - out = input_.clone()
377 - torch.distributed.all_reduce(out, group=self.device_group)
378 - return out
313 + return self.device_communicator.all_reduce(input_)
314
315 def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
316 world_size = self.world_size
@@ -385,40 +320,7 @@ class GroupCoordinator:
320 assert -input_.dim() <= dim < input_.dim(), (
321 f"Invalid dim ({dim}) for input tensor with shape {input_.size()}")
322
388 - # For TPUs, use TPU communicator.
389 - tpu_comm = self.tpu_communicator
390 - if tpu_comm is not None and not tpu_comm.disabled:
391 - return tpu_comm.all_gather(input_, dim)
392 -
393 - # For HPUs, use HPU communicator.
394 - hpu_comm = self.hpu_communicator
395 - if hpu_comm is not None and not hpu_comm.disabled:
396 - return hpu_comm.all_gather(input_, dim)
397 -
398 - if dim < 0:
399 - # Convert negative dim to positive.
400 - dim += input_.dim()
401 - input_size = input_.size()
402 - # NOTE: we have to use concat-style all-gather here,
403 - # stack-style all-gather has compatibility issues with
404 - # torch.compile . see https://github.com/pytorch/pytorch/issues/138795
405 - output_size = (input_size[0] * world_size, ) + input_size[1:]
406 - # Allocate output tensor.
407 - output_tensor = torch.empty(output_size,
408 - dtype=input_.dtype,
409 - device=input_.device)
410 - # All-gather.
411 - torch.distributed.all_gather_into_tensor(output_tensor,
412 - input_,
413 - group=self.device_group)
414 - # Reshape
415 - output_tensor = output_tensor.reshape((world_size, ) + input_size)
416 - output_tensor = output_tensor.movedim(0, dim)
417 - output_tensor = output_tensor.reshape(input_size[:dim] +
418 - (world_size *
419 - input_size[dim], ) +
420 - input_size[dim + 1:])
421 - return output_tensor
323 + return self.device_communicator.all_gather(input_, dim)
324
325 def gather(self,
326 input_: torch.Tensor,
@@ -433,30 +335,7 @@ class GroupCoordinator:
335 # Bypass the function if we are using only 1 GPU.
336 if world_size == 1:
337 return input_
436 - assert -input_.dim() <= dim < input_.dim(), (
437 - f"Invalid dim ({dim}) for input tensor with shape {input_.size()}")
438 - if dim < 0:
439 - # Convert negative dim to positive.
440 - dim += input_.dim()
441 - if self.xpu_communicator is not None and \
442 - not self.xpu_communicator.disabled:
443 - return self.xpu_communicator.gather(input_, self.rank_in_group,
444 - dst, dim)
445 - # Allocate output tensor.
446 - if self.rank_in_group == dst:
447 - gather_list = [torch.empty_like(input_) for _ in range(world_size)]
448 - else:
449 - gather_list = None
450 - # Gather.
451 - torch.distributed.gather(input_,
452 - gather_list,
453 - dst=self.ranks[dst],
454 - group=self.device_group)
455 - if self.rank_in_group == dst:
456 - output_tensor = torch.cat(gather_list, dim=dim)
457 - else:
458 - output_tensor = None
459 - return output_tensor
338 + return self.device_communicator.gather(input_, dst, dim)
339
340 def broadcast(self, input_: torch.Tensor, src: int = 0):
341 """Broadcast the input tensor.
@@ -798,14 +677,7 @@ class GroupCoordinator:
677 def send(self, tensor: torch.Tensor, dst: Optional[int] = None) -> None:
678 """Sends a tensor to the destination rank in a non-blocking way"""
679 """NOTE: `dst` is the local rank of the destination rank."""
801 - if dst is None:
802 - dst = (self.rank_in_group + 1) % self.world_size
803 -
804 - pynccl_comm = self.pynccl_comm
805 - if pynccl_comm is not None and not pynccl_comm.disabled:
806 - pynccl_comm.send(tensor, dst)
807 - else:
808 - torch.distributed.send(tensor, self.ranks[dst], self.device_group)
680 + self.device_communicator.send(tensor, dst)
681
682 def recv(self,
683 size: torch.Size,
@@ -813,16 +685,7 @@ class GroupCoordinator:
685 src: Optional[int] = None) -> torch.Tensor:
686 """Receives a tensor from the source rank."""
687 """NOTE: `src` is the local rank of the source rank."""
816 - if src is None:
817 - src = (self.rank_in_group - 1) % self.world_size
818 -
819 - tensor = torch.empty(size, dtype=dtype, device=self.device)
820 - pynccl_comm = self.pynccl_comm
821 - if pynccl_comm is not None and not pynccl_comm.disabled:
822 - pynccl_comm.recv(tensor, src)
823 - else:
824 - torch.distributed.recv(tensor, self.ranks[src], self.device_group)
825 - return tensor
688 + return self.device_communicator.recv(size, dtype, src)
689
690 def destroy(self):
691 if self.device_group is not None:
@@ -831,10 +694,8 @@ class GroupCoordinator:
694 if self.cpu_group is not None:
695 torch.distributed.destroy_process_group(self.cpu_group)
696 self.cpu_group = None
834 - if self.pynccl_comm is not None:
835 - self.pynccl_comm = None
836 - if self.ca_comm is not None:
837 - self.ca_comm = None
697 + if self.device_communicator is not None:
698 + self.device_communicator.destroy()
699 if self.mq_broadcaster is not None:
700 self.mq_broadcaster = None
701
@@ -853,11 +714,7 @@ def init_world_group(ranks: List[int], local_rank: int,
714 group_ranks=[ranks],
715 local_rank=local_rank,
716 torch_distributed_backend=backend,
856 - use_pynccl=False,
857 - use_custom_allreduce=False,
858 - use_tpu_communicator=False,
859 - use_hpu_communicator=False,
860 - use_xpu_communicator=False,
717 + use_device_communicator=False,
718 group_name="world",
719 )
720
@@ -866,23 +723,15 @@ def init_model_parallel_group(
723 group_ranks: List[List[int]],
724 local_rank: int,
725 backend: str,
869 - use_custom_allreduce: Optional[bool] = None,
726 use_message_queue_broadcaster: bool = False,
727 group_name: Optional[str] = None,
728 ) -> GroupCoordinator:
873 - if use_custom_allreduce is None:
874 - use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE
875 - from vllm.platforms import current_platform
729 +
730 return GroupCoordinator(
731 group_ranks=group_ranks,
732 local_rank=local_rank,
733 torch_distributed_backend=backend,
880 - use_pynccl=current_platform.is_cuda_alike(),
881 - use_custom_allreduce=current_platform.is_cuda_alike()
882 - and use_custom_allreduce,
883 - use_tpu_communicator=True,
884 - use_hpu_communicator=True,
885 - use_xpu_communicator=True,
734 + use_device_communicator=True,
735 use_message_queue_broadcaster=use_message_queue_broadcaster,
736 group_name=group_name,
737 )
@@ -1053,11 +902,9 @@ def initialize_model_parallel(
902 for i in range(num_pipeline_model_parallel_groups):
903 ranks = list(range(i, world_size, num_pipeline_model_parallel_groups))
904 group_ranks.append(ranks)
1056 - # pipeline parallel does not need custom allreduce
905 _PP = init_model_parallel_group(group_ranks,
906 get_world_group().local_rank,
907 backend,
1060 - use_custom_allreduce=False,
908 group_name="pp")
909
910
vllm/platforms/cpu.py
+7
@@ -146,3 +146,10 @@ class CpuPlatform(Platform):
146 @classmethod
147 def get_punica_wrapper(cls) -> str:
148 return "vllm.lora.punica_wrapper.punica_cpu.PunicaWrapperCPU"
149 +
150 + @classmethod
151 + def get_device_communicator_cls(cls) -> str:
152 + """
153 + Get device specific communicator class for distributed communication.
154 + """
155 + return "vllm.distributed.device_communicators.cpu_communicator.CpuCommunicator" # noqa
vllm/platforms/cuda.py
+4
@@ -233,6 +233,10 @@ class CudaPlatformBase(Platform):
233 def get_punica_wrapper(cls) -> str:
234 return "vllm.lora.punica_wrapper.punica_gpu.PunicaWrapperGPU"
235
236 + @classmethod
237 + def get_device_communicator_cls(cls) -> str:
238 + return "vllm.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa
239 +
240
241 # NVML utils
242 # Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
vllm/platforms/hpu.py
+4
@@ -88,3 +88,7 @@ class HpuPlatform(Platform):
88 @classmethod
89 def get_punica_wrapper(cls) -> str:
90 return "vllm.lora.punica_wrapper.punica_hpu.PunicaWrapperHPU"
91 +
92 + @classmethod
93 + def get_device_communicator_cls(cls) -> str:
94 + return "vllm.distributed.device_communicators.hpu_communicator.HpuCommunicator" # noqa
vllm/platforms/interface.py
+7
@@ -322,6 +322,13 @@ class Platform:
322 """
323 raise NotImplementedError
324
325 + @classmethod
326 + def get_device_communicator_cls(cls) -> str:
327 + """
328 + Get device specific communicator class for distributed communication.
329 + """
330 + return "vllm.distributed.device_communicator.base_device_communicator.DeviceCommunicatorBase" # noqa
331 +
332
333 class UnspecifiedPlatform(Platform):
334 _enum = PlatformEnum.UNSPECIFIED
vllm/platforms/rocm.py
+4
@@ -186,3 +186,7 @@ class RocmPlatform(Platform):
186 torch.cuda.reset_peak_memory_stats(device)
187 return torch.cuda.mem_get_info(device)[1] - torch.cuda.mem_get_info(
188 device)[0]
189 +
190 + @classmethod
191 + def get_device_communicator_cls(cls) -> str:
192 + return "vllm.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa
vllm/platforms/tpu.py
+4
@@ -115,3 +115,7 @@ class TpuPlatform(Platform):
115 def is_pin_memory_available(cls):
116 logger.warning("Pin memory is not supported on TPU.")
117 return False
118 +
119 + @classmethod
120 + def get_device_communicator_cls(cls) -> str:
121 + return "vllm.distributed.device_communicators.tpu_communicator.TpuCommunicator" # noqa