diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index f8d4e9fb28cc..fa3a702fb0af 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -451,7 +451,14 @@ def reduce_scatterv( 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) diff --git a/vllm/distributed/device_communicators/pynccl.py b/vllm/distributed/device_communicators/pynccl.py index 9f305c718f9d..9415d215e1b3 100644 --- a/vllm/distributed/device_communicators/pynccl.py +++ b/vllm/distributed/device_communicators/pynccl.py @@ -319,6 +319,66 @@ def reduce_scatterv( 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