@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b316ac6589
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 13, 03:15 UTC

[V1] Support MP Executor for multi node distributed inference (#23691)

Signed-off-by: Lu Fang <fanglu@fb.com> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Signed-off-by: Lucia Fang <fanglu@fb.com> Signed-off-by: Lucia Fang <116399278+luccafong@users.noreply.github.com> Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: Nick Hill <nhill@redhat.com>

Lucia Fang committed Nov 16, 2025 at 01:01 UTC b316ac658985f542618316b4285bd213dfdde046
10 files changed +930 -82
tests/distributed/test_multiproc_executor.py new
+437
@@ -0,0 +1,437 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +"""
5 +Integration tests for MultiprocExecutor at the executor level.
6 +This test directly tests the executor without going through the LLM interface,
7 +focusing on executor initialization, RPC calls, and distributed execution.
8 +"""
9 +
10 +import multiprocessing
11 +import os
12 +
13 +from tests.utils import multi_gpu_test
14 +from vllm.config import VllmConfig
15 +from vllm.engine.arg_utils import EngineArgs
16 +from vllm.utils import get_open_port
17 +from vllm.v1.core.sched.output import SchedulerOutput
18 +from vllm.v1.executor.multiproc_executor import MultiprocExecutor
19 +
20 +MODEL = "facebook/opt-125m"
21 +
22 +
23 +def create_vllm_config(
24 + tensor_parallel_size: int = 1,
25 + pipeline_parallel_size: int = 1,
26 + max_model_len: int = 256,
27 + gpu_memory_utilization: float = 0.3,
28 + distributed_executor_backend: str = "mp",
29 + nnodes: int = 1,
30 + node_rank: int = 0,
31 + master_port: int = 0,
32 +) -> VllmConfig:
33 + """Create a VllmConfig for testing using EngineArgs."""
34 + engine_args = EngineArgs(
35 + model=MODEL,
36 + tensor_parallel_size=tensor_parallel_size,
37 + pipeline_parallel_size=pipeline_parallel_size,
38 + max_model_len=max_model_len,
39 + gpu_memory_utilization=gpu_memory_utilization,
40 + distributed_executor_backend=distributed_executor_backend,
41 + enforce_eager=True,
42 + )
43 + vllm_config = engine_args.create_engine_config()
44 +
45 + # Override distributed node settings if needed
46 + if nnodes > 1 or node_rank > 0:
47 + vllm_config.parallel_config.nnodes = nnodes
48 + vllm_config.parallel_config.node_rank = node_rank
49 + vllm_config.parallel_config.master_port = master_port
50 + if nnodes > 1:
51 + vllm_config.parallel_config.disable_custom_all_reduce = True
52 +
53 + return vllm_config
54 +
55 +
56 +def create_test_scheduler_output(num_requests: int = 1) -> SchedulerOutput:
57 + """Create a minimal SchedulerOutput for testing."""
58 + # This is a simplified version - in practice you'd need proper
59 + # SchedulerOutput construction based on the actual vLLM v1 API
60 + return SchedulerOutput(
61 + scheduled_new_reqs=[],
62 + scheduled_resumed_reqs=[],
63 + scheduled_running_reqs=[],
64 + num_scheduled_tokens={},
65 + total_num_scheduled_tokens=0,
66 + )
67 +
68 +
69 +def test_multiproc_executor_initialization():
70 + """Test that MultiprocExecutor can be initialized with proper config."""
71 + vllm_config = create_vllm_config(
72 + tensor_parallel_size=1,
73 + pipeline_parallel_size=1,
74 + )
75 +
76 + # Create executor - this should initialize workers
77 + executor = MultiprocExecutor(vllm_config=vllm_config)
78 +
79 + # Verify executor properties
80 + assert executor.world_size == 1, "World size should be 1 for single GPU"
81 + assert executor.local_world_size == 1, "Local world size should be 1"
82 + assert hasattr(executor, "workers"), "Executor should have workers"
83 + assert len(executor.workers) == 1, "Should have 1 worker for single GPU"
84 +
85 + # Clean up
86 + executor.shutdown()
87 +
88 +
89 +@multi_gpu_test(num_gpus=2)
90 +def test_multiproc_executor_initialization_tensor_parallel():
91 + """Test MultiprocExecutor initialization with tensor parallelism."""
92 + vllm_config = create_vllm_config(
93 + tensor_parallel_size=2,
94 + pipeline_parallel_size=1,
95 + )
96 +
97 + # Create executor
98 + executor = MultiprocExecutor(vllm_config=vllm_config)
99 +
100 + # Verify executor properties
101 + assert executor.world_size == 2, "World size should be 2 for TP=2"
102 + assert executor.local_world_size == 2, "Local world size should be 2"
103 + assert len(executor.workers) == 2, "Should have 2 workers for TP=2"
104 +
105 + # Verify output rank calculation
106 + output_rank = executor._get_output_rank()
107 + assert output_rank == 0, "Output rank should be 0 for TP=2, PP=1"
108 +
109 + # Clean up
110 + executor.shutdown()
111 +
112 +
113 +@multi_gpu_test(num_gpus=2)
114 +def test_multiproc_executor_collective_rpc():
115 + """Test collective RPC calls to all workers."""
116 + vllm_config = create_vllm_config(
117 + tensor_parallel_size=2,
118 + pipeline_parallel_size=1,
119 + )
120 +
121 + # Create executor
122 + executor = MultiprocExecutor(vllm_config=vllm_config)
123 +
124 + try:
125 + # Test check_health RPC - should work without errors
126 + executor.check_health()
127 +
128 + # Test that RPC works correctly
129 + # Note: We're just testing that the RPC mechanism works,
130 + # not testing actual model execution here
131 + assert not executor.is_failed, "Executor should not be in failed state"
132 +
133 + finally:
134 + # Clean up
135 + executor.shutdown()
136 +
137 +
138 +def test_multiproc_executor_failure_callback():
139 + """Test failure callback registration and invocation."""
140 + vllm_config = create_vllm_config(
141 + tensor_parallel_size=1,
142 + pipeline_parallel_size=1,
143 + )
144 +
145 + executor = MultiprocExecutor(vllm_config=vllm_config)
146 +
147 + try:
148 + # Test callback registration
149 + callback_invoked = []
150 +
151 + def test_callback():
152 + callback_invoked.append(True)
153 +
154 + # Register callback
155 + executor.register_failure_callback(test_callback)
156 +
157 + # Callback should not be invoked yet
158 + assert len(callback_invoked) == 0, "Callback should not be invoked immediately"
159 +
160 + # Simulate failure
161 + executor.is_failed = True
162 +
163 + # Register another callback - should be invoked immediately
164 + executor.register_failure_callback(test_callback)
165 + assert len(callback_invoked) == 1, (
166 + "Callback should be invoked when executor is failed"
167 + )
168 +
169 + finally:
170 + # Clean up
171 + executor.shutdown()
172 +
173 +
174 +@multi_gpu_test(num_gpus=2)
175 +def test_multiproc_executor_worker_monitor():
176 + """Test that worker monitor is set up correctly."""
177 + vllm_config = create_vllm_config(
178 + tensor_parallel_size=2,
179 + pipeline_parallel_size=1,
180 + )
181 +
182 + executor = MultiprocExecutor(vllm_config=vllm_config)
183 +
184 + try:
185 + # Verify all worker processes are alive
186 + for worker in executor.workers:
187 + assert worker.proc.is_alive(), f"Worker rank {worker.rank} should be alive"
188 +
189 + # Verify executor is not in failed state
190 + assert not executor.is_failed, "Executor should not be in failed state"
191 +
192 + finally:
193 + # Clean up
194 + executor.shutdown()
195 +
196 + # After shutdown, workers should be terminated
197 + import time
198 +
199 + time.sleep(0.5) # Give processes time to terminate
200 + for worker in executor.workers:
201 + assert not worker.proc.is_alive(), (
202 + f"Worker rank {worker.rank} should terminate after shutdown"
203 + )
204 +
205 +
206 +@multi_gpu_test(num_gpus=2)
207 +def test_multiproc_executor_get_response_message_queues():
208 + """Test message queue retrieval for different ranks."""
209 + vllm_config = create_vllm_config(
210 + tensor_parallel_size=2,
211 + pipeline_parallel_size=1,
212 + )
213 +
214 + executor = MultiprocExecutor(vllm_config=vllm_config)
215 +
216 + try:
217 + # Get all message queues
218 + all_queues = executor.get_response_mqs()
219 + assert len(all_queues) == 2, "Should have 2 message queues for 2 workers"
220 +
221 + # Get message queue for specific rank
222 + rank0_queue = executor.get_response_mqs(unique_reply_rank=0)
223 + assert len(rank0_queue) == 1, "Should have 1 message queue for rank 0"
224 +
225 + rank1_queue = executor.get_response_mqs(unique_reply_rank=1)
226 + assert len(rank1_queue) == 1, "Should have 1 message queue for rank 1"
227 +
228 + finally:
229 + # Clean up
230 + executor.shutdown()
231 +
232 +
233 +def test_multiproc_executor_shutdown_cleanup():
234 + """Test that shutdown properly cleans up resources."""
235 + vllm_config = create_vllm_config(
236 + tensor_parallel_size=1,
237 + pipeline_parallel_size=1,
238 + )
239 +
240 + executor = MultiprocExecutor(vllm_config=vllm_config)
241 +
242 + # Verify executor is set up
243 + assert hasattr(executor, "workers"), "Executor should have workers"
244 + assert len(executor.workers) > 0, "Should have at least one worker"
245 +
246 + # Shutdown
247 + executor.shutdown()
248 +
249 + # Verify cleanup
250 + import time
251 +
252 + time.sleep(0.5) # Give processes time to terminate
253 +
254 + for worker in executor.workers:
255 + assert not worker.proc.is_alive(), "Worker processes should be terminated"
256 +
257 + # Verify shutdown event is set
258 + assert executor.shutdown_event.is_set(), "Shutdown event should be set"
259 +
260 + # Multiple shutdowns should be safe (idempotent)
261 + executor.shutdown()
262 + executor.shutdown()
263 +
264 +
265 +@multi_gpu_test(num_gpus=4)
266 +def test_multiproc_executor_pipeline_parallel():
267 + """Test MultiprocExecutor with pipeline parallelism."""
268 + vllm_config = create_vllm_config(
269 + tensor_parallel_size=2,
270 + pipeline_parallel_size=2,
271 + )
272 +
273 + executor = MultiprocExecutor(vllm_config=vllm_config)
274 +
275 + try:
276 + # Verify executor properties
277 + assert executor.world_size == 4, "World size should be 4 for TP=2, PP=2"
278 + assert len(executor.workers) == 4, "Should have 4 workers"
279 +
280 + # Verify output rank calculation
281 + # For TP=2, PP=2: output should be from the last PP stage (ranks 2-3)
282 + # Specifically rank 2 (first rank of last PP stage)
283 + output_rank = executor._get_output_rank()
284 + assert output_rank == 2, "Output rank should be 2 (first rank of last PP stage)"
285 +
286 + # Verify max_concurrent_batches for pipeline parallel
287 + assert executor.max_concurrent_batches == 2, (
288 + "Max concurrent batches should equal PP size"
289 + )
290 +
291 + finally:
292 + # Clean up
293 + executor.shutdown()
294 +
295 +
296 +def test_multiproc_executor_properties():
297 + """Test various executor properties and configurations."""
298 + vllm_config = create_vllm_config(
299 + tensor_parallel_size=1,
300 + pipeline_parallel_size=1,
301 + )
302 +
303 + executor = MultiprocExecutor(vllm_config=vllm_config)
304 +
305 + try:
306 + # Test supports_pp property
307 + assert MultiprocExecutor.supports_pp is True, (
308 + "MultiprocExecutor should support pipeline parallelism"
309 + )
310 +
311 + # Test world_size calculation
312 + assert executor.world_size == (
313 + executor.parallel_config.tensor_parallel_size
314 + * executor.parallel_config.pipeline_parallel_size
315 + ), "World size should equal TP * PP"
316 +
317 + # Test local_world_size calculation
318 + assert executor.local_world_size == (
319 + executor.parallel_config.world_size // executor.parallel_config.nnodes
320 + ), "Local world size should be world_size / nnodes"
321 +
322 + finally:
323 + # Clean up
324 + executor.shutdown()
325 +
326 +
327 +@multi_gpu_test(num_gpus=4)
328 +def test_multiproc_executor_multi_node():
329 + """
330 + Test MultiprocExecutor with multi-node configuration.
331 + This simulates 2 nodes with TP=4:
332 + - Node 0 (rank 0): Uses GPUs 0,1 (CUDA_VISIBLE_DEVICES=0,1) with TP=2
333 + - Node 1 (rank 1): Uses GPUs 2,3 (CUDA_VISIBLE_DEVICES=2,3) with TP=2
334 + Total world_size = 4, nnodes = 2
335 + """
336 + port = get_open_port()
337 + # symm_mem does not work for simulating multi instance in single node
338 + os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0"
339 +
340 + def run_node(node_rank: int, result_queue: multiprocessing.Queue, port: int):
341 + """Run a single node's executor."""
342 + executor = None
343 + try:
344 + # Set CUDA_VISIBLE_DEVICES for this node
345 + if node_rank == 0:
346 + os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
347 + else:
348 + os.environ["CUDA_VISIBLE_DEVICES"] = "2,3"
349 +
350 + # Create config for this node
351 + vllm_config = create_vllm_config(
352 + tensor_parallel_size=4, # Total TP across all nodes
353 + pipeline_parallel_size=1,
354 + nnodes=2, # 2 nodes
355 + node_rank=node_rank,
356 + master_port=port, # same port
357 + )
358 +
359 + # Create executor for this node
360 + executor = MultiprocExecutor(vllm_config=vllm_config)
361 +
362 + # Verify node-specific properties
363 + assert executor.world_size == 4, (
364 + f"World size should be 4 on node {node_rank}"
365 + )
366 + assert executor.local_world_size == 2, (
367 + f"Local world size should be 2 on node {node_rank}"
368 + )
369 + assert len(executor.workers) == 2, (
370 + f"Should have 2 local workers on node {node_rank}"
371 + )
372 +
373 + # Verify worker ranks are correct for this node
374 + expected_ranks = [node_rank * 2, node_rank * 2 + 1]
375 + actual_ranks = sorted([w.rank for w in executor.workers])
376 + assert actual_ranks == expected_ranks, (
377 + f"Node {node_rank} should have workers "
378 + f"with ranks {expected_ranks}, got {actual_ranks}"
379 + )
380 + # Verify all workers are alive
381 + for worker in executor.workers:
382 + assert worker.proc.is_alive(), (
383 + f"Worker rank {worker.rank} should be alive on node {node_rank}"
384 + )
385 + # executor.gen
386 + # Put success result in queue BEFORE shutdown to avoid hanging
387 + result_queue.put({"node": node_rank, "success": True})
388 + import time
389 +
390 + time.sleep(2)
391 + executor.shutdown()
392 + except Exception as e:
393 + # Put failure result in queue
394 + result_queue.put({"node": node_rank, "success": False, "error": str(e)})
395 + raise e
396 + finally:
397 + if executor is not None:
398 + executor.shutdown()
399 +
400 + # Create a queue to collect results from both processes
401 + result_queue: multiprocessing.Queue[dict[str, int | bool]] = multiprocessing.Queue()
402 +
403 + # Start both node processes
404 + processes = []
405 + for node_rank in range(2):
406 + p = multiprocessing.Process(
407 + target=run_node,
408 + args=(node_rank, result_queue, port),
409 + name=f"Node{node_rank}",
410 + )
411 + p.start()
412 + processes.append(p)
413 +
414 + # Wait for both processes to complete
415 + all_completed = True
416 + for p in processes:
417 + p.join(timeout=60)
418 + if p.is_alive():
419 + p.terminate()
420 + p.join(timeout=20)
421 + if p.is_alive():
422 + p.kill()
423 + p.join()
424 + all_completed = False
425 +
426 + # Check results from both nodes
427 + results: list[dict[str, int | bool]] = []
428 + while len(results) < 2:
429 + try:
430 + result = result_queue.get(timeout=1)
431 + results.append(result)
432 + except Exception:
433 + pass
434 + assert all_completed, "Not all processes completed successfully"
435 + assert len(results) == 2, f"Expected 2 results, got {len(results)}"
436 + assert results[0]["success"], f"Node 0 failed: {results[0]}"
437 + assert results[1]["success"], f"Node 1 failed: {results[1]}"
vllm/config/parallel.py
+40
@@ -210,6 +210,18 @@ class ParallelConfig:
210 class is dynamically inherited by the worker class. This is used to inject
211 new attributes and methods to the worker class for use in collective_rpc
212 calls."""
213 + master_addr: str = "127.0.0.1"
214 + """distributed master address for multi-node distributed
215 + inference when distributed_executor_backend is mp."""
216 + master_port: int = 29501
217 + """distributed master port for multi-node distributed
218 + inference when distributed_executor_backend is mp."""
219 + node_rank: int = 0
220 + """distributed node rank for multi-node distributed
221 + inference when distributed_executor_backend is mp."""
222 + nnodes: int = 1
223 + """num of nodes for multi-node distributed
224 + inference when distributed_executor_backend is mp."""
225
226 world_size: int = Field(init=False)
227 """world_size is TPxPP, it affects the number of workers we create."""
@@ -387,6 +399,23 @@ class ParallelConfig:
399 and self.data_parallel_size > 1
400 )
401
402 + @property
403 + def node_rank_within_dp(self) -> int:
404 + return self.node_rank % self.nnodes_within_dp
405 +
406 + @property
407 + def nnodes_within_dp(self) -> int:
408 + if self.nnodes == 1:
409 + return 1
410 + data_parallel_node_size = (
411 + self.data_parallel_size // self.data_parallel_size_local
412 + )
413 + return self.nnodes // data_parallel_node_size
414 +
415 + @property
416 + def local_world_size(self) -> int:
417 + return self.world_size // self.nnodes_within_dp
418 +
419 @staticmethod
420 def has_unfinished_dp(dp_group: ProcessGroup, has_unfinished: bool) -> bool:
421 tensor = torch.tensor([has_unfinished], dtype=torch.int32, device="cpu")
@@ -528,6 +557,8 @@ class ParallelConfig:
557 ray_found = ray_utils.ray_is_available()
558 if current_platform.is_tpu() and envs.VLLM_XLA_USE_SPMD:
559 backend = "uni"
560 + elif current_platform.is_cuda() and self.nnodes > 1:
561 + backend = "mp"
562 elif (
563 current_platform.is_cuda()
564 and cuda_device_count_stateless() < self.world_size
@@ -565,6 +596,10 @@ class ParallelConfig:
596 "max_parallel_loading_workers is currently "
597 "not supported and will be ignored."
598 )
599 + if self.distributed_executor_backend != "mp" and self.nnodes > 1:
600 + raise ValueError(
601 + "nnodes > 1 can only be set when distributed exectuor backend is mp."
602 + )
603
604 @property
605 def use_ray(self) -> bool:
@@ -607,6 +642,11 @@ class ParallelConfig:
642 "Disabled the custom all-reduce kernel because it is not "
643 "supported on current platform."
644 )
645 + if self.nnodes > 1:
646 + self.disable_custom_all_reduce = True
647 + logger.debug(
648 + "Disabled the custom all-reduce since we are running on multi-node."
649 + )
650 if self.ray_workers_use_nsight and not self.use_ray:
651 raise ValueError(
652 "Unable to use nsight profiling unless workers run with Ray."
vllm/distributed/device_communicators/shm_broadcast.py
+94 -16
@@ -8,7 +8,7 @@ from dataclasses import dataclass, field
8 from multiprocessing import shared_memory
9 from pickle import PickleBuffer
10 from threading import Event
11 -from typing import TYPE_CHECKING, Any
11 +from typing import TYPE_CHECKING, Any, cast
12 from unittest.mock import patch
13
14 import torch
@@ -602,13 +602,87 @@ class MessageQueue:
602 return obj
603 return self.dequeue()
604
605 + @staticmethod
606 + def create_from_process_group_single_reader(
607 + pg: ProcessGroup,
608 + max_chunk_bytes,
609 + max_chunks,
610 + reader_rank: int = 0,
611 + blocking: bool = False,
612 + ) -> tuple["MessageQueue", list[Handle]]:
613 + """
614 + Creates a MessageQueue for a process group with a single reader.
615 +
616 + This method is designed for scenarios where only one process (the reader)
617 + will consume messages, and all other processes are writers. It sets up
618 + the shared memory buffer and communication handles accordingly, and
619 + gathers the handles from all processes to the reader.
620 +
621 + Args:
622 + pg (ProcessGroup): The torch distributed process group.
623 + max_chunk_bytes (int): Maximum size in bytes for each chunk in the buffer.
624 + max_chunks (int): Maximum number of chunks in the buffer.
625 + reader_rank (int, optional): The global rank that will act as the reader.
626 + Defaults to 0.
627 + blocking (bool, optional): If True, blocks until all processes are ready.
628 + Defaults to False.
629 +
630 + Returns:
631 + tuple[MessageQueue, list[Handle]]:
632 + The MessageQueue instance for the calling process,
633 + and a list of handles (only non-empty for the reader process).
634 + """
635 + local_size = torch.cuda.device_count()
636 + rank = dist.get_rank()
637 + same_node = rank // local_size == reader_rank // local_size
638 + buffer_io = MessageQueue(
639 + n_reader=1,
640 + n_local_reader=1 if same_node else 0,
641 + max_chunk_bytes=max_chunk_bytes,
642 + max_chunks=max_chunks,
643 + )
644 + handle = buffer_io.export_handle()
645 + handles = [None] * dist.get_world_size(pg) if rank == reader_rank else None
646 + dist.gather_object(handle, handles, dst=reader_rank, group=pg)
647 + if blocking:
648 + buffer_io.wait_until_ready()
649 + return buffer_io, cast(list[Handle], handles or [])
650 +
651 @staticmethod
652 def create_from_process_group(
653 pg: ProcessGroup | StatelessProcessGroup,
654 max_chunk_bytes,
655 max_chunks,
610 - writer_rank=0,
656 + writer_rank: int = 0,
657 + external_writer_handle=None,
658 + blocking: bool = True,
659 ) -> "MessageQueue":
660 + """
661 + Creates a MessageQueue for a distributed process group with one writer and
662 + multiple readers.
663 +
664 + This method is designed for scenarios where one process (the writer) sends
665 + messages, and all other processes (the readers) receive messages. It sets up
666 + the shared memory buffer and socket communication handles accordingly, and
667 + broadcasts the handle from the writer to all readers.
668 +
669 + Args:
670 + pg (ProcessGroup | StatelessProcessGroup): The torch distributed process
671 + group.
672 + max_chunk_bytes (int): Maximum size in bytes for each chunk in the buffer.
673 + max_chunks (int): Maximum number of chunks in the buffer.
674 + writer_rank (int, optional): The global rank that will act as the writer.
675 + Defaults to 0.
676 + external_writer_handle (Handle, optional): Used when there is a handle
677 + from an external Message Queue. If provided, use this handle to init
678 + PG writer message queue instead of creating a new one. Defaults to None.
679 + blocking (bool, optional): If True, blocks until all processes are ready.
680 + Defaults to True.
681 +
682 + Returns:
683 + MessageQueue: The MessageQueue instance for the calling process.
684 +
685 + """
686 if isinstance(pg, ProcessGroup):
687 group_rank = dist.get_rank(pg)
688 group_world_size = dist.get_world_size(pg)
@@ -617,23 +691,26 @@ class MessageQueue:
691 group_rank = pg.rank
692 group_world_size = pg.world_size
693 global_ranks = list(range(pg.world_size))
620 -
694 from vllm.distributed.parallel_state import in_the_same_node_as
695
696 status = in_the_same_node_as(pg, source_rank=writer_rank)
624 - same_node_ranks = [i for i, s in enumerate(status) if s]
625 - n_reader = group_world_size - 1
626 - n_local_reader = len(same_node_ranks) - 1
627 - local_reader_ranks = [i for i in same_node_ranks if i != writer_rank]
628 - buffer_io: MessageQueue
697 if group_rank == writer_rank:
630 - buffer_io = MessageQueue(
631 - n_reader=n_reader,
632 - n_local_reader=n_local_reader,
633 - local_reader_ranks=local_reader_ranks,
634 - max_chunk_bytes=max_chunk_bytes,
635 - max_chunks=max_chunks,
636 - )
698 + if external_writer_handle is not None:
699 + buffer_io = MessageQueue.create_from_handle(
700 + external_writer_handle, group_rank
701 + )
702 + else:
703 + same_node_ranks = [i for i, s in enumerate(status) if s]
704 + n_reader = group_world_size - 1
705 + n_local_reader = len(same_node_ranks) - 1
706 + local_reader_ranks = [i for i in same_node_ranks if i != writer_rank]
707 + buffer_io = MessageQueue(
708 + n_reader=n_reader,
709 + n_local_reader=n_local_reader,
710 + local_reader_ranks=local_reader_ranks,
711 + max_chunk_bytes=max_chunk_bytes,
712 + max_chunks=max_chunks,
713 + )
714 handle = buffer_io.export_handle()
715 if isinstance(pg, ProcessGroup):
716 dist.broadcast_object_list(
@@ -651,5 +728,6 @@ class MessageQueue:
728 else:
729 handle = pg.broadcast_obj(None, writer_rank)
730 buffer_io = MessageQueue.create_from_handle(handle, group_rank)
654 - buffer_io.wait_until_ready()
731 + if blocking:
732 + buffer_io.wait_until_ready()
733 return buffer_io
vllm/distributed/parallel_state.py
+73 -4
@@ -385,6 +385,33 @@ class GroupCoordinator:
385 torch.ops._C, "init_shm_manager"
386 )
387
388 + def create_mq_broadcaster(
389 + self, writer_rank=0, external_writer_handle=None, blocking=True
390 + ):
391 + from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
392 +
393 + return MessageQueue.create_from_process_group(
394 + self.cpu_group,
395 + 1 << 22,
396 + 6,
397 + writer_rank=writer_rank,
398 + external_writer_handle=external_writer_handle,
399 + blocking=blocking,
400 + )
401 +
402 + def create_single_reader_mq_broadcasters(
403 + self, reader_rank_in_group=0, blocking=False
404 + ):
405 + from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
406 +
407 + return MessageQueue.create_from_process_group_single_reader(
408 + self.cpu_group,
409 + 1 << 22,
410 + 6,
411 + reader_rank=self.ranks[reader_rank_in_group],
412 + blocking=blocking,
413 + )
414 +
415 @property
416 def first_rank(self):
417 """Return the global rank of the first process in the group"""
@@ -997,6 +1024,7 @@ class GroupCoordinator:
1024
1025
1026 _WORLD: GroupCoordinator | None = None
1027 +_INNER_DP_WORLD: GroupCoordinator | None = None
1028 _NODE_COUNT: int | None = None
1029
1030
@@ -1005,6 +1033,11 @@ def get_world_group() -> GroupCoordinator:
1033 return _WORLD
1034
1035
1036 +def get_inner_dp_world_group() -> GroupCoordinator:
1037 + assert _INNER_DP_WORLD is not None, "inner dp world group is not initialized"
1038 + return _INNER_DP_WORLD
1039 +
1040 +
1041 def init_world_group(
1042 ranks: list[int], local_rank: int, backend: str
1043 ) -> GroupCoordinator:
@@ -1023,12 +1056,13 @@ def init_model_parallel_group(
1056 backend: str,
1057 use_message_queue_broadcaster: bool = False,
1058 group_name: str | None = None,
1059 + use_device_communicator: bool = True,
1060 ) -> GroupCoordinator:
1061 return GroupCoordinator(
1062 group_ranks=group_ranks,
1063 local_rank=local_rank,
1064 torch_distributed_backend=backend,
1031 - use_device_communicator=True,
1065 + use_device_communicator=use_device_communicator,
1066 use_message_queue_broadcaster=use_message_queue_broadcaster,
1067 group_name=group_name,
1068 )
@@ -1143,7 +1177,14 @@ def init_distributed_environment(
1177 from vllm.config import get_current_vllm_config
1178
1179 config = get_current_vllm_config()
1146 - if (
1180 + if config is not None and config.parallel_config.nnodes > 1:
1181 + parallel_config = config.parallel_config
1182 + ip = parallel_config.master_addr
1183 + rank = parallel_config.data_parallel_rank * world_size + rank
1184 + world_size = parallel_config.world_size_across_dp
1185 + port = parallel_config.master_port
1186 + distributed_init_method = get_distributed_init_method(ip, port)
1187 + elif (
1188 config is not None
1189 and config.parallel_config.data_parallel_size > 1
1190 and config.parallel_config.distributed_executor_backend != "external_launcher"
@@ -1164,6 +1205,14 @@ def init_distributed_environment(
1205 distributed_init_method,
1206 )
1207 if not torch.distributed.is_initialized():
1208 + logger.info(
1209 + "world_size=%d rank=%d local_rank=%d distributed_init_method=%s backend=%s",
1210 + world_size,
1211 + rank,
1212 + local_rank,
1213 + distributed_init_method,
1214 + backend,
1215 + )
1216 assert distributed_init_method is not None, (
1217 "distributed_init_method must be provided when initializing "
1218 "distributed environment"
@@ -1192,16 +1241,36 @@ def init_distributed_environment(
1241 # local rank not set, this usually happens in single-node
1242 # setting, where we can use rank as local rank
1243 local_rank = envs.LOCAL_RANK if distributed_init_method == "env://" else rank
1195 - global _WORLD, _NODE_COUNT
1244 + global _WORLD, _NODE_COUNT, _INNER_DP_WORLD
1245 if _WORLD is None:
1246 ranks = list(range(torch.distributed.get_world_size()))
1247 _WORLD = init_world_group(ranks, local_rank, backend)
1199 - _NODE_COUNT = _node_count(_WORLD.cpu_group)
1248 + if config.parallel_config.nnodes > 1:
1249 + _NODE_COUNT = config.parallel_config.nnodes
1250 + else:
1251 + _NODE_COUNT = _node_count(_WORLD.cpu_group)
1252 logger.debug("Detected %d nodes in the distributed environment", _NODE_COUNT)
1253 else:
1254 assert _WORLD.world_size == torch.distributed.get_world_size(), (
1255 "world group already initialized with a different world size"
1256 )
1257 + if config.parallel_config.nnodes_within_dp > 1:
1258 + if parallel_config.data_parallel_size > 1:
1259 + world_size_inner_dp = parallel_config.world_size
1260 + group_ranks = [
1261 + [dp_rank * world_size_inner_dp + i for i in range(world_size_inner_dp)]
1262 + for dp_rank in range(parallel_config.data_parallel_size)
1263 + ]
1264 + _INNER_DP_WORLD = init_model_parallel_group(
1265 + group_ranks,
1266 + get_world_group().local_rank,
1267 + backend,
1268 + use_message_queue_broadcaster=True,
1269 + group_name="inner_dp_world",
1270 + use_device_communicator=False,
1271 + )
1272 + else:
1273 + _INNER_DP_WORLD = _WORLD
1274
1275
1276 def initialize_model_parallel(
vllm/engine/arg_utils.py
+85 -6
@@ -384,6 +384,10 @@ class EngineArgs:
384 ) = ParallelConfig.distributed_executor_backend
385 # number of P/D disaggregation (or other disaggregation) workers
386 pipeline_parallel_size: int = ParallelConfig.pipeline_parallel_size
387 + master_addr: str = ParallelConfig.master_addr
388 + master_port: int = ParallelConfig.master_port
389 + nnodes: int = ParallelConfig.nnodes
390 + node_rank: int = ParallelConfig.node_rank
391 tensor_parallel_size: int = ParallelConfig.tensor_parallel_size
392 decode_context_parallel_size: int = ParallelConfig.decode_context_parallel_size
393 dcp_kv_cache_interleave_size: int = ParallelConfig.dcp_kv_cache_interleave_size
@@ -394,6 +398,7 @@ class EngineArgs:
398 data_parallel_address: str | None = None
399 data_parallel_rpc_port: int | None = None
400 data_parallel_hybrid_lb: bool = False
401 + data_parallel_external_lb: bool = False
402 data_parallel_backend: str = ParallelConfig.data_parallel_backend
403 enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel
404 all2all_backend: str | None = ParallelConfig.all2all_backend
@@ -749,6 +754,10 @@ class EngineArgs:
754 "-pp",
755 **parallel_kwargs["pipeline_parallel_size"],
756 )
757 + parallel_group.add_argument("--master-addr", **parallel_kwargs["master_addr"])
758 + parallel_group.add_argument("--master-port", **parallel_kwargs["master_port"])
759 + parallel_group.add_argument("--nnodes", "-n", **parallel_kwargs["nnodes"])
760 + parallel_group.add_argument("--node-rank", "-r", **parallel_kwargs["node_rank"])
761 parallel_group.add_argument(
762 "--tensor-parallel-size", "-tp", **parallel_kwargs["tensor_parallel_size"]
763 )
@@ -803,7 +812,14 @@ class EngineArgs:
812 help='Backend for data parallel, either "mp" or "ray".',
813 )
814 parallel_group.add_argument(
806 - "--data-parallel-hybrid-lb", **parallel_kwargs["data_parallel_hybrid_lb"]
815 + "--data-parallel-hybrid-lb",
816 + "-dph",
817 + **parallel_kwargs["data_parallel_hybrid_lb"],
818 + )
819 + parallel_group.add_argument(
820 + "--data-parallel-external-lb",
821 + "-dpe",
822 + **parallel_kwargs["data_parallel_external_lb"],
823 )
824 parallel_group.add_argument(
825 "--enable-expert-parallel", **parallel_kwargs["enable_expert_parallel"]
@@ -1428,12 +1444,56 @@ class EngineArgs:
1444 assert not headless or not self.data_parallel_hybrid_lb, (
1445 "data_parallel_hybrid_lb is not applicable in headless mode"
1446 )
1431 -
1432 - data_parallel_external_lb = self.data_parallel_rank is not None
1447 + assert not (self.data_parallel_hybrid_lb and self.data_parallel_external_lb), (
1448 + "data_parallel_hybrid_lb and data_parallel_external_lb cannot both be True."
1449 + )
1450 + assert self.data_parallel_backend == "mp" or self.nnodes == 1, (
1451 + "nnodes > 1 is only supported with data_parallel_backend=mp"
1452 + )
1453 + inferred_data_parallel_rank = 0
1454 + if self.nnodes > 1:
1455 + world_size = (
1456 + self.data_parallel_size
1457 + * self.pipeline_parallel_size
1458 + * self.tensor_parallel_size
1459 + )
1460 + world_size_within_dp = (
1461 + self.pipeline_parallel_size * self.tensor_parallel_size
1462 + )
1463 + local_world_size = world_size // self.nnodes
1464 + assert world_size % self.nnodes == 0, (
1465 + f"world_size={world_size} must be divisible by nnodes={self.nnodes}."
1466 + )
1467 + assert self.node_rank < self.nnodes, (
1468 + f"node_rank={self.node_rank} must be less than nnodes={self.nnodes}."
1469 + )
1470 + inferred_data_parallel_rank = (
1471 + self.node_rank * local_world_size
1472 + ) // world_size_within_dp
1473 + if self.data_parallel_size > 1 and self.data_parallel_external_lb:
1474 + self.data_parallel_rank = inferred_data_parallel_rank
1475 + logger.info(
1476 + "Inferred data_parallel_rank %d from node_rank %d for external lb",
1477 + self.data_parallel_rank,
1478 + self.node_rank,
1479 + )
1480 + elif self.data_parallel_size_local is None:
1481 + # Infer data parallel size local for internal dplb:
1482 + self.data_parallel_size_local = max(
1483 + local_world_size // world_size_within_dp, 1
1484 + )
1485 + data_parallel_external_lb = (
1486 + self.data_parallel_external_lb or self.data_parallel_rank is not None
1487 + )
1488 # Local DP rank = 1, use pure-external LB.
1489 if data_parallel_external_lb:
1490 + assert self.data_parallel_rank is not None, (
1491 + "data_parallel_rank or node_rank must be spefified if "
1492 + "data_parallel_external_lb is enable."
1493 + )
1494 assert self.data_parallel_size_local in (1, None), (
1436 - "data_parallel_size_local must be 1 when data_parallel_rank is set"
1495 + "data_parallel_size_local must be 1 or None when data_parallel_rank "
1496 + "is set"
1497 )
1498 data_parallel_size_local = 1
1499 # Use full external lb if we have local_size of 1.
@@ -1447,6 +1507,11 @@ class EngineArgs:
1507
1508 if self.data_parallel_hybrid_lb and data_parallel_size_local == 1:
1509 # Use full external lb if we have local_size of 1.
1510 + logger.warning(
1511 + "data_parallel_hybrid_lb is not eligible when "
1512 + "data_parallel_size_local = 1, autoswitch to "
1513 + "data_parallel_external_lb."
1514 + )
1515 data_parallel_external_lb = True
1516 self.data_parallel_hybrid_lb = False
1517
@@ -1454,7 +1519,15 @@ class EngineArgs:
1519 # Disable hybrid LB mode if set for a single node
1520 self.data_parallel_hybrid_lb = False
1521
1457 - self.data_parallel_rank = self.data_parallel_start_rank or 0
1522 + self.data_parallel_rank = (
1523 + self.data_parallel_start_rank or inferred_data_parallel_rank
1524 + )
1525 + if self.nnodes > 1:
1526 + logger.info(
1527 + "Inferred data_parallel_rank %d from node_rank %d",
1528 + self.data_parallel_rank,
1529 + self.node_rank,
1530 + )
1531 else:
1532 assert not self.data_parallel_hybrid_lb, (
1533 "data_parallel_size_local must be set to use data_parallel_hybrid_lb."
@@ -1484,7 +1557,9 @@ class EngineArgs:
1557 "data_parallel_backend can only be ray or mp, got %s",
1558 self.data_parallel_backend,
1559 )
1487 - data_parallel_address = ParallelConfig.data_parallel_master_ip
1560 + data_parallel_address = (
1561 + self.master_addr or ParallelConfig.data_parallel_master_ip
1562 + )
1563 else:
1564 data_parallel_address = self.data_parallel_address
1565
@@ -1517,6 +1592,10 @@ class EngineArgs:
1592 data_parallel_rank=self.data_parallel_rank or 0,
1593 data_parallel_external_lb=data_parallel_external_lb,
1594 data_parallel_size_local=data_parallel_size_local,
1595 + master_addr=self.master_addr,
1596 + master_port=self.master_port,
1597 + nnodes=self.nnodes,
1598 + node_rank=self.node_rank,
1599 data_parallel_master_ip=data_parallel_address,
1600 data_parallel_rpc_port=data_parallel_rpc_port,
1601 data_parallel_backend=self.data_parallel_backend,
vllm/entrypoints/cli/serve.py
+27 -4
@@ -24,6 +24,7 @@ from vllm.utils.system_utils import decorate_logs, set_process_title
24 from vllm.v1.engine.core import EngineCoreProc
25 from vllm.v1.engine.utils import CoreEngineProcManager, launch_core_engines
26 from vllm.v1.executor import Executor
27 +from vllm.v1.executor.multiproc_executor import MultiprocExecutor
28 from vllm.v1.metrics.prometheus import setup_multiprocess_prometheus
29 from vllm.v1.utils import APIServerProcessManager, wait_for_completion_or_failure
30
@@ -97,18 +98,40 @@ def run_headless(args: argparse.Namespace):
98 if local_engine_count <= 0:
99 raise ValueError("data_parallel_size_local must be > 0 in headless mode")
100
100 - host = parallel_config.data_parallel_master_ip
101 - port = engine_args.data_parallel_rpc_port # add to config too
102 - handshake_address = get_tcp_uri(host, port)
101 + shutdown_requested = False
102
103 # Catch SIGTERM and SIGINT to allow graceful shutdown.
104 def signal_handler(signum, frame):
105 + nonlocal shutdown_requested
106 logger.debug("Received %d signal.", signum)
107 - raise SystemExit
107 + if not shutdown_requested:
108 + shutdown_requested = True
109 + raise SystemExit
110
111 signal.signal(signal.SIGTERM, signal_handler)
112 signal.signal(signal.SIGINT, signal_handler)
113
114 + if parallel_config.node_rank_within_dp > 0:
115 + from vllm.version import __version__ as VLLM_VERSION
116 +
117 + # Run headless workers (for multi-node PP/TP).
118 + host = parallel_config.master_addr
119 + head_node_address = f"{host}:{parallel_config.master_port}"
120 + logger.info(
121 + "Launching vLLM (v%s) headless multiproc executor, "
122 + "with head node address %s for torch.distributed process group.",
123 + VLLM_VERSION,
124 + head_node_address,
125 + )
126 +
127 + executor = MultiprocExecutor(vllm_config, monitor_workers=False)
128 + executor.start_worker_monitor(inline=True)
129 + return
130 +
131 + host = parallel_config.data_parallel_master_ip
132 + port = parallel_config.data_parallel_rpc_port
133 + handshake_address = get_tcp_uri(host, port)
134 +
135 logger.info(
136 "Launching %d data parallel engine(s) in headless mode, "
137 "with head node address %s.",
vllm/v1/engine/utils.py
+12 -3
@@ -183,15 +183,19 @@ def set_device_control_env_var(
183 for engine subprocess.
184 """
185 world_size = vllm_config.parallel_config.world_size
186 + local_world_size = vllm_config.parallel_config.local_world_size
187 evar = current_platform.device_control_env_var
188
188 - value = get_device_indices(evar, local_dp_rank, world_size)
189 + value = get_device_indices(evar, local_dp_rank, world_size, local_world_size)
190 with patch.dict(os.environ, values=((evar, value),)):
191 yield
192
193
194 def get_device_indices(
194 - device_control_env_var: str, local_dp_rank: int, world_size: int
195 + device_control_env_var: str,
196 + local_dp_rank: int,
197 + world_size: int,
198 + local_world_size: int | None = None,
199 ):
200 """
201 Returns a comma-separated string of device indices for the specified
@@ -200,10 +204,15 @@ def get_device_indices(
204 For example, if world_size=2 and local_dp_rank=1, and there are 4 devices,
205 this will select devices 2 and 3 for local_dp_rank=1.
206 """
207 + if local_world_size is None:
208 + local_world_size = world_size
209 try:
210 value = ",".join(
211 str(current_platform.device_id_to_physical_device_id(i))
206 - for i in range(local_dp_rank * world_size, (local_dp_rank + 1) * world_size)
212 + for i in range(
213 + local_dp_rank * world_size,
214 + local_dp_rank * world_size + local_world_size,
215 + )
216 )
217 except IndexError as e:
218 raise Exception(
vllm/v1/executor/multiproc_executor.py
+150 -47
@@ -10,7 +10,7 @@ import time
10 import traceback
11 import weakref
12 from collections import deque
13 -from collections.abc import Callable
13 +from collections.abc import Callable, Sequence
14 from concurrent.futures import Future, InvalidStateError
15 from contextlib import suppress
16 from dataclasses import dataclass
@@ -34,6 +34,7 @@ from vllm.distributed.parallel_state import (
34 get_dcp_group,
35 get_dp_group,
36 get_ep_group,
37 + get_inner_dp_world_group,
38 get_pp_group,
39 get_tp_group,
40 )
@@ -90,6 +91,10 @@ class FutureWrapper(Future):
91 class MultiprocExecutor(Executor):
92 supports_pp: bool = True
93
94 + def __init__(self, vllm_config: VllmConfig, monitor_workers: bool = True):
95 + self.monitor_workers = monitor_workers
96 + super().__init__(vllm_config)
97 +
98 def _init_executor(self) -> None:
99 # Call self.shutdown at exit to clean up
100 # and ensure workers will be terminated.
@@ -99,6 +104,12 @@ class MultiprocExecutor(Executor):
104 self.failure_callback: FailureCallback | None = None
105
106 self.world_size = self.parallel_config.world_size
107 + assert self.world_size % self.parallel_config.nnodes_within_dp == 0, (
108 + f"global world_size ({self.parallel_config.world_size}) must be "
109 + f"divisible by nnodes_within_dp "
110 + f"({self.parallel_config.nnodes_within_dp}). "
111 + )
112 + self.local_world_size = self.parallel_config.local_world_size
113 tensor_parallel_size = self.parallel_config.tensor_parallel_size
114 pp_parallel_size = self.parallel_config.pipeline_parallel_size
115 assert self.world_size == tensor_parallel_size * pp_parallel_size, (
@@ -116,27 +127,37 @@ class MultiprocExecutor(Executor):
127 distributed_init_method = get_distributed_init_method(
128 get_loopback_ip(), get_open_port()
129 )
119 -
130 + self.rpc_broadcast_mq: MessageQueue | None = None
131 + scheduler_output_handle: Handle | None = None
132 # Initialize worker and set up message queues for SchedulerOutputs
133 # and ModelRunnerOutputs
122 - max_chunk_bytes = envs.VLLM_MQ_MAX_CHUNK_BYTES_MB * 1024 * 1024
123 - self.rpc_broadcast_mq = MessageQueue(
124 - self.world_size, self.world_size, max_chunk_bytes=max_chunk_bytes
125 - )
126 - scheduler_output_handle = self.rpc_broadcast_mq.export_handle()
127 -
134 + if self.parallel_config.node_rank_within_dp == 0:
135 + # For leader node within each dp rank,
136 + # each dp will have its own leader multiproc executor.
137 + max_chunk_bytes = envs.VLLM_MQ_MAX_CHUNK_BYTES_MB * 1024 * 1024
138 + self.rpc_broadcast_mq = MessageQueue(
139 + self.world_size,
140 + self.local_world_size,
141 + max_chunk_bytes=max_chunk_bytes,
142 + connect_ip=self.parallel_config.master_addr,
143 + )
144 + scheduler_output_handle = self.rpc_broadcast_mq.export_handle()
145 # Create workers
146 context = get_mp_context()
147 shared_worker_lock = context.Lock()
148 unready_workers: list[UnreadyWorkerProcHandle] = []
149 success = False
150 try:
134 - for rank in range(self.world_size):
151 + global_start_rank = (
152 + self.local_world_size * self.parallel_config.node_rank_within_dp
153 + )
154 + for local_rank in range(self.local_world_size):
155 + global_rank = global_start_rank + local_rank
156 unready_workers.append(
157 WorkerProc.make_worker_process(
158 vllm_config=self.vllm_config,
138 - local_rank=rank,
139 - rank=rank,
159 + local_rank=local_rank,
160 + rank=global_rank,
161 distributed_init_method=distributed_init_method,
162 input_shm_handle=scheduler_output_handle,
163 shared_worker_lock=shared_worker_lock,
@@ -145,15 +166,38 @@ class MultiprocExecutor(Executor):
166
167 # Workers must be created before wait_for_ready to avoid
168 # deadlock, since worker.init_device() does a device sync.
169 +
170 + # Wait for all local workers to be ready.
171 self.workers = WorkerProc.wait_for_ready(unready_workers)
172
173 + # Start background thread to monitor worker health if not in headless mode.
174 + if self.monitor_workers:
175 + self.start_worker_monitor()
176 +
177 + self.response_mqs = []
178 + # Only leader node have remote response mqs
179 + if self.parallel_config.node_rank_within_dp == 0:
180 + for rank in range(self.world_size):
181 + if rank < self.local_world_size:
182 + local_message_queue = self.workers[rank].worker_response_mq
183 + assert local_message_queue is not None
184 + self.response_mqs.append(local_message_queue)
185 + else:
186 + remote_message_queue = self.workers[0].peer_worker_response_mqs[
187 + rank
188 + ]
189 + assert remote_message_queue is not None
190 + self.response_mqs.append(remote_message_queue)
191 +
192 # Ensure message queues are ready. Will deadlock if re-ordered
193 # Must be kept consistent with the WorkerProc.
152 - self.rpc_broadcast_mq.wait_until_ready()
153 - for w in self.workers:
154 - w.worker_response_mq.wait_until_ready()
194
156 - self.start_worker_monitor()
195 + # Wait for all input mqs to be ready.
196 + if self.rpc_broadcast_mq is not None:
197 + self.rpc_broadcast_mq.wait_until_ready()
198 + # Wait for all remote response mqs to be ready.
199 + for response_mq in self.response_mqs:
200 + response_mq.wait_until_ready()
201 success = True
202 finally:
203 if not success:
@@ -168,7 +212,7 @@ class MultiprocExecutor(Executor):
212
213 self.output_rank = self._get_output_rank()
214
171 - def start_worker_monitor(self):
215 + def start_worker_monitor(self, inline=False) -> None:
216 workers = self.workers
217 self_ref = weakref.ref(self)
218
@@ -192,9 +236,13 @@ class MultiprocExecutor(Executor):
236 _self.failure_callback = None
237 callback()
238
195 - Thread(
196 - target=monitor_workers, daemon=True, name="MultiprocWorkerMonitor"
197 - ).start()
239 + if not inline:
240 + Thread(
241 + target=monitor_workers, daemon=True, name="MultiprocWorkerMonitor"
242 + ).start()
243 + return
244 +
245 + monitor_workers()
246
247 def register_failure_callback(self, callback: FailureCallback):
248 if self.is_failed:
@@ -247,7 +295,9 @@ class MultiprocExecutor(Executor):
295 ) -> Any | list[Any] | Future[Any | list[Any]]:
296 """Returns single result if unique_reply_rank and/or kv_output_aggregator
297 is provided, otherwise list."""
250 -
298 + assert self.rpc_broadcast_mq is not None, (
299 + "collective_rpc should not be called on follower node"
300 + )
301 if self.is_failed:
302 raise RuntimeError("Executor failed.")
303
@@ -269,20 +319,20 @@ class MultiprocExecutor(Executor):
319 send_method = cloudpickle.dumps(method, protocol=pickle.HIGHEST_PROTOCOL)
320 self.rpc_broadcast_mq.enqueue((send_method, args, kwargs, output_rank))
321
272 - workers = (
273 - (self.workers[output_rank],) if output_rank is not None else self.workers
274 - )
322 + response_mqs: Sequence[MessageQueue] = self.response_mqs
323 + if output_rank is not None:
324 + response_mqs = (response_mqs[output_rank],)
325
326 shutdown_event = self.shutdown_event
327
328 def get_response():
329 responses = []
280 - for w in workers:
330 + for mq in response_mqs:
331 dequeue_timeout = (
332 None if deadline is None else (deadline - time.monotonic())
333 )
334 try:
285 - status, result = w.worker_response_mq.dequeue(
335 + status, result = mq.dequeue(
336 timeout=dequeue_timeout, cancel=shutdown_event
337 )
338 except TimeoutError as e:
@@ -391,17 +441,26 @@ class UnreadyWorkerProcHandle:
441 class WorkerProcHandle:
442 proc: BaseProcess
443 rank: int
394 - worker_response_mq: MessageQueue # The worker process writes to this MQ
444 + # The worker process writes to this MQ in single-node mode
445 + worker_response_mq: MessageQueue | None
446 + # This is only non empty on driver node,
447 + # the peer worker process i writes to MQ
448 + # `peer_worker_response_mqs[i]`
449 + peer_worker_response_mqs: list[MessageQueue | None]
450 death_writer: Connection | None = None
451
452 @classmethod
453 def from_unready_handle(
399 - cls, unready_handle: UnreadyWorkerProcHandle, worker_response_mq: MessageQueue
454 + cls,
455 + unready_handle: UnreadyWorkerProcHandle,
456 + worker_response_mq: MessageQueue | None,
457 + peer_worker_response_mqs: list[MessageQueue | None],
458 ) -> "WorkerProcHandle":
459 return cls(
460 proc=unready_handle.proc,
461 rank=unready_handle.rank,
462 worker_response_mq=worker_response_mq,
463 + peer_worker_response_mqs=peer_worker_response_mqs,
464 death_writer=unready_handle.death_writer,
465 )
466
@@ -411,6 +470,38 @@ class WorkerProc:
470
471 READY_STR = "READY"
472
473 + def _init_message_queues(
474 + self, input_shm_handle: Handle, vllm_config: VllmConfig
475 + ) -> None:
476 + if vllm_config.parallel_config.nnodes_within_dp == 1:
477 + # Initialize MessageQueue for receiving SchedulerOutput
478 + self.rpc_broadcast_mq = MessageQueue.create_from_handle(
479 + input_shm_handle, self.worker.rank
480 + )
481 +
482 + # Initializes a message queue for sending the model output
483 + self.worker_response_mq: MessageQueue = MessageQueue(1, 1)
484 + self.peer_response_handles = []
485 + else:
486 + # Initialize remote MessageQueue for receiving SchedulerOutput across nodes
487 + self.rpc_broadcast_mq = get_inner_dp_world_group().create_mq_broadcaster(
488 + external_writer_handle=input_shm_handle,
489 + # Since there is external_writer_handle from executor proc,
490 + # where the ready signal from actual writer is sent out of the
491 + # create_mq_broadcaster method and after this setup, we make it
492 + # non blocking. The handshake will be triggered when
493 + # worker.rpc_broadcast_mq.wait_until_ready() is called
494 + blocking=False,
495 + )
496 + # Initializes remote message queue for sending the model output to the
497 + # driver worker, exposing peer_response_handles for driver worker
498 + # that include handles for all ranks
499 + self.worker_response_mq, self.peer_response_handles = (
500 + get_inner_dp_world_group().create_single_reader_mq_broadcasters(
501 + reader_rank_in_group=0
502 + )
503 + )
504 +
505 def __init__(
506 self,
507 vllm_config: VllmConfig,
@@ -421,13 +512,15 @@ class WorkerProc:
512 shared_worker_lock: LockType,
513 ):
514 self.rank = rank
424 - wrapper = WorkerWrapperBase(vllm_config=vllm_config, rpc_rank=rank)
515 + wrapper = WorkerWrapperBase(
516 + vllm_config=vllm_config, rpc_rank=local_rank, global_rank=rank
517 + )
518 # TODO: move `init_worker` to executor level as a collective rpc call
519 all_kwargs: list[dict] = [
520 {} for _ in range(vllm_config.parallel_config.world_size)
521 ]
522 is_driver_worker = rank % vllm_config.parallel_config.tensor_parallel_size == 0
430 - all_kwargs[rank] = {
523 + all_kwargs[local_rank] = {
524 "vllm_config": vllm_config,
525 "local_rank": local_rank,
526 "rank": rank,
@@ -438,14 +531,6 @@ class WorkerProc:
531 wrapper.init_worker(all_kwargs)
532 self.worker = wrapper
533
441 - # Initialize MessageQueue for receiving SchedulerOutput
442 - self.rpc_broadcast_mq = MessageQueue.create_from_handle(
443 - input_shm_handle, self.worker.rank
444 - )
445 -
446 - # Initializes a message queue for sending the model output
447 - self.worker_response_mq = MessageQueue(1, 1)
448 -
534 scheduler_config = vllm_config.scheduler_config
535 self.use_async_scheduling = scheduler_config.async_scheduling
536 if self.use_async_scheduling:
@@ -466,6 +551,7 @@ class WorkerProc:
551 )
552
553 # Load model
554 + self._init_message_queues(input_shm_handle, vllm_config)
555 self.worker.load_model()
556
557 # Enable environment variable cache (e.g. assume no more
@@ -512,6 +598,27 @@ class WorkerProc:
598 # death_reader in child will get EOFError
599 return UnreadyWorkerProcHandle(proc, rank, reader, death_writer)
600
601 + @staticmethod
602 + def wait_for_response_handle_ready(
603 + handles: dict[str, Any], proc_handle: UnreadyWorkerProcHandle
604 + ) -> WorkerProcHandle:
605 + response_handle = handles["handle"]
606 + worker_response_mq: MessageQueue | None = None
607 + if len(response_handle.local_reader_ranks) > 0:
608 + worker_response_mq = MessageQueue.create_from_handle(response_handle, 0)
609 + peer_response_handles = handles["peer_response_handles"]
610 + peer_worker_response_mqs = [
611 + MessageQueue.create_from_handle(handle, -1)
612 + if handle.remote_subscribe_addr is not None
613 + else None
614 + for handle in peer_response_handles
615 + ]
616 + return WorkerProcHandle.from_unready_handle(
617 + proc_handle,
618 + worker_response_mq,
619 + peer_worker_response_mqs=peer_worker_response_mqs,
620 + )
621 +
622 @staticmethod
623 def wait_for_ready(
624 unready_proc_handles: list[UnreadyWorkerProcHandle],
@@ -537,16 +644,10 @@ class WorkerProc:
644 if response["status"] != "READY":
645 raise e
646
540 - # Extract the message queue handle.
541 - worker_response_mq = MessageQueue.create_from_handle(
542 - response["handle"], 0
543 - )
544 - ready_proc_handles[unready_proc_handle.rank] = (
545 - WorkerProcHandle.from_unready_handle(
546 - unready_proc_handle, worker_response_mq
547 - )
647 + idx = unready_proc_handle.rank % len(ready_proc_handles)
648 + ready_proc_handles[idx] = WorkerProc.wait_for_response_handle_ready(
649 + response, unready_proc_handle
650 )
549 -
651 except EOFError:
652 e.__suppress_context__ = True
653 raise e from None
@@ -618,12 +719,14 @@ class WorkerProc:
719 {
720 "status": WorkerProc.READY_STR,
721 "handle": worker.worker_response_mq.export_handle(),
722 + "peer_response_handles": worker.peer_response_handles,
723 }
724 )
725
726 # Ensure message queues are ready. Will deadlock if re-ordered.
727 # Must be kept consistent with the Executor
626 - worker.rpc_broadcast_mq.wait_until_ready()
728 + if worker.rpc_broadcast_mq is not None:
729 + worker.rpc_broadcast_mq.wait_until_ready()
730 worker.worker_response_mq.wait_until_ready()
731 ready_writer.close()
732 ready_writer = None
vllm/v1/worker/gpu_worker.py
+9 -1
@@ -189,6 +189,7 @@ class Worker(WorkerBase):
189 and self.parallel_config.distributed_executor_backend
190 not in ["ray", "external_launcher"]
191 and self.vllm_config.parallel_config.data_parallel_backend != "ray"
192 + and self.vllm_config.parallel_config.nnodes_within_dp == 1
193 ):
194 # Use local DP rank if available, otherwise use global DP rank.
195 dp_local_rank = self.parallel_config.data_parallel_rank_local
@@ -205,7 +206,14 @@ class Worker(WorkerBase):
206 assert self.local_rank < torch.cuda.device_count(), (
207 f"DP adjusted local rank {self.local_rank} is out of bounds. "
208 )
208 -
209 + visible_device_count = (
210 + torch.cuda.device_count() if torch.cuda.is_available() else 0
211 + )
212 + assert self.parallel_config.local_world_size <= visible_device_count, (
213 + f"local_world_size ({self.parallel_config.local_world_size}) must be "
214 + f"less than or equal to the number of visible devices "
215 + f"({visible_device_count})."
216 + )
217 self.device = torch.device(f"cuda:{self.local_rank}")
218 current_platform.set_device(self.device)
219
vllm/v1/worker/worker_base.py
+3 -1
@@ -180,6 +180,7 @@ class WorkerWrapperBase:
180 self,
181 vllm_config: VllmConfig,
182 rpc_rank: int = 0,
183 + global_rank: int | None = None,
184 ) -> None:
185 """
186 Initialize the worker wrapper with the given vllm_config and rpc_rank.
@@ -192,6 +193,7 @@ class WorkerWrapperBase:
193 group.
194 """
195 self.rpc_rank = rpc_rank
196 + self.global_rank = self.rpc_rank if global_rank is None else global_rank
197 self.worker: WorkerBase | None = None
198
199 # do not store this `vllm_config`, `init_worker` will set the final
@@ -312,7 +314,7 @@ class WorkerWrapperBase:
314 assert self.worker is not None
315
316 def initialize_from_config(self, kv_cache_configs: list[Any]) -> None:
315 - kv_cache_config = kv_cache_configs[self.rpc_rank]
317 + kv_cache_config = kv_cache_configs[self.global_rank]
318 with set_current_vllm_config(self.vllm_config):
319 self.worker.initialize_from_config(kv_cache_config) # type: ignore
320