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

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

Browse Source
This commit is contained in:
Roy
2024-03-29 06:01:55 +08:00
committed by GitHub
parent a4075cba4d
commit 515386ef3c
7 changed files with 25 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
vllm/model_executor/parallel_utils/pynccl.py
View File

@@ -202,6 +202,7 @@ class NCCLCommunicator:
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 @@ class NCCLCommunicator:
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,

4
vllm/model_executor/parallel_utils/pynccl_utils.py
View File

@@ -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)