[Bugfix] Deterministic MoE combine (reduce_scatterv) under VLLM_BATCH_INVARIANT (#45683)

Signed-off-by: shijuzhao <shijuzhao@tencent.com>
Co-authored-by: shijuzhao <shijuzhao@tencent.com>

vllm/distributed/device_communicators/cuda_communicator.py

@@ -451,7 +451,14 @@ class CudaCommunicator(DeviceCommunicatorBase):
             output = torch.empty(
                 output_shape, dtype=input_tensor.dtype, device=input_tensor.device
             )
-            if sizes is not None and sizes.count(sizes[0]) != len(sizes):
+            use_deterministic_rs = envs.VLLM_BATCH_INVARIANT and world_size > 2
+            if use_deterministic_rs:
+                # Reduce to a fixed root (0) for determinism
+                reduced = torch.empty_like(input_tensor)
+                sizes = sizes if sizes else [chunk_size] * world_size
+                pynccl_comm.reduce(reduced, input_tensor, root=0)
+                pynccl_comm.scatter(output, reduced, sizes, root=0)
+            elif sizes is not None and sizes.count(sizes[0]) != len(sizes):
                 pynccl_comm.reduce_scatterv(output, input_tensor, sizes=sizes)
             else:
                 pynccl_comm.reduce_scatter(output, input_tensor)

vllm/distributed/device_communicators/pynccl.py

@@ -319,6 +319,66 @@ class PyNcclCommunicator:
             split_offset += split_size
         self.nccl.ncclGroupEnd()
 
+    def reduce(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        root: int,
+        op: ReduceOp = ReduceOp.SUM,
+        stream=None,
+    ):
+        if self.disabled:
+            return
+        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(input_tensor.data_ptr()),
+            buffer_type(output_tensor.data_ptr()),
+            input_tensor.numel(),
+            ncclDataTypeEnum.from_torch(input_tensor.dtype),
+            ncclRedOpTypeEnum.from_torch(op),
+            root,
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+
+    def scatter(
+        self,
+        output_tensor: torch.Tensor,
+        input_tensor: torch.Tensor,
+        sizes: list[int],
+        root: int = 0,
+        stream=None,
+    ):
+        if self.disabled:
+            return
+        assert output_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the output tensor is on {output_tensor.device}"
+        )
+        if stream is None:
+            stream = current_stream()
+        self.nccl.ncclGroupStart()
+        if self.rank == root:
+            split_offset = 0
+            for dst, split_size in enumerate(sizes):
+                if split_size == 0:
+                    continue
+
+                chunk = input_tensor[split_offset : split_offset + split_size, ...]
+                if dst == root:
+                    output_tensor.copy_(chunk)
+                else:
+                    self.send(chunk, dst, stream)
+                split_offset += split_size
+        elif sizes[self.rank] > 0:
+            self.recv(output_tensor, root, stream)
+        self.nccl.ncclGroupEnd()
+
     def send(self, tensor: torch.Tensor, dst: int, stream=None):
         if self.disabled:
             return