From c024cfdb02e32ef5c804d31b28846003894276aa Mon Sep 17 00:00:00 2001 From: siyu Date: Wed, 23 Sep 2026 02:19:40 +0800 Subject: [PATCH] [Fast Start] Support data parallelism in the weight cache daemon (#57386) Signed-off-by: liusy58 Signed-off-by: Isotr0py Co-authored-by: Isotr0py --- .../model_loader/test_weight_cache.py | 113 ++++++++++ .../model_loader/weight_cache/daemon.py | 204 ++++++++++++++---- .../model_loader/weight_cache/ipc_loader.py | 4 + .../model_loader/weight_cache/protocol.py | 6 + 4 files changed, 285 insertions(+), 42 deletions(-) diff --git a/tests/model_executor/model_loader/test_weight_cache.py b/tests/model_executor/model_loader/test_weight_cache.py index eb672bc35b..7b60aaf04b 100644 --- a/tests/model_executor/model_loader/test_weight_cache.py +++ b/tests/model_executor/model_loader/test_weight_cache.py @@ -251,3 +251,116 @@ def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase): assert cold_outputs == baseline_outputs assert warm_outputs == baseline_outputs assert restart_outputs == baseline_outputs + + +def _parallel(**kw): + from types import SimpleNamespace + + base = dict( + tensor_parallel_size=1, + pipeline_parallel_size=1, + data_parallel_size=1, + data_parallel_size_local=1, + data_parallel_rank=0, + nnodes=1, + node_rank=0, + ) + base.update(kw) + return SimpleNamespace(**base) + + +def test_daemon_places_tp_and_dp_ranks_on_local_gpus(): + """Each node serves a contiguous global-rank block (DP-major), so without + DP local GPU i is TP rank r*local+i, and single-node DP launchers offset + by --data-parallel-start-rank.""" + 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)] + + 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)] + + dp_tp = _parallel( + tensor_parallel_size=2, + data_parallel_size=4, + 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)] + + # 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)] + 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)] + + # 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)] + 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)] + + # 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)] + + +def test_daemon_rejects_unmappable_parallelism(): + from vllm.model_executor.model_loader.weight_cache.daemon import ( + _reject_unsupported_parallelism, + ) + + _reject_unsupported_parallelism( + _parallel( + data_parallel_size=16, data_parallel_size_local=4, data_parallel_rank=12 + ) + ) + # DP combines with --nnodes when the world size divides evenly. + _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)) + with pytest.raises(ValueError, match="evenly divide"): + _reject_unsupported_parallelism(_parallel(tensor_parallel_size=3, nnodes=2)) + with pytest.raises(ValueError, match="evenly divide"): + _reject_unsupported_parallelism( + _parallel(tensor_parallel_size=2, data_parallel_size=3, nnodes=4) + ) + with pytest.raises(ValueError, match="exceeds"): + _reject_unsupported_parallelism( + _parallel( + data_parallel_size=16, data_parallel_size_local=4, data_parallel_rank=13 + ) + ) + + +def test_weight_cache_key_distinguishes_dp_ranks(): + from dataclasses import replace + + from vllm.model_executor.model_loader.weight_cache.protocol import WeightCacheKey + + key = WeightCacheKey( + checkpoint="ckpt", + model_arch="Arch", + tp_size=1, + tp_rank=0, + dtype="bf16", + quantization=None, + quant_config_hash="h", + revision=None, + vllm_version="v", + dp_size=16, + dp_rank=3, + ) + 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"] diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index 1eeff81f65..70fb31d579 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/vllm/model_executor/model_loader/weight_cache/daemon.py @@ -17,8 +17,8 @@ Engines then load from the daemons with: vllm serve /path/to/model --tensor-parallel-size 4 \\ --load-format ipc_cache -Only tensor and expert parallelism are supported; pipeline and data parallelism -are rejected at launch. +Tensor, expert and data parallelism are supported; pipeline parallelism is +rejected at launch. 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 @@ -42,6 +42,32 @@ The global TP rank of local GPU ``i`` on node ``r`` is assignment, so each engine worker maps its shard from the daemon on its own node. +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 +``start_rank + i // tp_size`` and TP rank ``i % tp_size``, and all +``dp_size * tp_size`` daemons form one world group on +``--data-parallel-address``/``--weight-cache-master-port`` so the expert +shards are laid out exactly as in the engine: + + # node r (4 local GPUs) + python -m vllm.model_executor.model_loader.weight_cache.daemon \\ + --model /path/to/model --tensor-parallel-size 1 --enable-expert-parallel \\ + --data-parallel-size 16 --data-parallel-size-local 4 \\ + --data-parallel-start-rank 4r --data-parallel-address 10.0.0.1 \\ + --weight-cache-master-port 29600 + +Data parallelism also combines with multi-node tensor parallelism: pass both +flag sets (``--nnodes``/``--node-rank``/``--master-addr`` and the DP flags with +a reachable ``--data-parallel-address``). Each node then serves a contiguous +block of the ``dp_size * tp_size`` global ranks, e.g. TP8 x DP2 on 4 nodes: + + # node r (4 local GPUs) + python -m vllm.model_executor.model_loader.weight_cache.daemon \\ + --model /path/to/model --tensor-parallel-size 8 --enable-expert-parallel \\ + --nnodes 4 --node-rank r --master-addr 10.0.0.1 \\ + --data-parallel-size 2 --data-parallel-address 10.0.0.1 \\ + --weight-cache-master-port 29600 + With MTP, EAGLE or EAGLE3 speculative decoding the launcher additionally starts a draft daemon group that caches the draft model. It uses its own cache key, Unix sockets (``*_draft.sock``) and rendezvous port @@ -66,6 +92,7 @@ import torch from vllm.config import ( ParallelConfig, VllmConfig, + replace, set_current_vllm_config, ) from vllm.distributed import ( @@ -98,6 +125,7 @@ from vllm.platforms import current_platform from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.utils.network_utils import get_distributed_init_method, get_open_port from vllm.utils.torch_utils import set_default_torch_dtype +from vllm.v1.worker.workspace import init_workspace_manager logger = init_logger("vllm.model_executor.model_loader.weight_cache.daemon") @@ -158,7 +186,25 @@ class WeightCacheDaemon: distributed_init_method: str, socket_dir: str | None = None, is_draft: bool = False, + dp_rank: int = 0, ): + parallel_config = vllm_config.parallel_config + if parallel_config.data_parallel_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 + local_world = world_size // parallel_config.nnodes + start_dp_rank = parallel_config.node_rank * local_world // tp_size + else: + start_dp_rank = parallel_config.data_parallel_rank + parallel_config = replace( + parallel_config, + data_parallel_rank=dp_rank, + data_parallel_rank_local=dp_rank - start_dp_rank, + ) + vllm_config = replace(vllm_config, parallel_config=parallel_config) self.vllm_config = vllm_config # A draft daemon builds the draft with the target's VllmConfig, like # the engine does; only the ModelConfig differs. @@ -169,6 +215,11 @@ class WeightCacheDaemon: else: self.model_config = vllm_config.model_config self.tp_rank = tp_rank + 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.local_rank = local_rank self.distributed_init_method = distributed_init_method self.socket_dir = socket_dir @@ -179,28 +230,32 @@ class WeightCacheDaemon: # mutate hf_config.quantization_config. self.cache_config = WeightCacheKey.from_model_config( self.model_config, - tp_size=vllm_config.parallel_config.tensor_parallel_size, + tp_size=self.tp_size, tp_rank=tp_rank, + dp_size=self.dp_size, + dp_rank=dp_rank, is_draft=is_draft, ) def load_model(self) -> None: - tp_size = self.cache_config.tp_size torch.accelerator.set_device_index(self.local_rank) + # Outside the config context so the daemon keeps its own rendezvous + # instead of the engine's; model parallel groups (TP/DP/EP) are then + # carved out of this world group like in the engine. init_distributed_environment( - world_size=tp_size, - rank=self.tp_rank, + world_size=self.world_size, + rank=self.global_rank, distributed_init_method=self.distributed_init_method, local_rank=self.local_rank, backend=current_platform.dist_backend, ) with set_current_vllm_config(self.vllm_config): - ensure_model_parallel_initialized(tp_size, 1) + ensure_model_parallel_initialized(self.tp_size, 1) self.model = self.get_model() logger.info( "Weight cache %s daemon rank %d loaded model", self.role, - self.tp_rank, + self.global_rank, ) def get_model(self) -> torch.nn.Module: @@ -219,6 +274,9 @@ class WeightCacheDaemon: target_device = torch.device( device_config.device if load_config.device is None else load_config.device ) + # Attention backends may take workspace at construction, like in the + # engine's worker init. + init_workspace_manager(target_device) with set_default_torch_dtype(model_config.dtype): with target_device: model = loader.create_model(vllm_config, model_config) @@ -255,7 +313,7 @@ class WeightCacheDaemon: logger.info( "Weight cache %s daemon rank %d serving on %s", self.role, - self.tp_rank, + self.global_rank, socket_path, ) if ready_callback is not None: @@ -347,7 +405,7 @@ class WeightCacheDaemon: logger.info_once( "Weight cache %s daemon rank %d sent %d tensors (+%d aliases) to engine", self.role, - self.tp_rank, + self.global_rank, len(entries), len(aliases), ) @@ -358,7 +416,7 @@ class WeightCacheDaemon: logger.info( "Weight cache %s daemon rank %d released cached weights", self.role, - self.tp_rank, + self.global_rank, ) send_msg(conn, {"status": "ok"}) @@ -371,6 +429,7 @@ def _run_daemon( socket_dir: str | None, ready_queue: "multiprocessing.Queue[tuple[str, int]]", is_draft: bool = False, + dp_rank: int = 0, ) -> None: daemon = WeightCacheDaemon( vllm_config, @@ -379,9 +438,34 @@ def _run_daemon( distributed_init_method, socket_dir, is_draft, + dp_rank, ) daemon.load_model() - daemon.serve_forever(ready_callback=lambda: ready_queue.put((daemon.role, tp_rank))) + daemon.serve_forever( + ready_callback=lambda: ready_queue.put((daemon.role, daemon.global_rank)) + ) + + +def plan_local_ranks(parallel_config: ParallelConfig) -> list[tuple[int, int, int]]: + """``(local_rank, dp_rank, tp_rank)`` for every GPU this launcher serves. + + Global ranks enumerate DP then TP: ``global = dp_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 + 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) + ] def get_draft_daemon_config(vllm_config: VllmConfig) -> VllmConfig | None: @@ -393,17 +477,33 @@ def get_draft_daemon_config(vllm_config: VllmConfig) -> VllmConfig | None: def _reject_unsupported_parallelism(parallel_config: ParallelConfig) -> None: - """Reject parallelism modes other than tensor/expert parallelism.""" - unsupported = { - "pipeline parallelism": parallel_config.pipeline_parallel_size > 1, - "data parallelism": parallel_config.data_parallel_size > 1, - } - for name, enabled in unsupported.items(): - if enabled: + """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" + ) + dp_size = parallel_config.data_parallel_size + tp_size = parallel_config.tensor_parallel_size + if parallel_config.nnodes > 1: + world_size = dp_size * tp_size + if world_size % parallel_config.nnodes != 0: raise ValueError( - f"The weight cache daemon only supports tensor and expert " - f"parallelism; {name} is not supported" + 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})" ) + return + if dp_size == 1: + return + end_rank = ( + parallel_config.data_parallel_rank + parallel_config.data_parallel_size_local + ) + if end_rank > dp_size: + raise ValueError( + "--data-parallel-start-rank + --data-parallel-size-local " + f"({end_rank}) exceeds --data-parallel-size ({dp_size})" + ) def main() -> None: @@ -448,18 +548,38 @@ def main() -> None: parallel_config = vllm_config.parallel_config _reject_unsupported_parallelism(parallel_config) tp_size = parallel_config.tensor_parallel_size + dp_size = parallel_config.data_parallel_size - nnodes = parallel_config.nnodes - node_rank = parallel_config.node_rank - local_world_size = tp_size // nnodes + placements = plan_local_ranks(parallel_config) + local_world_size = len(placements) + if dp_size == 1: + nnodes = parallel_config.nnodes + node_rank = parallel_config.node_rank + master_addr = parallel_config.master_addr + else: + master_addr = parallel_config.data_parallel_master_ip + if parallel_config.nnodes > 1: + nnodes = parallel_config.nnodes + node_rank = parallel_config.node_rank + else: + nnodes = dp_size // parallel_config.data_parallel_size_local + node_rank = parallel_config.data_parallel_rank // ( + parallel_config.data_parallel_size_local + ) + if nnodes > 1 and master_addr in ("127.0.0.1", "localhost"): + raise ValueError( + "Data parallelism across nodes requires a reachable " + "--data-parallel-address for the daemon rendezvous" + ) - # The daemon forms its own TP group and holds it open while serving, so it - # needs a rendezvous port distinct from the engine's --master-port. All - # nodes must agree on it; single-node can auto-pick a free port. The master - # address is reused from the engine's --master-addr (node_rank 0). - master_addr = parallel_config.master_addr + # The daemon forms its own world group and holds it open while serving, so + # it needs a rendezvous port distinct from the engine's. All nodes must + # agree on it; single-node can auto-pick a free port. The master address is + # the engine's --master-addr (TP across nodes) or --data-parallel-address. if nnodes > 1 and args.weight_cache_master_port is None: - raise ValueError("--weight-cache-master-port is required when --nnodes > 1") + raise ValueError( + "--weight-cache-master-port is required when the daemons span nodes" + ) master_port = args.weight_cache_master_port or get_open_port() distributed_init_method = get_distributed_init_method(master_addr, master_port) @@ -487,32 +607,32 @@ def main() -> None: ctx = multiprocessing.get_context("spawn") ready_queue: multiprocessing.Queue[tuple[str, int]] = ctx.Queue() - # Global TP rank of local GPU i on this node; local index == device index. - global_ranks = [ - node_rank * local_world_size + local_rank - for local_rank in range(local_world_size) - ] + # Local index == device index; global rank enumerates DP then TP. expected_ready = { - (format_daemon_role(is_draft), global_rank) - for (is_draft, _, _), global_rank in product(groups, global_ranks) + (format_daemon_role(is_draft), dp_rank * tp_size + tp_rank) + for (is_draft, _, _), (_, dp_rank, tp_rank) in product(groups, placements) } procs = [ ctx.Process( target=_run_daemon, args=( - global_rank, + tp_rank, local_rank, config, init_method, args.weight_cache_socket_dir, ready_queue, is_draft, + dp_rank, ), - name=f"vllm-weight-cache-{format_daemon_role(is_draft)}-{global_rank}", - ) - for (is_draft, config, init_method), (local_rank, global_rank) in product( - groups, enumerate(global_ranks) + 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, + tp_rank, + ) in product(groups, placements) ] for proc in procs: proc.start() 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 e22fc2b029..8e37ce3ad6 100644 --- a/vllm/model_executor/model_loader/weight_cache/ipc_loader.py +++ b/vllm/model_executor/model_loader/weight_cache/ipc_loader.py @@ -13,6 +13,7 @@ import torch.nn as nn from vllm.config import ModelConfig, VllmConfig from vllm.config.load import LoadConfig from vllm.distributed import ( + get_dp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) @@ -285,10 +286,13 @@ class IpcModelLoader(BaseModelLoader): _register(alias_name, obj, isinstance(obj, nn.Parameter)) def _fetch_entries(self, model_config: ModelConfig) -> WeightCacheState: + dp_group = get_dp_group() cache_config = WeightCacheKey.from_model_config( model_config, tp_size=get_tensor_model_parallel_world_size(), tp_rank=get_tensor_model_parallel_rank(), + dp_size=dp_group.world_size, + dp_rank=dp_group.rank_in_group, is_draft=self.is_draft, ) return self._request_state(cache_config) diff --git a/vllm/model_executor/model_loader/weight_cache/protocol.py b/vllm/model_executor/model_loader/weight_cache/protocol.py index 3803030cd4..6b5bbc69e5 100644 --- a/vllm/model_executor/model_loader/weight_cache/protocol.py +++ b/vllm/model_executor/model_loader/weight_cache/protocol.py @@ -291,6 +291,8 @@ class WeightCacheKey: vllm_version: str is_draft: bool = False """Daemon group the weights come from; False is the target model.""" + dp_size: int = 1 + dp_rank: int = 0 @classmethod def from_model_config( @@ -300,6 +302,8 @@ class WeightCacheKey: tp_rank: int, *, is_draft: bool = False, + dp_size: int = 1, + dp_rank: int = 0, ) -> "WeightCacheKey": """Build the fingerprint for a model configuration. @@ -327,6 +331,8 @@ class WeightCacheKey: revision=model_config.revision, vllm_version=vllm.version.__version__, is_draft=is_draft, + dp_size=dp_size, + dp_rank=dp_rank, ) def mismatched_fields(self, other: "WeightCacheKey") -> list[str]: