EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Fix All-Gather under world size one (#7219)

Browse Source
This commit is contained in:
Cheng Wan
2025-06-20 14:57:34 -07:00
committed by GitHub
parent dea2b84bc3
commit 8609e637a9
1 changed files with 12 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
python/sglang/srt/distributed/parallel_state.py
View File

@@ -523,17 +523,25 @@ class GroupCoordinator:
self, self,
input_: torch.Tensor, input_: torch.Tensor,
dim: int = -1, dim: int = -1,
tensor_list: List[torch.Tensor] = None, output_tensor_list: Optional[List[torch.Tensor]] = None,
) -> torch.Tensor: ) -> torch.Tensor:
world_size = self.world_size world_size = self.world_size
# Bypass the function if we are using only 1 GPU. # Bypass the function if we are using only 1 GPU.
if world_size == 1: if world_size == 1:
if output_tensor_list is not None:
logger.warning(
"Performing in-place all-gather with a group size of 1. "
"This may be unnecessary; consider bypassing it for better efficiency."
)
output_tensor_list[0].copy_(input_)
return None
else:
return input_ return input_
if tensor_list is not None: if output_tensor_list is not None:
# TODO(ch-wan): support other backends # TODO(ch-wan): support other backends
return torch.distributed.all_gather( return torch.distributed.all_gather(
tensor_list, input_, group=self.device_group output_tensor_list, input_, group=self.device_group
) )
assert ( assert (