choose reduce scatter over reduce

IwakuraRein · IwakuraRein · commit 529031a4807d · 2025-11-26T10:49:17.000-08:00
Signed-off-by: Siyuan Fu &lt;siyuanf@nvidia.com&gt;
diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py
@@ -225,7 +225,7 @@ def reduce_scatterv(
             output_shape, dtype=input_tensor.dtype, device=input_tensor.device
         )
 
-        if sizes is not None:
+        if 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)

Original file line number	Diff line number	Diff line change
`@@ -225,7 +225,7 @@ def reduce_scatterv(`
`225`	`225`	`output_shape, dtype=input_tensor.dtype, device=input_tensor.device`
`226`	`226`	`)`
`227`	`227`
`228`		`- if sizes is not None:`
	`228`	`+ if sizes is not None and sizes.count(sizes[0]) != len(sizes):`
`229`	`229`	`pynccl_comm.reduce_scatterv(output, input_tensor, sizes=sizes)`
`230`	`230`	`else:`
`231`	`231`	`pynccl_comm.reduce_scatter(output, input_tensor)`