[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>
siyu committed
Sep 23, 2026 at 02:19 UTC
c024cfdb02e32ef5c804d31b28846003894276aa
4 files changed
+286
-43
tests/model_executor/model_loader/test_weight_cache.py
+113
@@ -251,3 +251,116 @@ def test_ipc_cache_cold_start_and_warm_restart(vllm_runner, case: ModelCase):
251
assert cold_outputs == baseline_outputs
252
assert warm_outputs == baseline_outputs
253
assert restart_outputs == baseline_outputs
254
+
255
+
256
+def _parallel(**kw):
257
+ from types import SimpleNamespace
258
+
259
+ base = dict(
260
+ tensor_parallel_size=1,
261
+ pipeline_parallel_size=1,
262
+ data_parallel_size=1,
263
+ data_parallel_size_local=1,
264
+ data_parallel_rank=0,
265
+ nnodes=1,
266
+ node_rank=0,
267
+ )
268
+ base.update(kw)
269
+ return SimpleNamespace(**base)
270
+
271
+
272
+def test_daemon_places_tp_and_dp_ranks_on_local_gpus():
273
+ """Each node serves a contiguous global-rank block (DP-major), so without
274
+ DP local GPU i is TP rank r*local+i, and single-node DP launchers offset
275
+ by --data-parallel-start-rank."""
276
+ from vllm.model_executor.model_loader.weight_cache.daemon import plan_local_ranks
277
+
278
+ tp_second_node = _parallel(tensor_parallel_size=8, nnodes=2, node_rank=1)
279
+ assert plan_local_ranks(tp_second_node) == [(i, 0, 4 + i) for i in range(4)]
280
+
281
+ dp_third_node = _parallel(
282
+ data_parallel_size=16, data_parallel_size_local=4, data_parallel_rank=8
283
+ )
284
+ assert plan_local_ranks(dp_third_node) == [(i, 8 + i, 0) for i in range(4)]
285
+
286
+ dp_tp = _parallel(
287
+ tensor_parallel_size=2,
288
+ data_parallel_size=4,
289
+ data_parallel_size_local=2,
290
+ data_parallel_rank=2,
291
+ )
292
+ assert plan_local_ranks(dp_tp) == [(0, 2, 0), (1, 2, 1), (2, 3, 0), (3, 3, 1)]
293
+
294
+ # DP + nnodes: node-local TP replicas
295
+ dp_tp_node0 = _parallel(tensor_parallel_size=8, data_parallel_size=2, nnodes=2)
296
+ assert plan_local_ranks(dp_tp_node0) == [(i, 0, i) for i in range(8)]
297
+ dp_tp_node1 = _parallel(
298
+ tensor_parallel_size=8, data_parallel_size=2, nnodes=2, node_rank=1
299
+ )
300
+ assert plan_local_ranks(dp_tp_node1) == [(i, 1, i) for i in range(8)]
301
+
302
+ # DP + nnodes: TP group spans nodes (TP8 x DP2 on 4 nodes)
303
+ tp_span_node1 = _parallel(
304
+ tensor_parallel_size=8, data_parallel_size=2, nnodes=4, node_rank=1
305
+ )
306
+ assert plan_local_ranks(tp_span_node1) == [(i, 0, 4 + i) for i in range(4)]
307
+ tp_span_node2 = _parallel(
308
+ tensor_parallel_size=8, data_parallel_size=2, nnodes=4, node_rank=2
309
+ )
310
+ assert plan_local_ranks(tp_span_node2) == [(i, 1, i) for i in range(4)]
311
+
312
+ # DP + nnodes: several DP replicas per node (TP4 x DP4 on 2 nodes)
313
+ dp_multi_node0 = _parallel(tensor_parallel_size=4, data_parallel_size=4, nnodes=2)
314
+ assert plan_local_ranks(dp_multi_node0) == [(i, i // 4, i % 4) for i in range(8)]
315
+
316
+
317
+def test_daemon_rejects_unmappable_parallelism():
318
+ from vllm.model_executor.model_loader.weight_cache.daemon import (
319
+ _reject_unsupported_parallelism,
320
+ )
321
+
322
+ _reject_unsupported_parallelism(
323
+ _parallel(
324
+ data_parallel_size=16, data_parallel_size_local=4, data_parallel_rank=12
325
+ )
326
+ )
327
+ # DP combines with --nnodes when the world size divides evenly.
328
+ _reject_unsupported_parallelism(
329
+ _parallel(tensor_parallel_size=4, data_parallel_size=4, nnodes=2)
330
+ )
331
+ with pytest.raises(ValueError, match="pipeline"):
332
+ _reject_unsupported_parallelism(_parallel(pipeline_parallel_size=2))
333
+ with pytest.raises(ValueError, match="evenly divide"):
334
+ _reject_unsupported_parallelism(_parallel(tensor_parallel_size=3, nnodes=2))
335
+ with pytest.raises(ValueError, match="evenly divide"):
336
+ _reject_unsupported_parallelism(
337
+ _parallel(tensor_parallel_size=2, data_parallel_size=3, nnodes=4)
338
+ )
339
+ with pytest.raises(ValueError, match="exceeds"):
340
+ _reject_unsupported_parallelism(
341
+ _parallel(
342
+ data_parallel_size=16, data_parallel_size_local=4, data_parallel_rank=13
343
+ )
344
+ )
345
+
346
+
347
+def test_weight_cache_key_distinguishes_dp_ranks():
348
+ from dataclasses import replace
349
+
350
+ from vllm.model_executor.model_loader.weight_cache.protocol import WeightCacheKey
351
+
352
+ key = WeightCacheKey(
353
+ checkpoint="ckpt",
354
+ model_arch="Arch",
355
+ tp_size=1,
356
+ tp_rank=0,
357
+ dtype="bf16",
358
+ quantization=None,
359
+ quant_config_hash="h",
360
+ revision=None,
361
+ vllm_version="v",
362
+ dp_size=16,
363
+ dp_rank=3,
364
+ )
365
+ assert key.mismatched_fields(replace(key, dp_rank=4)) == ["dp_rank"]
366
+ assert key.mismatched_fields(replace(key, dp_size=8, dp_rank=3)) == ["dp_size"]
vllm/model_executor/model_loader/weight_cache/daemon.py
+163
-43
@@ -17,8 +17,8 @@ Engines then load from the daemons with:
17
vllm serve /path/to/model --tensor-parallel-size 4 \\
18
--load-format ipc_cache
19
20
-Only tensor and expert parallelism are supported; pipeline and data parallelism
21
-are rejected at launch.
20
+Tensor, expert and data parallelism are supported; pipeline parallelism is
21
+rejected at launch.
22
23
For multi-node tensor parallelism, run one launcher per node with a shared
24
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
42
assignment, so each engine worker maps its shard from the daemon on its own
43
node.
44
45
+For data parallelism (e.g. a TP1 x DP16 x EP decode fleet) run one launcher per
46
+node with the engine's DP placement flags. Local GPU ``i`` serves DP rank
47
+``start_rank + i // tp_size`` and TP rank ``i % tp_size``, and all
48
+``dp_size * tp_size`` daemons form one world group on
49
+``--data-parallel-address``/``--weight-cache-master-port`` so the expert
50
+shards are laid out exactly as in the engine:
51
+
52
+ # node r (4 local GPUs)
53
+ python -m vllm.model_executor.model_loader.weight_cache.daemon \\
54
+ --model /path/to/model --tensor-parallel-size 1 --enable-expert-parallel \\
55
+ --data-parallel-size 16 --data-parallel-size-local 4 \\
56
+ --data-parallel-start-rank 4r --data-parallel-address 10.0.0.1 \\
57
+ --weight-cache-master-port 29600
58
+
59
+Data parallelism also combines with multi-node tensor parallelism: pass both
60
+flag sets (``--nnodes``/``--node-rank``/``--master-addr`` and the DP flags with
61
+a reachable ``--data-parallel-address``). Each node then serves a contiguous
62
+block of the ``dp_size * tp_size`` global ranks, e.g. TP8 x DP2 on 4 nodes:
63
+
64
+ # node r (4 local GPUs)
65
+ python -m vllm.model_executor.model_loader.weight_cache.daemon \\
66
+ --model /path/to/model --tensor-parallel-size 8 --enable-expert-parallel \\
67
+ --nnodes 4 --node-rank r --master-addr 10.0.0.1 \\
68
+ --data-parallel-size 2 --data-parallel-address 10.0.0.1 \\
69
+ --weight-cache-master-port 29600
70
+
71
With MTP, EAGLE or EAGLE3 speculative decoding the launcher additionally
72
starts a draft daemon group that caches the draft model. It uses its own cache
73
key, Unix sockets (``*_draft.sock``) and rendezvous port
@@ -66,6 +92,7 @@ import torch
92
from vllm.config import (
93
ParallelConfig,
94
VllmConfig,
95
+ replace,
96
set_current_vllm_config,
97
)
98
from vllm.distributed import (
@@ -98,6 +125,7 @@ from vllm.platforms import current_platform
125
from vllm.utils.argparse_utils import FlexibleArgumentParser
126
from vllm.utils.network_utils import get_distributed_init_method, get_open_port
127
from vllm.utils.torch_utils import set_default_torch_dtype
128
+from vllm.v1.worker.workspace import init_workspace_manager
129
130
logger = init_logger("vllm.model_executor.model_loader.weight_cache.daemon")
131
@@ -158,7 +186,25 @@ class WeightCacheDaemon:
186
distributed_init_method: str,
187
socket_dir: str | None = None,
188
is_draft: bool = False,
189
+ dp_rank: int = 0,
190
):
191
+ parallel_config = vllm_config.parallel_config
192
+ if parallel_config.data_parallel_size > 1:
193
+ # The engine sets these on each DP rank's workers; the MoE/EP
194
+ # layers read them to place their expert shards.
195
+ tp_size = parallel_config.tensor_parallel_size
196
+ if parallel_config.nnodes > 1:
197
+ world_size = parallel_config.data_parallel_size * tp_size
198
+ local_world = world_size // parallel_config.nnodes
199
+ start_dp_rank = parallel_config.node_rank * local_world // tp_size
200
+ else:
201
+ start_dp_rank = parallel_config.data_parallel_rank
202
+ parallel_config = replace(
203
+ parallel_config,
204
+ data_parallel_rank=dp_rank,
205
+ data_parallel_rank_local=dp_rank - start_dp_rank,
206
+ )
207
+ vllm_config = replace(vllm_config, parallel_config=parallel_config)
208
self.vllm_config = vllm_config
209
# A draft daemon builds the draft with the target's VllmConfig, like
210
# the engine does; only the ModelConfig differs.
@@ -169,6 +215,11 @@ class WeightCacheDaemon:
215
else:
216
self.model_config = vllm_config.model_config
217
self.tp_rank = tp_rank
218
+ self.dp_rank = dp_rank
219
+ self.tp_size = parallel_config.tensor_parallel_size
220
+ self.dp_size = parallel_config.data_parallel_size
221
+ self.global_rank = dp_rank * self.tp_size + tp_rank
222
+ self.world_size = self.dp_size * self.tp_size
223
self.local_rank = local_rank
224
self.distributed_init_method = distributed_init_method
225
self.socket_dir = socket_dir
@@ -179,28 +230,32 @@ class WeightCacheDaemon:
230
# mutate hf_config.quantization_config.
231
self.cache_config = WeightCacheKey.from_model_config(
232
self.model_config,
182
- tp_size=vllm_config.parallel_config.tensor_parallel_size,
233
+ tp_size=self.tp_size,
234
tp_rank=tp_rank,
235
+ dp_size=self.dp_size,
236
+ dp_rank=dp_rank,
237
is_draft=is_draft,
238
)
239
240
def load_model(self) -> None:
188
- tp_size = self.cache_config.tp_size
241
torch.accelerator.set_device_index(self.local_rank)
242
+ # Outside the config context so the daemon keeps its own rendezvous
243
+ # instead of the engine's; model parallel groups (TP/DP/EP) are then
244
+ # carved out of this world group like in the engine.
245
init_distributed_environment(
191
- world_size=tp_size,
192
- rank=self.tp_rank,
246
+ world_size=self.world_size,
247
+ rank=self.global_rank,
248
distributed_init_method=self.distributed_init_method,
249
local_rank=self.local_rank,
250
backend=current_platform.dist_backend,
251
)
252
with set_current_vllm_config(self.vllm_config):
198
- ensure_model_parallel_initialized(tp_size, 1)
253
+ ensure_model_parallel_initialized(self.tp_size, 1)
254
self.model = self.get_model()
255
logger.info(
256
"Weight cache %s daemon rank %d loaded model",
257
self.role,
203
- self.tp_rank,
258
+ self.global_rank,
259
)
260
261
def get_model(self) -> torch.nn.Module:
@@ -219,6 +274,9 @@ class WeightCacheDaemon:
274
target_device = torch.device(
275
device_config.device if load_config.device is None else load_config.device
276
)
277
+ # Attention backends may take workspace at construction, like in the
278
+ # engine's worker init.
279
+ init_workspace_manager(target_device)
280
with set_default_torch_dtype(model_config.dtype):
281
with target_device:
282
model = loader.create_model(vllm_config, model_config)
@@ -255,7 +313,7 @@ class WeightCacheDaemon:
313
logger.info(
314
"Weight cache %s daemon rank %d serving on %s",
315
self.role,
258
- self.tp_rank,
316
+ self.global_rank,
317
socket_path,
318
)
319
if ready_callback is not None:
@@ -347,7 +405,7 @@ class WeightCacheDaemon:
405
logger.info_once(
406
"Weight cache %s daemon rank %d sent %d tensors (+%d aliases) to engine",
407
self.role,
350
- self.tp_rank,
408
+ self.global_rank,
409
len(entries),
410
len(aliases),
411
)
@@ -358,7 +416,7 @@ class WeightCacheDaemon:
416
logger.info(
417
"Weight cache %s daemon rank %d released cached weights",
418
self.role,
361
- self.tp_rank,
419
+ self.global_rank,
420
)
421
send_msg(conn, {"status": "ok"})
422
@@ -371,6 +429,7 @@ def _run_daemon(
429
socket_dir: str | None,
430
ready_queue: "multiprocessing.Queue[tuple[str, int]]",
431
is_draft: bool = False,
432
+ dp_rank: int = 0,
433
) -> None:
434
daemon = WeightCacheDaemon(
435
vllm_config,
@@ -379,9 +438,34 @@ def _run_daemon(
438
distributed_init_method,
439
socket_dir,
440
is_draft,
441
+ dp_rank,
442
)
443
daemon.load_model()
384
- daemon.serve_forever(ready_callback=lambda: ready_queue.put((daemon.role, tp_rank)))
444
+ daemon.serve_forever(
445
+ ready_callback=lambda: ready_queue.put((daemon.role, daemon.global_rank))
446
+ )
447
+
448
+
449
+def plan_local_ranks(parallel_config: ParallelConfig) -> list[tuple[int, int, int]]:
450
+ """``(local_rank, dp_rank, tp_rank)`` for every GPU this launcher serves.
451
+
452
+ Global ranks enumerate DP then TP: ``global = dp_rank * tp_size + tp_rank``.
453
+ With ``--nnodes`` the engine hands each node a contiguous block of global
454
+ ranks, so node ``r`` serves ``node_rank * local + i``; without it a
455
+ launcher's block starts at ``--data-parallel-start-rank * tp_size``.
456
+ """
457
+ tp_size = parallel_config.tensor_parallel_size
458
+ world_size = parallel_config.data_parallel_size * tp_size
459
+ if parallel_config.nnodes > 1:
460
+ local_world_size = world_size // parallel_config.nnodes
461
+ base = parallel_config.node_rank * local_world_size
462
+ else:
463
+ local_world_size = parallel_config.data_parallel_size_local * tp_size
464
+ base = parallel_config.data_parallel_rank * tp_size
465
+ return [
466
+ (i, (base + i) // tp_size, (base + i) % tp_size)
467
+ for i in range(local_world_size)
468
+ ]
469
470
471
def get_draft_daemon_config(vllm_config: VllmConfig) -> VllmConfig | None:
@@ -393,17 +477,33 @@ def get_draft_daemon_config(vllm_config: VllmConfig) -> VllmConfig | None:
477
478
479
def _reject_unsupported_parallelism(parallel_config: ParallelConfig) -> None:
396
- """Reject parallelism modes other than tensor/expert parallelism."""
397
- unsupported = {
398
- "pipeline parallelism": parallel_config.pipeline_parallel_size > 1,
399
- "data parallelism": parallel_config.data_parallel_size > 1,
400
- }
401
- for name, enabled in unsupported.items():
402
- if enabled:
480
+ """Reject pipeline parallelism and placements the daemon cannot map."""
481
+ if parallel_config.pipeline_parallel_size > 1:
482
+ raise ValueError(
483
+ "The weight cache daemon only supports tensor, expert and data "
484
+ "parallelism; pipeline parallelism is not supported"
485
+ )
486
+ dp_size = parallel_config.data_parallel_size
487
+ tp_size = parallel_config.tensor_parallel_size
488
+ if parallel_config.nnodes > 1:
489
+ world_size = dp_size * tp_size
490
+ if world_size % parallel_config.nnodes != 0:
491
raise ValueError(
404
- f"The weight cache daemon only supports tensor and expert "
405
- f"parallelism; {name} is not supported"
492
+ f"--nnodes ({parallel_config.nnodes}) must evenly divide the "
493
+ f"daemon world size ({world_size} = data-parallel-size "
494
+ f"{dp_size} x tensor-parallel-size {tp_size})"
495
)
496
+ return
497
+ if dp_size == 1:
498
+ return
499
+ end_rank = (
500
+ parallel_config.data_parallel_rank + parallel_config.data_parallel_size_local
501
+ )
502
+ if end_rank > dp_size:
503
+ raise ValueError(
504
+ "--data-parallel-start-rank + --data-parallel-size-local "
505
+ f"({end_rank}) exceeds --data-parallel-size ({dp_size})"
506
+ )
507
508
509
def main() -> None:
@@ -448,18 +548,38 @@ def main() -> None:
548
parallel_config = vllm_config.parallel_config
549
_reject_unsupported_parallelism(parallel_config)
550
tp_size = parallel_config.tensor_parallel_size
551
+ dp_size = parallel_config.data_parallel_size
552
+
553
+ placements = plan_local_ranks(parallel_config)
554
+ local_world_size = len(placements)
555
+ if dp_size == 1:
556
+ nnodes = parallel_config.nnodes
557
+ node_rank = parallel_config.node_rank
558
+ master_addr = parallel_config.master_addr
559
+ else:
560
+ master_addr = parallel_config.data_parallel_master_ip
561
+ if parallel_config.nnodes > 1:
562
+ nnodes = parallel_config.nnodes
563
+ node_rank = parallel_config.node_rank
564
+ else:
565
+ nnodes = dp_size // parallel_config.data_parallel_size_local
566
+ node_rank = parallel_config.data_parallel_rank // (
567
+ parallel_config.data_parallel_size_local
568
+ )
569
+ if nnodes > 1 and master_addr in ("127.0.0.1", "localhost"):
570
+ raise ValueError(
571
+ "Data parallelism across nodes requires a reachable "
572
+ "--data-parallel-address for the daemon rendezvous"
573
+ )
574
452
- nnodes = parallel_config.nnodes
453
- node_rank = parallel_config.node_rank
454
- local_world_size = tp_size // nnodes
455
-
456
- # The daemon forms its own TP group and holds it open while serving, so it
457
- # needs a rendezvous port distinct from the engine's --master-port. All
458
- # nodes must agree on it; single-node can auto-pick a free port. The master
459
- # address is reused from the engine's --master-addr (node_rank 0).
460
- master_addr = parallel_config.master_addr
575
+ # The daemon forms its own world group and holds it open while serving, so
576
+ # it needs a rendezvous port distinct from the engine's. All nodes must
577
+ # agree on it; single-node can auto-pick a free port. The master address is
578
+ # the engine's --master-addr (TP across nodes) or --data-parallel-address.
579
if nnodes > 1 and args.weight_cache_master_port is None:
462
- raise ValueError("--weight-cache-master-port is required when --nnodes > 1")
580
+ raise ValueError(
581
+ "--weight-cache-master-port is required when the daemons span nodes"
582
+ )
583
master_port = args.weight_cache_master_port or get_open_port()
584
distributed_init_method = get_distributed_init_method(master_addr, master_port)
585
@@ -487,32 +607,32 @@ def main() -> None:
607
608
ctx = multiprocessing.get_context("spawn")
609
ready_queue: multiprocessing.Queue[tuple[str, int]] = ctx.Queue()
490
- # Global TP rank of local GPU i on this node; local index == device index.
491
- global_ranks = [
492
- node_rank * local_world_size + local_rank
493
- for local_rank in range(local_world_size)
494
- ]
610
+ # Local index == device index; global rank enumerates DP then TP.
611
expected_ready = {
496
- (format_daemon_role(is_draft), global_rank)
497
- for (is_draft, _, _), global_rank in product(groups, global_ranks)
612
+ (format_daemon_role(is_draft), dp_rank * tp_size + tp_rank)
613
+ for (is_draft, _, _), (_, dp_rank, tp_rank) in product(groups, placements)
614
}
615
procs = [
616
ctx.Process(
617
target=_run_daemon,
618
args=(
503
- global_rank,
619
+ tp_rank,
620
local_rank,
621
config,
622
init_method,
623
args.weight_cache_socket_dir,
624
ready_queue,
625
is_draft,
626
+ dp_rank,
627
),
511
- name=f"vllm-weight-cache-{format_daemon_role(is_draft)}-{global_rank}",
512
- )
513
- for (is_draft, config, init_method), (local_rank, global_rank) in product(
514
- groups, enumerate(global_ranks)
628
+ name=f"vllm-weight-cache-{format_daemon_role(is_draft)}-"
629
+ f"{dp_rank * tp_size + tp_rank}",
630
)
631
+ for (is_draft, config, init_method), (
632
+ local_rank,
633
+ dp_rank,
634
+ tp_rank,
635
+ ) in product(groups, placements)
636
]
637
for proc in procs:
638
proc.start()
vllm/model_executor/model_loader/weight_cache/ipc_loader.py
+4
@@ -13,6 +13,7 @@ import torch.nn as nn
13
from vllm.config import ModelConfig, VllmConfig
14
from vllm.config.load import LoadConfig
15
from vllm.distributed import (
16
+ get_dp_group,
17
get_tensor_model_parallel_rank,
18
get_tensor_model_parallel_world_size,
19
)
@@ -285,10 +286,13 @@ class IpcModelLoader(BaseModelLoader):
286
_register(alias_name, obj, isinstance(obj, nn.Parameter))
287
288
def _fetch_entries(self, model_config: ModelConfig) -> WeightCacheState:
289
+ dp_group = get_dp_group()
290
cache_config = WeightCacheKey.from_model_config(
291
model_config,
292
tp_size=get_tensor_model_parallel_world_size(),
293
tp_rank=get_tensor_model_parallel_rank(),
294
+ dp_size=dp_group.world_size,
295
+ dp_rank=dp_group.rank_in_group,
296
is_draft=self.is_draft,
297
)
298
return self._request_state(cache_config)
vllm/model_executor/model_loader/weight_cache/protocol.py
+6
@@ -291,6 +291,8 @@ class WeightCacheKey:
291
vllm_version: str
292
is_draft: bool = False
293
"""Daemon group the weights come from; False is the target model."""
294
+ dp_size: int = 1
295
+ dp_rank: int = 0
296
297
@classmethod
298
def from_model_config(
@@ -300,6 +302,8 @@ class WeightCacheKey:
302
tp_rank: int,
303
*,
304
is_draft: bool = False,
305
+ dp_size: int = 1,
306
+ dp_rank: int = 0,
307
) -> "WeightCacheKey":
308
"""Build the fingerprint for a model configuration.
309
@@ -327,6 +331,8 @@ class WeightCacheKey:
331
revision=model_config.revision,
332
vllm_version=vllm.version.__version__,
333
is_draft=is_draft,
334
+ dp_size=dp_size,
335
+ dp_rank=dp_rank,
336
)
337
338
def mismatched_fields(self, other: "WeightCacheKey") -> list[str]: