[Fast Start] Support data parallelism in the weight cache daemon (#57386)

Signed-off-by: liusy58 <mg21330037@smail.nju.edu.cn> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
2026-09-22 18:19:40 +00:00
parent ca831d1c55
commit c024cfdb02
4 changed files with 285 additions and 42 deletions
@@ -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"]
@@ -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()
@@ -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)
@@ -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]: