biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

fix pynccl reduce_scatter (#23648)

Browse Source 
Co-authored-by: hongchao <hongchao@msh.team>
This commit is contained in:
yzds
2025-08-27 09:21:11 +08:00
committed by GitHub
parent 6891205b16
commit c7c80af084
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
vllm/distributed/device_communicators/cuda_communicator.py
View File

@@ -152,7 +152,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
dtype=input_tensor.dtype, dtype=input_tensor.dtype,
device=input_tensor.device) device=input_tensor.device)
pynccl_comm.reduce_scatter(output, input_) pynccl_comm.reduce_scatter(output, input_tensor)
# Reshape before returning # Reshape before returning
return output.movedim(0, dim).contiguous() return output.movedim(0, dim).contiguous()
@@ -186,9 +186,9 @@ class CudaCommunicator(DeviceCommunicatorBase):
device=input_tensor.device) device=input_tensor.device)
if sizes is not None: if sizes is not None:
pynccl_comm.reduce_scatterv(output, input_, sizes=sizes) pynccl_comm.reduce_scatterv(output, input_tensor, sizes=sizes)
else: else:
pynccl_comm.reduce_scatter(output, input_) pynccl_comm.reduce_scatter(output, input_tensor)
# Reshape before returning # Reshape before returning
return output.movedim(0, dim).contiguous() return output.movedim(0, dim).contiguous()