From af5b4857e1353c01fd6bf41bc3cb9f84dc82dd89 Mon Sep 17 00:00:00 2001 From: siyu Date: Tue, 29 Sep 2026 17:37:43 +0800 Subject: [PATCH] [Fast Start] Support PP (#55477) Signed-off-by: Isotr0py Co-authored-by: Isotr0py Co-authored-by: Xun Sun --- .../model_loader/test_weight_cache.py | 36 +++++--- vllm/entrypoints/cli/preload.py | 22 +++-- .../model_loader/weight_cache/daemon.py | 84 +++++++++++-------- .../model_loader/weight_cache/ipc_loader.py | 4 + .../model_loader/weight_cache/protocol.py | 6 ++ 5 files changed, 100 insertions(+), 52 deletions(-) diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index e4ac86e0eb..7706522f6e 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/tests/model_executor/model_loader/test_weight_cache.py @@ -250,12 +250,12 @@ def test_daemon_places_tp_and_dp_ranks_on_local_gpus(): from vllm.model_executor.model_loader.weight_cache.daemon import plan_local_ranks tp_second_node = _parallel(tensor_parallel_size=8, nnodes=2, node_rank=1) - assert plan_local_ranks(tp_second_node) == [(i, 0, 4 + i) for i in range(4)] + assert plan_local_ranks(tp_second_node) == [(i, 0, 0, 4 + i) for i in range(4)] dp_third_node = _parallel( data_parallel_size=16, data_parallel_size_local=4, data_parallel_rank=8 ) - assert plan_local_ranks(dp_third_node) == [(i, 8 + i, 0) for i in range(4)] + assert plan_local_ranks(dp_third_node) == [(i, 8 + i, 0, 0) for i in range(4)] dp_tp = _parallel( tensor_parallel_size=2, @@ -263,29 +263,43 @@ def test_daemon_places_tp_and_dp_ranks_on_local_gpus(): data_parallel_size_local=2, data_parallel_rank=2, ) - assert plan_local_ranks(dp_tp) == [(0, 2, 0), (1, 2, 1), (2, 3, 0), (3, 3, 1)] + assert plan_local_ranks(dp_tp) == [ + (0, 2, 0, 0), + (1, 2, 0, 1), + (2, 3, 0, 0), + (3, 3, 0, 1), + ] # DP + nnodes: node-local TP replicas dp_tp_node0 = _parallel(tensor_parallel_size=8, data_parallel_size=2, nnodes=2) - assert plan_local_ranks(dp_tp_node0) == [(i, 0, i) for i in range(8)] + assert plan_local_ranks(dp_tp_node0) == [(i, 0, 0, i) for i in range(8)] dp_tp_node1 = _parallel( tensor_parallel_size=8, data_parallel_size=2, nnodes=2, node_rank=1 ) - assert plan_local_ranks(dp_tp_node1) == [(i, 1, i) for i in range(8)] + assert plan_local_ranks(dp_tp_node1) == [(i, 1, 0, i) for i in range(8)] # DP + nnodes: TP group spans nodes (TP8 x DP2 on 4 nodes) tp_span_node1 = _parallel( tensor_parallel_size=8, data_parallel_size=2, nnodes=4, node_rank=1 ) - assert plan_local_ranks(tp_span_node1) == [(i, 0, 4 + i) for i in range(4)] + assert plan_local_ranks(tp_span_node1) == [(i, 0, 0, 4 + i) for i in range(4)] tp_span_node2 = _parallel( tensor_parallel_size=8, data_parallel_size=2, nnodes=4, node_rank=2 ) - assert plan_local_ranks(tp_span_node2) == [(i, 1, i) for i in range(4)] + assert plan_local_ranks(tp_span_node2) == [(i, 1, 0, i) for i in range(4)] # DP + nnodes: several DP replicas per node (TP4 x DP4 on 2 nodes) dp_multi_node0 = _parallel(tensor_parallel_size=4, data_parallel_size=4, nnodes=2) - assert plan_local_ranks(dp_multi_node0) == [(i, i // 4, i % 4) for i in range(8)] + assert plan_local_ranks(dp_multi_node0) == [(i, i // 4, 0, i % 4) for i in range(8)] + + # PP adds one daemon placement per stage while preserving TP rank order. + pp_tp = _parallel(tensor_parallel_size=2, pipeline_parallel_size=2) + assert plan_local_ranks(pp_tp) == [ + (0, 0, 0, 0), + (1, 0, 0, 1), + (2, 0, 1, 0), + (3, 0, 1, 1), + ] def test_daemon_rejects_unmappable_parallelism(): @@ -302,8 +316,7 @@ def test_daemon_rejects_unmappable_parallelism(): _reject_unsupported_parallelism( _parallel(tensor_parallel_size=4, data_parallel_size=4, nnodes=2) ) - with pytest.raises(ValueError, match="pipeline"): - _reject_unsupported_parallelism(_parallel(pipeline_parallel_size=2)) + _reject_unsupported_parallelism(_parallel(pipeline_parallel_size=2)) with pytest.raises(ValueError, match="evenly divide"): _reject_unsupported_parallelism(_parallel(tensor_parallel_size=3, nnodes=2)) with pytest.raises(ValueError, match="evenly divide"): @@ -335,6 +348,9 @@ def test_weight_cache_key_distinguishes_dp_ranks(): vllm_version="v", dp_size=16, dp_rank=3, + pp_size=2, + pp_rank=1, ) assert key.mismatched_fields(replace(key, dp_rank=4)) == ["dp_rank"] assert key.mismatched_fields(replace(key, dp_size=8, dp_rank=3)) == ["dp_size"] + assert key.mismatched_fields(replace(key, pp_rank=0)) == ["pp_rank"] diff --git a/vllm/entrypoints/cli/preload.py b/vllm/entrypoints/cli/preload.py index 50734a052b..9eab7421f1 100644 --- a/vllm/entrypoints/cli/preload.py +++ b/vllm/entrypoints/cli/preload.py @@ -170,6 +170,7 @@ class PreloadSubcommand(CLISubcommand): parallel_config = vllm_config.parallel_config _reject_unsupported_parallelism(parallel_config) tp_size = parallel_config.tensor_parallel_size + pp_size = parallel_config.pipeline_parallel_size dp_size = parallel_config.data_parallel_size placements = plan_local_ranks(parallel_config) @@ -230,16 +231,21 @@ class PreloadSubcommand(CLISubcommand): ctx = multiprocessing.get_context("spawn") ready_queue: multiprocessing.Queue[tuple[str, int]] = ctx.Queue() - # Local index == device index; global rank enumerates DP then TP. + # Local index == device index; global rank enumerates DP, PP, then TP. expected_ready = { - (format_daemon_role(is_draft), dp_rank * tp_size + tp_rank) - for (is_draft, _, _), (_, dp_rank, tp_rank) in product(groups, placements) + ( + format_daemon_role(is_draft), + dp_rank * pp_size * tp_size + pp_rank * tp_size + tp_rank, + ) + for (is_draft, _, _), (_, dp_rank, pp_rank, tp_rank) in product( + groups, placements + ) } procs = [ ctx.Process( target=_run_daemon, args=( - tp_rank, + dp_rank * pp_size * tp_size + pp_rank * tp_size + tp_rank, local_rank, config, init_method, @@ -247,13 +253,17 @@ class PreloadSubcommand(CLISubcommand): ready_queue, is_draft, dp_rank, + pp_rank, + ), + name=( + f"vllm-weight-cache-{format_daemon_role(is_draft)}-" + f"{(dp_rank * pp_size * tp_size + pp_rank * tp_size + tp_rank)}" ), - name=f"vllm-weight-cache-{format_daemon_role(is_draft)}-" - f"{dp_rank * tp_size + tp_rank}", ) for (is_draft, config, init_method), ( local_rank, dp_rank, + pp_rank, tp_rank, ) in product(groups, placements) ] diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index 561f693ae7..b56cbbea8f 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -17,8 +17,7 @@ Engines then load from the daemons with: vllm serve /path/to/model --tensor-parallel-size 4 \\ --load-format ipc_cache -Tensor, expert and data parallelism are supported; pipeline parallelism is -rejected at launch. +Tensor, expert, data and pipeline parallelism are supported. For multi-node tensor parallelism, run one launcher per node with a shared rendezvous so the global TP group forms across nodes (CUDA IPC handles are @@ -40,7 +39,8 @@ node-local, so each node serves only its local GPUs' shards). Reuse the same The global TP rank of local GPU ``i`` on node ``r`` is ``r * (tp_size // nnodes) + i``, matching vLLM's contiguous per-node rank assignment, so each engine worker maps its shard from the daemon on its own -node. +node. With pipeline parallelism each node additionally runs one daemon per PP +stage; the global rank is ``pp_rank * tp_size + tp_rank``. For data parallelism (e.g. a TP1 x DP16 x EP decode fleet) run one launcher per node with the engine's DP placement flags. Local GPU ``i`` serves DP rank @@ -168,27 +168,30 @@ def export_entries( class WeightCacheDaemon: - """Per-GPU process that loads one TP shard and serves CUDA IPC handles.""" + """Per-GPU process that loads one TP/PP shard and serves CUDA IPC handles.""" def __init__( self, vllm_config: VllmConfig, - tp_rank: int, + global_rank: int, local_rank: int, distributed_init_method: str, socket_dir: str | None = None, is_draft: bool = False, dp_rank: int = 0, + pp_rank: int = 0, ): parallel_config = vllm_config.parallel_config - if parallel_config.data_parallel_size > 1: + self.tp_size = parallel_config.tensor_parallel_size + self.pp_size = parallel_config.pipeline_parallel_size + self.dp_size = parallel_config.data_parallel_size + if self.dp_size > 1: # The engine sets these on each DP rank's workers; the MoE/EP # layers read them to place their expert shards. - tp_size = parallel_config.tensor_parallel_size if parallel_config.nnodes > 1: - world_size = parallel_config.data_parallel_size * tp_size + world_size = self.dp_size * self.tp_size local_world = world_size // parallel_config.nnodes - start_dp_rank = parallel_config.node_rank * local_world // tp_size + start_dp_rank = parallel_config.node_rank * local_world // self.tp_size else: start_dp_rank = parallel_config.data_parallel_rank parallel_config = replace( @@ -206,12 +209,12 @@ class WeightCacheDaemon: self.model_config = spec.draft_model_config else: self.model_config = vllm_config.model_config - self.tp_rank = tp_rank + self.global_rank = global_rank + rank_in_dp = global_rank - dp_rank * self.pp_size * self.tp_size + self.tp_rank = rank_in_dp % self.tp_size + self.pp_rank = rank_in_dp // self.tp_size self.dp_rank = dp_rank - self.tp_size = parallel_config.tensor_parallel_size - self.dp_size = parallel_config.data_parallel_size - self.global_rank = dp_rank * self.tp_size + tp_rank - self.world_size = self.dp_size * self.tp_size + self.world_size = self.dp_size * self.pp_size * self.tp_size self.local_rank = local_rank self.distributed_init_method = distributed_init_method self.socket_dir = socket_dir @@ -223,7 +226,9 @@ class WeightCacheDaemon: self.cache_config = WeightCacheKey.from_model_config( self.model_config, tp_size=self.tp_size, - tp_rank=tp_rank, + tp_rank=self.tp_rank, + pp_size=self.pp_size, + pp_rank=self.pp_rank, dp_size=self.dp_size, dp_rank=dp_rank, is_draft=is_draft, @@ -242,7 +247,7 @@ class WeightCacheDaemon: backend=current_platform.dist_backend, ) with set_current_vllm_config(self.vllm_config): - ensure_model_parallel_initialized(self.tp_size, 1) + ensure_model_parallel_initialized(self.tp_size, self.pp_size) self.model = self.get_model() logger.info( "Weight cache %s daemon rank %d loaded model", @@ -414,7 +419,7 @@ class WeightCacheDaemon: def _run_daemon( - tp_rank: int, + global_rank: int, local_rank: int, vllm_config: VllmConfig, distributed_init_method: str, @@ -422,15 +427,17 @@ def _run_daemon( ready_queue: "multiprocessing.Queue[tuple[str, int]]", is_draft: bool = False, dp_rank: int = 0, + pp_rank: int = 0, ) -> None: daemon = WeightCacheDaemon( vllm_config, - tp_rank, + global_rank, local_rank, distributed_init_method, socket_dir, is_draft, dp_rank, + pp_rank, ) daemon.load_model() daemon.serve_forever( @@ -438,26 +445,34 @@ def _run_daemon( ) -def plan_local_ranks(parallel_config: ParallelConfig) -> list[tuple[int, int, int]]: - """``(local_rank, dp_rank, tp_rank)`` for every GPU this launcher serves. +def plan_local_ranks( + parallel_config: ParallelConfig, +) -> list[tuple[int, int, int, int]]: + """``(local_rank, dp_rank, pp_rank, tp_rank)`` for every local daemon. - Global ranks enumerate DP then TP: ``global = dp_rank * tp_size + tp_rank``. + Global ranks enumerate DP, PP, then TP: ``global = dp_rank * pp_size * + tp_size + pp_rank * tp_size + tp_rank``. With ``--nnodes`` the engine hands each node a contiguous block of global ranks, so node ``r`` serves ``node_rank * local + i``; without it a launcher's block starts at ``--data-parallel-start-rank * tp_size``. """ tp_size = parallel_config.tensor_parallel_size - world_size = parallel_config.data_parallel_size * tp_size + pp_size = parallel_config.pipeline_parallel_size + world_size = parallel_config.data_parallel_size * pp_size * tp_size if parallel_config.nnodes > 1: local_world_size = world_size // parallel_config.nnodes base = parallel_config.node_rank * local_world_size else: - local_world_size = parallel_config.data_parallel_size_local * tp_size - base = parallel_config.data_parallel_rank * tp_size - return [ - (i, (base + i) // tp_size, (base + i) % tp_size) - for i in range(local_world_size) - ] + local_world_size = parallel_config.data_parallel_size_local * pp_size * tp_size + base = parallel_config.data_parallel_rank * pp_size * tp_size + placements = [] + for local_rank in range(local_world_size): + global_rank = base + local_rank + rank_in_dp = global_rank % (pp_size * tp_size) + dp_rank = global_rank // (pp_size * tp_size) + pp_rank, tp_rank = divmod(rank_in_dp, tp_size) + placements.append((local_rank, dp_rank, pp_rank, tp_rank)) + return placements def get_draft_daemon_config(vllm_config: VllmConfig) -> VllmConfig | None: @@ -469,21 +484,18 @@ def get_draft_daemon_config(vllm_config: VllmConfig) -> VllmConfig | None: def _reject_unsupported_parallelism(parallel_config: ParallelConfig) -> None: - """Reject pipeline parallelism and placements the daemon cannot map.""" - if parallel_config.pipeline_parallel_size > 1: - raise ValueError( - "The weight cache daemon only supports tensor, expert and data " - "parallelism; pipeline parallelism is not supported" - ) + """Reject placements the daemon cannot map.""" dp_size = parallel_config.data_parallel_size tp_size = parallel_config.tensor_parallel_size + pp_size = parallel_config.pipeline_parallel_size if parallel_config.nnodes > 1: - world_size = dp_size * tp_size + world_size = dp_size * pp_size * tp_size if world_size % parallel_config.nnodes != 0: raise ValueError( f"--nnodes ({parallel_config.nnodes}) must evenly divide the " f"daemon world size ({world_size} = data-parallel-size " - f"{dp_size} x tensor-parallel-size {tp_size})" + f"{dp_size} x pipeline-parallel-size {pp_size} x " + f"tensor-parallel-size {tp_size})" ) return if dp_size == 1: diff --git a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py index 508e95c83d..185c11b83a 100644 --- a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py +++ b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py @@ -15,6 +15,7 @@ from vllm.config import ModelConfig, VllmConfig from vllm.config.load import LoadConfig from vllm.distributed import ( get_dp_group, + get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) @@ -289,10 +290,13 @@ class IpcModelLoader(BaseModelLoader): def _fetch_entries(self, model_config: ModelConfig) -> WeightCacheState: dp_group = get_dp_group() + pp_group = get_pp_group() cache_config = WeightCacheKey.from_model_config( model_config, tp_size=get_tensor_model_parallel_world_size(), tp_rank=get_tensor_model_parallel_rank(), + pp_size=pp_group.world_size, + pp_rank=pp_group.rank_in_group, dp_size=dp_group.world_size, dp_rank=dp_group.rank_in_group, is_draft=self.is_draft, diff --git a/vllm/model_executor/model_loader/weight_cache/protocol.py b/vllm/model_executor/model_loader/weight_cache/protocol.py index 6b5bbc69e5..72053805cd 100644 --- a/vllm/model_executor/model_loader/weight_cache/protocol.py +++ b/vllm/model_executor/model_loader/weight_cache/protocol.py @@ -289,6 +289,8 @@ class WeightCacheKey: quant_config_hash: str revision: str | None vllm_version: str + pp_size: int = 1 + pp_rank: int = 0 is_draft: bool = False """Daemon group the weights come from; False is the target model.""" dp_size: int = 1 @@ -301,6 +303,8 @@ class WeightCacheKey: tp_size: int, tp_rank: int, *, + pp_size: int = 1, + pp_rank: int = 0, is_draft: bool = False, dp_size: int = 1, dp_rank: int = 0, @@ -325,6 +329,8 @@ class WeightCacheKey: model_arch=arch, tp_size=tp_size, tp_rank=tp_rank, + pp_size=pp_size, + pp_rank=pp_rank, dtype=str(model_config.dtype), quantization=model_config.quantization, quant_config_hash=_hash_quant_config(quant_config),