diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 95fd8787af..c0eb67cac1 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -292,6 +292,10 @@ class ParallelConfig: timeout parameter. If None, PyTorch's default timeout is used (600s for NCCL). Increase this for multi-node setups where model downloads may be slow.""" + cpu_distributed_timeout_seconds: int | None = None + """Timeout (in seconds) for cpu communication groups. If None, PyTorch's + default timeout is used (1800s for gloo).""" + world_size: int = Field(init=False) """world_size is TPxPP, it affects the number of workers we create.""" diff --git a/vllm/distributed/parallel_state.py b/vllm/distributed/parallel_state.py index 2a9452fc25..712167c601 100644 --- a/vllm/distributed/parallel_state.py +++ b/vllm/distributed/parallel_state.py @@ -335,6 +335,10 @@ class GroupCoordinator: self_device_group = None self_cpu_group = None + from vllm.distributed.utils import get_cpu_distributed_timeout_or_none + + timeout = get_cpu_distributed_timeout_or_none() + for ranks in group_ranks: device_group = torch.distributed.new_group( ranks, backend=torch_distributed_backend @@ -342,7 +346,9 @@ class GroupCoordinator: # a group with `gloo` backend, to allow direct coordination between # processes through the CPU. with suppress_stdout(): - cpu_group = torch.distributed.new_group(ranks, backend="gloo") + cpu_group = torch.distributed.new_group( + ranks, backend="gloo", timeout=timeout + ) if self.rank in ranks: self.ranks = ranks self.world_size = len(ranks) diff --git a/vllm/distributed/utils.py b/vllm/distributed/utils.py index 9991ab1ddc..ba0419a280 100644 --- a/vllm/distributed/utils.py +++ b/vllm/distributed/utils.py @@ -491,6 +491,16 @@ def get_cached_tcp_store_client(host: str, port: int) -> TCPStore: return TCPStore(host, port, is_master=False, wait_for_workers=False) +def get_cpu_distributed_timeout_or_none() -> timedelta | None: + from vllm.config import get_current_vllm_config_or_none + + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return None + timeout_seconds = vllm_config.parallel_config.cpu_distributed_timeout_seconds + return timedelta(seconds=timeout_seconds) if timeout_seconds is not None else None + + def init_gloo_process_group( prefix_store: PrefixStore, group_rank: int, @@ -570,6 +580,10 @@ def stateless_init_torch_distributed_process_group( init_method = get_tcp_uri(host, port) backend = Backend(backend) # it is basically string timeout = _get_default_timeout(backend) + if backend == "gloo": + gloo_timeout = get_cpu_distributed_timeout_or_none() + if gloo_timeout is not None: + timeout = gloo_timeout if listen_socket is not None: store = create_tcp_store( diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 0e25a4a35c..aa82b70eb9 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -456,6 +456,9 @@ class EngineArgs: nnodes: int = ParallelConfig.nnodes node_rank: int = ParallelConfig.node_rank distributed_timeout_seconds: int | None = ParallelConfig.distributed_timeout_seconds + cpu_distributed_timeout_seconds: int | None = ( + ParallelConfig.cpu_distributed_timeout_seconds + ) numa_bind: bool = ParallelConfig.numa_bind numa_bind_nodes: list[int] | None = ParallelConfig.numa_bind_nodes numa_bind_cpus: list[str] | None = ParallelConfig.numa_bind_cpus @@ -951,6 +954,10 @@ class EngineArgs: "--distributed-timeout-seconds", **parallel_kwargs["distributed_timeout_seconds"], ) + parallel_group.add_argument( + "--cpu-distributed-timeout-seconds", + **parallel_kwargs["cpu_distributed_timeout_seconds"], + ) parallel_group.add_argument("--numa-bind", **parallel_kwargs["numa_bind"]) parallel_group.add_argument( "--numa-bind-nodes", **parallel_kwargs["numa_bind_nodes"] @@ -1950,6 +1957,7 @@ class EngineArgs: nnodes=self.nnodes, node_rank=self.node_rank, distributed_timeout_seconds=self.distributed_timeout_seconds, + cpu_distributed_timeout_seconds=self.cpu_distributed_timeout_seconds, data_parallel_master_ip=data_parallel_address, data_parallel_rpc_port=data_parallel_rpc_port, data_parallel_backend=self.data_parallel_backend,