[Fast Start] Fast loader support nnode>1 (#55468)
Signed-off-by: siyu <mg21330037@smail.nju.edu.cn> Co-authored-by: Isotr0py <2037008807@qq.com>
siyu committed
Sep 14, 2026 at 11:05 UTC
52dd0d7562adb3c2d556ce6fe7ca7c3226e1976b
1 file changed
+72
-14
vllm/model_executor/model_loader/weight_cache/daemon.py
+72
-14
@@ -17,8 +17,30 @@ 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
21
-parallelism are rejected at launch.
20
+Only tensor and expert parallelism are supported; pipeline and data parallelism
21
+are 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
25
+node-local, so each node serves only its local GPUs' shards). Reuse the same
26
+``--nnodes``/``--node-rank``/``--master-addr`` flags you pass the engine, plus a
27
+``--weight-cache-master-port`` distinct from the engine's ``--master-port``:
28
+
29
+ # node 0 (8 local GPUs)
30
+ python -m vllm.model_executor.model_loader.weight_cache.daemon \\
31
+ --model /path/to/model --tensor-parallel-size 16 \\
32
+ --nnodes 2 --node-rank 0 --master-addr 10.0.0.1 \\
33
+ --weight-cache-master-port 29600
34
+ # node 1 (8 local GPUs)
35
+ python -m vllm.model_executor.model_loader.weight_cache.daemon \\
36
+ --model /path/to/model --tensor-parallel-size 16 \\
37
+ --nnodes 2 --node-rank 1 --master-addr 10.0.0.1 \\
38
+ --weight-cache-master-port 29600
39
+
40
+The global TP rank of local GPU ``i`` on node ``r`` is
41
+``r * (tp_size // nnodes) + i``, matching vLLM's contiguous per-node rank
42
+assignment, so each engine worker maps its shard from the daemon on its own
43
+node.
44
"""
45
46
import contextlib
@@ -132,11 +154,13 @@ class WeightCacheDaemon:
154
self,
155
vllm_config: VllmConfig,
156
tp_rank: int,
157
+ local_rank: int,
158
distributed_init_method: str,
159
socket_dir: str | None = None,
160
):
161
self.vllm_config = vllm_config
162
self.tp_rank = tp_rank
163
+ self.local_rank = local_rank
164
self.distributed_init_method = distributed_init_method
165
self.socket_dir = socket_dir
166
self.model: torch.nn.Module | None = None
@@ -152,12 +176,12 @@ class WeightCacheDaemon:
176
177
def load_model(self) -> None:
178
tp_size = self.cache_config.tp_size
155
- torch.accelerator.set_device_index(self.tp_rank)
179
+ torch.accelerator.set_device_index(self.local_rank)
180
init_distributed_environment(
181
world_size=tp_size,
182
rank=self.tp_rank,
183
distributed_init_method=self.distributed_init_method,
160
- local_rank=self.tp_rank,
184
+ local_rank=self.local_rank,
185
backend=current_platform.dist_backend,
186
)
187
with set_current_vllm_config(self.vllm_config):
@@ -307,13 +331,14 @@ class WeightCacheDaemon:
331
332
def _run_daemon(
333
tp_rank: int,
334
+ local_rank: int,
335
vllm_config: VllmConfig,
336
distributed_init_method: str,
337
socket_dir: str | None,
338
ready_queue: "multiprocessing.Queue[int]",
339
) -> None:
340
daemon = WeightCacheDaemon(
316
- vllm_config, tp_rank, distributed_init_method, socket_dir
341
+ vllm_config, tp_rank, local_rank, distributed_init_method, socket_dir
342
)
343
daemon.load_model()
344
daemon.serve_forever(ready_callback=lambda: ready_queue.put(tp_rank))
@@ -348,6 +373,15 @@ def main() -> None:
373
default=None,
374
help="Directory for the daemon Unix sockets (default: tempdir).",
375
)
376
+ parser.add_argument(
377
+ "--weight-cache-master-port",
378
+ type=int,
379
+ default=None,
380
+ help="Rendezvous port for the daemon's own TP group. Must differ from "
381
+ "the engine's --master-port (the daemon holds its group open while "
382
+ "serving) and match across nodes. Required when --nnodes > 1; defaults "
383
+ "to a free port for single-node.",
384
+ )
385
args = parser.parse_args()
386
engine_args = EngineArgs.from_cli_args(args)
387
vllm_config = engine_args.create_engine_config()
@@ -363,22 +397,41 @@ def main() -> None:
397
_reject_unsupported_parallelism(parallel_config)
398
tp_size = parallel_config.tensor_parallel_size
399
366
- distributed_init_method = get_distributed_init_method("127.0.0.1", get_open_port())
400
+ nnodes = parallel_config.nnodes
401
+ node_rank = parallel_config.node_rank
402
+ local_world_size = tp_size // nnodes
403
+
404
+ # The daemon forms its own TP group and holds it open while serving, so it
405
+ # needs a rendezvous port distinct from the engine's --master-port. All
406
+ # nodes must agree on it; single-node can auto-pick a free port. The master
407
+ # address is reused from the engine's --master-addr (node_rank 0).
408
+ master_addr = parallel_config.master_addr
409
+ if nnodes > 1 and args.weight_cache_master_port is None:
410
+ raise ValueError("--weight-cache-master-port is required when --nnodes > 1")
411
+ master_port = args.weight_cache_master_port or get_open_port()
412
+ distributed_init_method = get_distributed_init_method(master_addr, master_port)
413
+
414
ctx = multiprocessing.get_context("spawn")
415
ready_queue: multiprocessing.Queue[int] = ctx.Queue()
416
+ # Global TP rank of local GPU i on this node; local index == device index.
417
+ global_ranks = [
418
+ node_rank * local_world_size + local_rank
419
+ for local_rank in range(local_world_size)
420
+ ]
421
procs = [
422
ctx.Process(
423
target=_run_daemon,
424
args=(
373
- rank,
425
+ global_rank,
426
+ local_rank,
427
vllm_config,
428
distributed_init_method,
429
args.weight_cache_socket_dir,
430
ready_queue,
431
),
379
- name=f"vllm-weight-cache-daemon-{rank}",
432
+ name=f"vllm-weight-cache-daemon-{global_rank}",
433
)
381
- for rank in range(tp_size)
434
+ for local_rank, global_rank in enumerate(global_ranks)
435
]
436
for proc in procs:
437
proc.start()
@@ -391,7 +444,7 @@ def main() -> None:
444
signal.signal(signal.SIGTERM, _shutdown)
445
446
ready_ranks: set[int] = set()
394
- while len(ready_ranks) < tp_size:
447
+ while len(ready_ranks) < local_world_size:
448
try:
449
ready_ranks.add(ready_queue.get(timeout=1.0))
450
except queue.Empty:
@@ -407,11 +460,16 @@ def main() -> None:
460
for proc in procs:
461
proc.join()
462
sys.exit(max((p.exitcode or 0) for p in procs))
410
- logger.info(
411
- "Weight cache daemon ready: all %d ranks serving in %s",
412
- tp_size,
413
- args.weight_cache_socket_dir or "the default socket dir",
463
+ socket_dir_msg = args.weight_cache_socket_dir or "the default socket dir"
464
+ logger.info_once(
465
+ "===== Weight cache daemon READY: node %d/%d serving %d local rank(s) "
466
+ "in %s =====",
467
+ node_rank,
468
+ nnodes,
469
+ local_world_size,
470
+ socket_dir_msg,
471
)
472
+
473
for proc in procs:
474
proc.join()
475
sys.exit(max(proc.exitcode or 0 for proc in procs))