[XPU][UT]Bugfix when the process can't see all the world_size meet accuracy issue. (#57779)
Signed-off-by: yisheng <yi.sheng@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
YiSheng5 committed
Sep 21, 2026 at 06:50 UTC
17e50b9b761023d5f2499f062507df1ff49092a4
1 file changed
+16
-4
vllm/distributed/device_communicators/xpu_communicator.py
+16
-4
@@ -99,12 +99,24 @@ class XpuCommunicator(DeviceCommunicatorBase):
99
output = torch.empty(
100
output_shape, dtype=input_tensor.dtype, device=input_tensor.device
101
)
102
- if sizes is not None and sizes.count(sizes[0]) != len(sizes):
103
- # if inputs shape in different ranks is not the same using reduce_scatter
102
+ if sizes is None or sizes.count(sizes[0]) == len(sizes):
103
+ dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group)
104
+ elif torch.xpu.device_count() < self.world_size:
105
+ # uneven reduce_scatter is not supported when the
106
+ # process cannot see every device in the group (
107
+ # e.g. external-LBDP).
108
+ max_size = max(sizes)
109
+ padded = input_tensor.new_zeros(
110
+ (world_size * max_size,) + input_tensor.shape[1:]
111
+ )
112
+ for rank, split in enumerate(input_tensor.split(sizes, dim=0)):
113
+ padded[rank * max_size : rank * max_size + sizes[rank]] = split
114
+ padded_output = input_tensor.new_empty((max_size,) + input_tensor.shape[1:])
115
+ dist.reduce_scatter_tensor(padded_output, padded, group=self.device_group)
116
+ output = padded_output[:chunk_size]
117
+ else:
118
input_splits = list(input_tensor.split(sizes, dim=0))
119
dist.reduce_scatter(output, input_splits, group=self.device_group)
106
- else:
107
- dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group)
120
# Reshape before returning
121
return output.movedim(0, dim).contiguous()
122