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 15, 17:15 UTC
[V1] [Feature] Collective RPC (#15444)
Signed-off-by: wwl2755 <wangwenlong2755@gmail.com>
wwl2755 committed
Mar 29, 2025 at 05:39 UTC
94744ba41a2807cb195e4a41a85d4d49f6867967
7 files changed
+86
-10
.buildkite/test-pipeline.yaml
+3
-3
@@ -150,8 +150,8 @@ steps:
150
# TODO: create a dedicated test section for multi-GPU example tests
151
# when we have multiple distributed example tests
152
- pushd ../examples/offline_inference
153
- - VLLM_ENABLE_V1_MULTIPROCESSING=0 python3 rlhf.py
154
- - VLLM_ENABLE_V1_MULTIPROCESSING=0 RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
153
+ - python3 rlhf.py
154
+ - RAY_DEDUP_LOGS=0 python3 rlhf_colocate.py
155
- popd
156
157
- label: Metrics, Tracing Test # 10min
@@ -520,7 +520,7 @@ steps:
520
- vllm/v1/engine/
521
commands:
522
- TP_SIZE=1 DP_SIZE=2 pytest -v -s v1/test_async_llm_dp.py
523
- - VLLM_ENABLE_V1_MULTIPROCESSING=0 pytest -v -s entrypoints/llm/test_collective_rpc.py
523
+ - pytest -v -s entrypoints/llm/test_collective_rpc.py
524
- pytest -v -s ./compile/test_basic_correctness.py
525
- pytest -v -s ./compile/test_wrapper.py
526
- VLLM_TEST_SAME_HOST=1 torchrun --nproc-per-node=4 distributed/test_same_node.py | grep 'Same node test passed'
vllm/engine/llm_engine.py
+11
-2
@@ -7,8 +7,8 @@ from collections import deque
7
from contextlib import contextmanager
8
from dataclasses import dataclass
9
from functools import partial
10
-from typing import (TYPE_CHECKING, Callable, ClassVar, Deque, Dict, Iterable,
11
- List, Mapping, NamedTuple, Optional)
10
+from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Deque, Dict,
11
+ Iterable, List, Mapping, NamedTuple, Optional)
12
from typing import Sequence as GenericSequence
13
from typing import Set, Type, Union, cast, overload
14
@@ -67,6 +67,7 @@ _LOCAL_LOGGING_INTERVAL_SEC = 5
67
68
_G = TypeVar("_G", bound=BaseTokenizerGroup, default=BaseTokenizerGroup)
69
_O = TypeVar("_O", RequestOutput, PoolingRequestOutput)
70
+_R = TypeVar("_R", default=Any)
71
72
73
@dataclass
@@ -2123,6 +2124,14 @@ class LLMEngine:
2124
2125
return sampling_params
2126
2127
+ def collective_rpc(self,
2128
+ method: Union[str, Callable[..., _R]],
2129
+ timeout: Optional[float] = None,
2130
+ args: tuple = (),
2131
+ kwargs: Optional[dict[str, Any]] = None) -> list[_R]:
2132
+ return self.model_executor.collective_rpc(method, timeout, args,
2133
+ kwargs)
2134
+
2135
2136
if envs.is_set("VLLM_USE_V1") and envs.VLLM_USE_V1:
2137
from vllm.v1.engine.llm_engine import LLMEngine as V1LLMEngine
vllm/entrypoints/llm.py
+2
-2
@@ -492,8 +492,8 @@ class LLM:
492
It is recommended to use this API to only pass control messages,
493
and set up data-plane communication to pass data.
494
"""
495
- executor = self.llm_engine.model_executor
496
- return executor.collective_rpc(method, timeout, args, kwargs)
495
+
496
+ return self.llm_engine.collective_rpc(method, timeout, args, kwargs)
497
498
def apply_model(self, func: Callable[[nn.Module], _R]) -> list[_R]:
499
"""
vllm/v1/engine/core.py
+11
-1
@@ -8,7 +8,7 @@ import time
8
from concurrent.futures import Future
9
from inspect import isclass, signature
10
from logging import DEBUG
11
-from typing import Any, Optional
11
+from typing import Any, Callable, Optional, TypeVar, Union
12
13
import msgspec
14
import psutil
@@ -43,6 +43,8 @@ logger = init_logger(__name__)
43
44
POLLING_TIMEOUT_S = 2.5
45
46
+_R = TypeVar('_R') # Return type for collective_rpc
47
+
48
49
class EngineCore:
50
"""Inner loop of vLLM's Engine."""
@@ -280,6 +282,14 @@ class EngineCore:
282
def pin_lora(self, lora_id: int) -> bool:
283
return self.model_executor.pin_lora(lora_id)
284
285
+ def collective_rpc(self,
286
+ method: Union[str, Callable[..., _R]],
287
+ timeout: Optional[float] = None,
288
+ args: tuple = (),
289
+ kwargs: Optional[dict[str, Any]] = None) -> list[_R]:
290
+ return self.model_executor.collective_rpc(method, timeout, args,
291
+ kwargs)
292
+
293
294
class EngineCoreProc(EngineCore):
295
"""ZMQ-wrapper for running EngineCore in background process."""
vllm/v1/engine/core_client.py
+42
-1
@@ -12,7 +12,7 @@ from collections.abc import Awaitable, Sequence
12
from concurrent.futures import Future
13
from dataclasses import dataclass, field
14
from threading import Thread
15
-from typing import Any, Callable, Optional, Union
15
+from typing import Any, Callable, Optional, TypeVar, Union
16
17
import zmq
18
import zmq.asyncio
@@ -33,6 +33,8 @@ logger = init_logger(__name__)
33
34
AnyFuture = Union[asyncio.Future[Any], Future[Any]]
35
36
+_R = TypeVar('_R') # Return type for collective_rpc
37
+
38
39
class EngineCoreClient(ABC):
40
"""
@@ -117,6 +119,13 @@ class EngineCoreClient(ABC):
119
def pin_lora(self, lora_id: int) -> bool:
120
raise NotImplementedError
121
122
+ def collective_rpc(self,
123
+ method: Union[str, Callable[..., _R]],
124
+ timeout: Optional[float] = None,
125
+ args: tuple = (),
126
+ kwargs: Optional[dict[str, Any]] = None) -> list[_R]:
127
+ raise NotImplementedError
128
+
129
async def get_output_async(self) -> EngineCoreOutputs:
130
raise NotImplementedError
131
@@ -153,6 +162,14 @@ class EngineCoreClient(ABC):
162
async def pin_lora_async(self, lora_id: int) -> bool:
163
raise NotImplementedError
164
165
+ async def collective_rpc_async(
166
+ self,
167
+ method: Union[str, Callable[..., _R]],
168
+ timeout: Optional[float] = None,
169
+ args: tuple = (),
170
+ kwargs: Optional[dict[str, Any]] = None) -> list[_R]:
171
+ raise NotImplementedError
172
+
173
174
class InprocClient(EngineCoreClient):
175
"""
@@ -210,6 +227,13 @@ class InprocClient(EngineCoreClient):
227
def pin_lora(self, lora_id: int) -> bool:
228
return self.engine_core.pin_lora(lora_id)
229
230
+ def collective_rpc(self,
231
+ method: Union[str, Callable[..., _R]],
232
+ timeout: Optional[float] = None,
233
+ args: tuple = (),
234
+ kwargs: Optional[dict[str, Any]] = None) -> list[_R]:
235
+ return self.engine_core.collective_rpc(method, timeout, args, kwargs)
236
+
237
238
class CoreEngine:
239
"""One per data parallel rank."""
@@ -505,6 +529,14 @@ class SyncMPClient(MPClient):
529
def execute_dummy_batch(self) -> None:
530
self.call_utility("execute_dummy_batch")
531
532
+ def collective_rpc(self,
533
+ method: Union[str, Callable[..., _R]],
534
+ timeout: Optional[float] = None,
535
+ args: tuple = (),
536
+ kwargs: Optional[dict[str, Any]] = None) -> list[_R]:
537
+ return self.call_utility("collective_rpc", method, timeout, args,
538
+ kwargs)
539
+
540
541
class AsyncMPClient(MPClient):
542
"""Asyncio-compatible client for multi-proc EngineCore."""
@@ -636,6 +668,15 @@ class AsyncMPClient(MPClient):
668
async def pin_lora_async(self, lora_id: int) -> bool:
669
return await self.call_utility_async("pin_lora", lora_id)
670
671
+ async def collective_rpc_async(
672
+ self,
673
+ method: Union[str, Callable[..., _R]],
674
+ timeout: Optional[float] = None,
675
+ args: tuple = (),
676
+ kwargs: Optional[dict[str, Any]] = None) -> list[_R]:
677
+ return await self.call_utility_async("collective_rpc", method, timeout,
678
+ args, kwargs)
679
+
680
681
class DPAsyncMPClient(AsyncMPClient):
682
"""Asyncio-compatible client for multi-proc, multi-engine (data parallel)
vllm/v1/engine/llm_engine.py
+9
-1
@@ -2,7 +2,7 @@
2
3
from collections.abc import Mapping
4
from copy import copy
5
-from typing import Optional, Union
5
+from typing import Any, Callable, Optional, Union
6
7
from typing_extensions import TypeVar
8
@@ -32,6 +32,7 @@ from vllm.v1.executor.abstract import Executor
32
logger = init_logger(__name__)
33
34
_G = TypeVar("_G", bound=BaseTokenizerGroup, default=BaseTokenizerGroup)
35
+_R = TypeVar("_R", default=Any)
36
37
38
class LLMEngine:
@@ -282,6 +283,13 @@ class LLMEngine:
283
"""Prevent an adapter from being evicted."""
284
return self.engine_core.pin_lora(lora_id)
285
286
+ def collective_rpc(self,
287
+ method: Union[str, Callable[..., _R]],
288
+ timeout: Optional[float] = None,
289
+ args: tuple = (),
290
+ kwargs: Optional[dict[str, Any]] = None) -> list[_R]:
291
+ return self.engine_core.collective_rpc(method, timeout, args, kwargs)
292
+
293
def __del__(self):
294
if dp_group := getattr(self, "dp_group", None):
295
stateless_destroy_torch_distributed_process_group(dp_group)
vllm/v1/serial_utils.py
+8
@@ -1,13 +1,16 @@
1
# SPDX-License-Identifier: Apache-2.0
2
3
import pickle
4
+from types import FunctionType
5
from typing import Any, Optional
6
7
+import cloudpickle
8
import torch
9
from msgspec import msgpack
10
11
CUSTOM_TYPE_TENSOR = 1
12
CUSTOM_TYPE_PICKLE = 2
13
+CUSTOM_TYPE_CLOUDPICKLE = 3
14
15
16
class MsgpackEncoder:
@@ -41,6 +44,9 @@ def custom_enc_hook(obj: Any) -> Any:
44
# https://gist.github.com/tlrmchlsmth/8067f1b24a82b6e2f90450e7764fa103 # noqa: E501
45
return msgpack.Ext(CUSTOM_TYPE_TENSOR, pickle.dumps(obj.numpy()))
46
47
+ if isinstance(obj, FunctionType):
48
+ return msgpack.Ext(CUSTOM_TYPE_CLOUDPICKLE, cloudpickle.dumps(obj))
49
+
50
return msgpack.Ext(CUSTOM_TYPE_PICKLE, pickle.dumps(obj))
51
52
@@ -49,5 +55,7 @@ def custom_ext_hook(code: int, data: memoryview) -> Any:
55
return torch.from_numpy(pickle.loads(data))
56
if code == CUSTOM_TYPE_PICKLE:
57
return pickle.loads(data)
58
+ if code == CUSTOM_TYPE_CLOUDPICKLE:
59
+ return cloudpickle.loads(data)
60
61
raise NotImplementedError(f"Extension type code {code} is not supported")