Use NCCL instead of ray for control-plane communication to remove serialization overhead (#2221)

vllm/model_executor/parallel_utils/communication_op.py

@@ -1,6 +1,7 @@
 import torch
 
 from vllm.model_executor.parallel_utils.parallel_state import (
+    get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
     get_tensor_model_parallel_group,
 )
@@ -45,3 +46,61 @@ def tensor_model_parallel_all_gather(input_, dim=-1):
                                           (world_size * input_size[dim], ) +
                                           input_size[dim + 1:])
     return output_tensor
+
+
+def tensor_model_parallel_gather(input_, dst=0, dim=-1):
+    """Gather the input tensor across model parallel group.
+
+    NOTE: We assume that the input tensor is on the same device across
+    all the ranks.
+    """
+    world_size = get_tensor_model_parallel_world_size()
+    # Bypass the function if we are using only 1 GPU.
+    if world_size == 1:
+        return input_
+    assert -input_.dim() <= dim < input_.dim(), (
+        f"Invalid dim ({dim}) for input tensor with shape {input_.size()}")
+    if dim < 0:
+        # Convert negative dim to positive.
+        dim += input_.dim()
+    # Allocate output tensor.
+    if get_tensor_model_parallel_rank() == dst:
+        gather_list = [torch.empty_like(input_) for _ in range(world_size)]
+    else:
+        gather_list = None
+    # Gather.
+    torch.distributed.gather(input_,
+                             gather_list,
+                             dst=dst,
+                             group=get_tensor_model_parallel_group())
+    if get_tensor_model_parallel_rank() == dst:
+        output_tensor = torch.cat(gather_list, dim=dim)
+    else:
+        output_tensor = None
+    return output_tensor
+
+
+def broadcast(input_, src=0):
+    """Broadcast the input tensor."""
+    world_size = torch.distributed.get_world_size()
+    assert 0 <= src < world_size, f"Invalid src rank ({src})"
+
+    # Bypass the function if we are using only 1 GPU.
+    if world_size == 1:
+        return input_
+    # Broadcast.
+    torch.distributed.broadcast(input_, src=src)
+    return input_
+
+
+def broadcast_object_list(obj_list, src=0):
+    """Broadcast the input object list."""
+    world_size = torch.distributed.get_world_size()
+    assert 0 <= src < world_size, f"Invalid src rank ({src})"
+
+    # Bypass the function if we are using only 1 GPU.
+    if world_size == 1:
+        return obj_list
+    # Broadcast.
+    torch.distributed.broadcast_object_list(obj_list, src=src)
+    return obj_list