Separate pynccl functions

Signed-off-by: Trevor Morris <tmorris@nvidia.com>

Author: trevor-m

vllm/distributed/device_communicators/cuda_communicator.py

@@ -152,7 +152,10 @@ def reduce_scatterv(self,
                              dtype=input_tensor.dtype,
                              device=input_tensor.device)
 
-        pynccl_comm.reduce_scatter(output, input_, sizes=sizes)
+        if sizes is not None:
+            pynccl_comm.reduce_scatterv(output, input_, sizes=sizes)
+        else:
+            pynccl_comm.reduce_scatter(output, input_)
 
         # Reshape before returning
         return output.movedim(0, dim).contiguous()
@@ -222,7 +225,10 @@ def _all_gather_single(input_: torch.Tensor,
             output_tensor = torch.empty(output_size,
                                         dtype=input_.dtype,
                                         device=input_.device)
-            pynccl_comm.all_gather(output_tensor, input_, sizes=sizes)
+            if sizes is not None:
+                pynccl_comm.all_gatherv(output_tensor, input_, sizes=sizes)
+            else:
+                pynccl_comm.all_gather(output_tensor, input_)
             return output_tensor
 
         if isinstance(input_, torch.Tensor):

vllm/distributed/device_communicators/pynccl.py

@@ -135,8 +135,7 @@ def all_reduce(self,
     def all_gather(self,
                    output_tensor: torch.Tensor,
                    input_tensor: torch.Tensor,
-                   stream=None,
-                   sizes: Optional[list[int]] = None):
+                   stream=None):
         if self.disabled:
             return
         # nccl communicator created on a specific device
@@ -147,37 +146,51 @@ def all_gather(self,
             f"but the input tensor is on {input_tensor.device}")
         if stream is None:
             stream = current_stream()
-        if sizes is not None:
-            assert output_tensor.shape[0] == sum(sizes)
-            split_offset = 0
-            self.nccl.ncclGroupStart()
-            for root, split_size in enumerate(sizes):
-                dst_slice = output_tensor[split_offset:split_offset +
-                                          split_size]
-                self.nccl.ncclBroadcast(
-                    buffer_type(input_tensor.data_ptr()),
-                    buffer_type(dst_slice.data_ptr()),
-                    dst_slice.numel(),
-                    ncclDataTypeEnum.from_torch(input_tensor.dtype),
-                    root,
-                    self.comm,
-                    cudaStream_t(stream.cuda_stream),
-                )
-                split_offset += split_size
-            self.nccl.ncclGroupEnd()
-        else:
-            self.nccl.ncclAllGather(
+        self.nccl.ncclAllGather(
+            buffer_type(input_tensor.data_ptr()),
+            buffer_type(output_tensor.data_ptr()), input_tensor.numel(),
+            ncclDataTypeEnum.from_torch(input_tensor.dtype), self.comm,
+            cudaStream_t(stream.cuda_stream))
+
+    def all_gatherv(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        sizes: list[int],
+        stream=None,
+    ):
+        if self.disabled:
+            return
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}")
+        if stream is None:
+            stream = current_stream()
+        assert output_tensor.shape[0] == sum(sizes)
+        split_offset = 0
+        self.nccl.ncclGroupStart()
+        for root, split_size in enumerate(sizes):
+            dst_slice = output_tensor[split_offset:split_offset + split_size]
+            self.nccl.ncclBroadcast(
                 buffer_type(input_tensor.data_ptr()),
-                buffer_type(output_tensor.data_ptr()), input_tensor.numel(),
-                ncclDataTypeEnum.from_torch(input_tensor.dtype), self.comm,
-                cudaStream_t(stream.cuda_stream))
+                buffer_type(dst_slice.data_ptr()),
+                dst_slice.numel(),
+                ncclDataTypeEnum.from_torch(input_tensor.dtype),
+                root,
+                self.comm,
+                cudaStream_t(stream.cuda_stream),
+            )
+            split_offset += split_size
+        self.nccl.ncclGroupEnd()
 
     def reduce_scatter(self,
                        output_tensor: torch.Tensor,
                        input_tensor: torch.Tensor,
                        op: ReduceOp = ReduceOp.SUM,
-                       stream=None,
-                       sizes: Optional[list[int]] = None):
+                       stream=None):
         if self.disabled:
             return
         # nccl communicator created on a specific device
@@ -188,29 +201,44 @@ def reduce_scatter(self,
             f"but the input tensor is on {input_tensor.device}")
         if stream is None:
             stream = current_stream()
+        self.nccl.ncclReduceScatter(
+            buffer_type(input_tensor.data_ptr()),
+            buffer_type(output_tensor.data_ptr()), output_tensor.numel(),
+            ncclDataTypeEnum.from_torch(input_tensor.dtype),
+            ncclRedOpTypeEnum.from_torch(op), self.comm,
+            cudaStream_t(stream.cuda_stream))
 
-        if sizes is not None:
-            split_offset = 0
-            self.nccl.ncclGroupStart()
-            for root, split_size in enumerate(sizes):
-                chunk = input_tensor[split_offset:split_offset + split_size,
-                                     ...]
+    def reduce_scatterv(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        sizes: list[int],
+        op: ReduceOp = ReduceOp.SUM,
+        stream=None,
+    ):
+        if self.disabled:
+            return
+        # nccl communicator created on a specific device
+        # will only work on tensors on the same device
+        # otherwise it will cause "illegal memory access"
+        assert input_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {input_tensor.device}")
+        if stream is None:
+            stream = current_stream()
 
-                self.nccl.ncclReduce(
-                    buffer_type(chunk.data_ptr()),
-                    buffer_type(output_tensor.data_ptr()), chunk.numel(),
-                    ncclDataTypeEnum.from_torch(input_tensor.dtype),
-                    ncclRedOpTypeEnum.from_torch(op), root, self.comm,
-                    cudaStream_t(stream.cuda_stream))
-                split_offset += split_size
-            self.nccl.ncclGroupEnd()
-        else:
-            self.nccl.ncclReduceScatter(
-                buffer_type(input_tensor.data_ptr()),
-                buffer_type(output_tensor.data_ptr()), output_tensor.numel(),
+        split_offset = 0
+        self.nccl.ncclGroupStart()
+        for root, split_size in enumerate(sizes):
+            chunk = input_tensor[split_offset:split_offset + split_size, ...]
+            self.nccl.ncclReduce(
+                buffer_type(chunk.data_ptr()),
+                buffer_type(output_tensor.data_ptr()), chunk.numel(),
                 ncclDataTypeEnum.from_torch(input_tensor.dtype),
-                ncclRedOpTypeEnum.from_torch(op), self.comm,
+                ncclRedOpTypeEnum.from_torch(op), root, self.comm,
                 cudaStream_t(stream.cuda_stream))
+            split_offset += split_size
+        self.nccl.ncclGroupEnd()
 
     def send(self, tensor: torch.Tensor, dst: int, stream=None):
         if self.disabled: