iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Fast Start] Support PP (#55477)

Browse Source 
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Xun Sun <UNIDY2002@outlook.com>
This commit is contained in:
siyu
2026-09-29 09:37:43 +00:00
committed by GitHub
co-authored by Isotr0py Xun Sun
parent 77e52645e9
commit af5b4857e1
5 changed files with 100 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/model_executor/model_loader/test_weight_cache.py
+26 -10
View File
@@ -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"]
vllm/entrypoints/cli/preload.py
+16 -6
View File
@@ -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)
]
vllm/model_executor/model_loader/weight_cache/daemon.py
+48 -36
View File
@@ -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:
vllm/model_executor/model_loader/weight_cache/ipc_loader.py
+4
View File
@@ -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,
vllm/model_executor/model_loader/weight_cache/protocol.py
+6
View File
@@ -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),