[Core] Support multi-node inference(eager and cuda graph)

tests/distributed/test_comm_ops.py

@@ -24,7 +24,7 @@ def all_reduce_test_worker(tensor_parallel_size: int, rank: int,
     del os.environ["CUDA_VISIBLE_DEVICES"]
     device = torch.device(f"cuda:{rank}")
     torch.cuda.set_device(device)
-    init_test_distributed_environment(1, tensor_parallel_size, rank,
+    init_test_distributed_environment(1, tensor_parallel_size, rank, rank,
                                       distributed_init_port)
     num_elements = 8
     all_tensors = [
@@ -46,7 +46,7 @@ def all_gather_test_worker(tensor_parallel_size: int, rank: int,
     del os.environ["CUDA_VISIBLE_DEVICES"]
     device = torch.device(f"cuda:{rank}")
     torch.cuda.set_device(device)
-    init_test_distributed_environment(1, tensor_parallel_size, rank,
+    init_test_distributed_environment(1, tensor_parallel_size, rank, rank,
                                       distributed_init_port)
     num_dimensions = 3
     tensor_size = list(range(2, num_dimensions + 2))
@@ -74,7 +74,7 @@ def broadcast_tensor_dict_test_worker(tensor_parallel_size: int, rank: int,
     del os.environ["CUDA_VISIBLE_DEVICES"]
     device = torch.device(f"cuda:{rank}")
     torch.cuda.set_device(device)
-    init_test_distributed_environment(1, tensor_parallel_size, rank,
+    init_test_distributed_environment(1, tensor_parallel_size, rank, rank,
                                       distributed_init_port)
     test_dict = {
         "a": torch.arange(8, dtype=torch.float32, device="cuda"),

tests/distributed/test_custom_all_reduce.py

@@ -23,7 +23,7 @@ def graph_allreduce(world_size, rank, distributed_init_port):
     del os.environ["CUDA_VISIBLE_DEVICES"]
     device = torch.device(f"cuda:{rank}")
     torch.cuda.set_device(device)
-    init_test_distributed_environment(1, world_size, rank,
+    init_test_distributed_environment(1, world_size, rank, rank,
                                       distributed_init_port)
 
     custom_ar.init_custom_ar()
@@ -58,7 +58,7 @@ def eager_allreduce(world_size, rank, distributed_init_port):
     del os.environ["CUDA_VISIBLE_DEVICES"]
     device = torch.device(f"cuda:{rank}")
     torch.cuda.set_device(device)
-    init_test_distributed_environment(1, world_size, rank,
+    init_test_distributed_environment(1, world_size, rank, rank,
                                       distributed_init_port)
 
     sz = 1024

vllm/executor/ray_gpu_executor.py

@@ -188,8 +188,6 @@ def _init_workers_ray(self, placement_group: "PlacementGroup",
             is_driver_worker=True,
         )
 
-        # FIXME(woosuk): We are not properly initializing pynccl when
-        # we have multiple nodes.
         self._run_workers("init_device")
         self._run_workers(
             "load_model",

vllm/model_executor/parallel_utils/pynccl.py

@@ -202,6 +202,7 @@ def __init__(
         init_method=None,
         timeout=datetime.timedelta(seconds=10),
         world_size: int = -1,
+        local_rank: int = -1,
         rank: int = -1,
         store=None,
         group_name: str = "",
@@ -219,25 +220,22 @@ def __init__(
                                     store=store,
                                     group_name=group_name,
                                     pg_options=pg_options)
-        self.world_size = dist.get_world_size()
-        self.rank = dist.get_rank()
-        torch.cuda.set_device(self.rank)
-        if self.rank == 0:
+        torch.cuda.set_device(local_rank)
+        if rank == 0:
             self.unique_id = ncclGetUniqueId()
         else:
             self.unique_id = NcclUniqueId()
-        tensor = torch.ByteTensor(list(self.unique_id.internal)).cuda(
-            self.rank)
+        tensor = torch.ByteTensor(list(
+            self.unique_id.internal)).cuda(local_rank)
         dist.broadcast(tensor, src=0)
         byte_list = tensor.cpu().tolist()
-        self.unique_id = NcclUniqueId()
         for i, byte in enumerate(byte_list):
             self.unique_id.internal[i] = byte
         self.comm = ctypes.c_void_p()
-        result = _c_ncclCommInitRank(ctypes.byref(self.comm), self.world_size,
-                                     self.unique_id, self.rank)
+        result = _c_ncclCommInitRank(ctypes.byref(self.comm), world_size,
+                                     self.unique_id, rank)
         assert result == 0
-        self.stream = torch.cuda.Stream(device=f"cuda:{self.rank}")
+        self.stream = torch.cuda.Stream(device=f"cuda:{local_rank}")
 
     def all_reduce(self,
                    tensor: torch.Tensor,

vllm/model_executor/parallel_utils/pynccl_utils.py

@@ -36,11 +36,13 @@ def set_pynccl_stream(stream: torch.cuda.Stream):
         pass
 
 
-def init_process_group(world_size: int, rank: int, init_method: str) -> None:
+def init_process_group(world_size: int, local_rank: int, rank: int,
+                       init_method: str) -> None:
     assert not is_initialized()
     global comm
     comm = NCCLCommunicator(init_method=init_method,
                             world_size=world_size,
+                            local_rank=local_rank,
                             rank=rank)
 
 

vllm/test_utils.py

@@ -8,6 +8,7 @@
 def init_test_distributed_environment(
     pipeline_parallel_size: int,
     tensor_parallel_size: int,
+    local_rank: int,
     rank: int,
     distributed_init_port: str,
 ) -> None:
@@ -16,7 +17,10 @@ def init_test_distributed_environment(
                                      worker_use_ray=True)
     distributed_init_method = f"tcp://localhost:{distributed_init_port}"
     init_distributed_environment(
-        parallel_config, rank, distributed_init_method=distributed_init_method)
+        parallel_config,
+        local_rank,
+        rank,
+        distributed_init_method=distributed_init_method)
 
 
 def multi_process_tensor_parallel(

vllm/worker/worker.py

@@ -97,8 +97,8 @@ def init_device(self) -> None:
             raise RuntimeError(
                 f"Not support device type: {self.device_config.device}")
         # Initialize the distributed environment.
-        init_distributed_environment(self.parallel_config, self.rank,
-                                     self.distributed_init_method)
+        init_distributed_environment(self.parallel_config, self.local_rank,
+                                     self.rank, self.distributed_init_method)
         # Set random seed.
         set_random_seed(self.model_config.seed)
 
@@ -249,6 +249,7 @@ def get_cache_block_size_bytes(self, block_size: int,
 
 def init_distributed_environment(
     parallel_config: ParallelConfig,
+    local_rank: int,
     rank: int,
     distributed_init_method: Optional[str] = None,
 ) -> None:
@@ -282,9 +283,9 @@ def init_distributed_environment(
     elif parallel_config.world_size > 1:
         # NOTE(woosuk): We don't initialize pynccl process group when world size
         # is 1.
-        # TODO(woosuk): Support multi-node connection.
         pynccl_utils.init_process_group(
             world_size=parallel_config.world_size,
+            local_rank=local_rank,
             rank=rank,
             init_method=distributed_init_method,
         )