@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3e472d882a
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 15, 22:15 UTC

[core] set up data parallel communication (#13591)

Signed-off-by: youkaichao <youkaichao@gmail.com>

youkaichao committed Feb 22, 2025 at 19:28 UTC 3e472d882a6071813bf6e683f5b9269e0d1d9678
17 files changed +416 -28
.buildkite/test-pipeline.yaml
+2
@@ -134,7 +134,9 @@ steps:
134 - tests/compile/test_basic_correctness
135 - examples/offline_inference/rlhf.py
136 - examples/offline_inference/rlhf_colocate.py
137 + - tests/examples/offline_inference/data_parallel.py
138 commands:
139 + - VLLM_USE_V1=1 python3 ../examples/offline_inference/data_parallel.py
140 - pytest -v -s distributed/test_utils.py
141 - pytest -v -s compile/test_basic_correctness.py
142 - pytest -v -s distributed/test_pynccl.py
examples/offline_inference/data_parallel.py new
+76
@@ -0,0 +1,76 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# usage: VLLM_USE_V1=1 python examples/offline_inference/data_parallel.py
3 +# we need to have a launcher to create multiple data parallel
4 +# ranks. And each rank will create a vLLM instance to process its own prompts.
5 +import os
6 +
7 +from vllm import LLM, SamplingParams
8 +from vllm.utils import get_open_port
9 +
10 +
11 +def main(dp_size, dp_rank, dp_master_ip, dp_master_port, GPUs_per_dp_rank):
12 + os.environ["VLLM_DP_RANK"] = str(dp_rank)
13 + os.environ["VLLM_DP_SIZE"] = str(dp_size)
14 + os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip
15 + os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port)
16 + # set devices for each dp_rank
17 + os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(
18 + str(i) for i in range(dp_rank * GPUs_per_dp_rank, (dp_rank + 1) *
19 + GPUs_per_dp_rank))
20 +
21 + # Sample prompts.
22 + prompts = [
23 + "Hello, my name is",
24 + "The president of the United States is",
25 + "The capital of France is",
26 + "The future of AI is",
27 + ]
28 +
29 + # with DP, each rank should process different prompts.
30 + # usually all the DP ranks process a full dataset,
31 + # and each rank processes a different part of the dataset.
32 + promts_per_rank = len(prompts) // dp_size
33 + start = dp_rank * promts_per_rank
34 + end = start + promts_per_rank
35 + prompts = prompts[start:end]
36 + if len(prompts) == 0:
37 + # if any rank has no prompts to process,
38 + # we need to set a placeholder prompt
39 + prompts = ["Placeholder"]
40 + print(f"DP rank {dp_rank} needs to process {len(prompts)} prompts")
41 +
42 + # Create a sampling params object.
43 + # since we are doing data parallel, every rank can have different
44 + # sampling params. here we set different max_tokens for different
45 + # ranks for demonstration.
46 + sampling_params = SamplingParams(temperature=0.8,
47 + top_p=0.95,
48 + max_tokens=16 * (dp_rank + 1))
49 +
50 + # Create an LLM.
51 + llm = LLM(model="facebook/opt-125m", tensor_parallel_size=2, enforce_eager=True)
52 + outputs = llm.generate(prompts, sampling_params)
53 + # Print the outputs.
54 + for output in outputs:
55 + prompt = output.prompt
56 + generated_text = output.outputs[0].text
57 + print(
58 + f"DP rank {dp_rank}, Prompt: {prompt!r}, "
59 + f"Generated text: {generated_text!r}")
60 +
61 +
62 +if __name__ == "__main__":
63 + from multiprocessing import Process
64 + dp_size = 2
65 + GPUs_per_dp_rank = 2
66 + dp_master_ip = "127.0.0.1"
67 + dp_master_port = get_open_port()
68 + procs = []
69 + for i in range(dp_size):
70 + proc = Process(target=main,
71 + args=(dp_size, i, dp_master_ip, dp_master_port,
72 + GPUs_per_dp_rank))
73 + proc.start()
74 + procs.append(proc)
75 + for proc in procs:
76 + proc.join()
vllm/config.py
+57
@@ -16,6 +16,7 @@ from typing import (TYPE_CHECKING, Any, Callable, ClassVar, Counter, Dict,
16
17 import torch
18 from pydantic import BaseModel, Field, PrivateAttr
19 +from torch.distributed import ProcessGroup, ReduceOp
20 from transformers import PretrainedConfig
21
22 import vllm.envs as envs
@@ -1296,6 +1297,11 @@ class ParallelConfig:
1297
1298 pipeline_parallel_size: int = 1 # Number of pipeline parallel groups.
1299 tensor_parallel_size: int = 1 # Number of tensor parallel groups.
1300 + data_parallel_size: int = 1 # Number of data parallel groups.
1301 + data_parallel_rank: int = 0 # Rank of the data parallel group.
1302 + # IP of the data parallel master.
1303 + data_parallel_master_ip: str = "127.0.0.1"
1304 + data_parallel_master_port: int = 29500 # Port of the data parallel master.
1305
1306 # Maximum number of multiple batches
1307 # when load model sequentially. To avoid RAM OOM when using tensor
@@ -1329,10 +1335,55 @@ class ParallelConfig:
1335 worker_cls: str = "auto"
1336 sd_worker_cls: str = "auto"
1337
1338 + # world_size is TPxPP, it affects the number of workers we create.
1339 world_size: int = field(init=False)
1340 + # world_size_across_dp is TPxPPxDP, it is the size of the world
1341 + # including data parallelism.
1342 + world_size_across_dp: int = field(init=False)
1343
1344 rank: int = 0
1345
1346 + def get_next_dp_init_port(self) -> int:
1347 + """
1348 + We might need to initialize process groups in multiple
1349 + processes that is related to data parallelism,
1350 + e.g. both in the worker and in the engine, which
1351 + can live in different processes. To avoid port conflicts, we
1352 + increment the port number each time we need to initialize a
1353 + new process group related to data parallelism.
1354 + """
1355 + answer = self.data_parallel_master_port
1356 + self.data_parallel_master_port += 1
1357 + return answer
1358 +
1359 + def stateless_init_dp_group(self) -> "ProcessGroup":
1360 + from vllm.distributed.utils import (
1361 + stateless_init_torch_distributed_process_group)
1362 +
1363 + # use gloo since the engine process might not have cuda device
1364 + dp_group = stateless_init_torch_distributed_process_group(
1365 + self.data_parallel_master_ip,
1366 + self.get_next_dp_init_port(),
1367 + self.data_parallel_rank,
1368 + self.data_parallel_size,
1369 + backend="gloo")
1370 +
1371 + return dp_group
1372 +
1373 + @staticmethod
1374 + def has_unfinished_dp(dp_group: "ProcessGroup",
1375 + has_unfinished: bool) -> bool:
1376 + tensor = torch.tensor([has_unfinished],
1377 + dtype=torch.int32,
1378 + device="cpu")
1379 + # dp rank 0: has_unfinished_seqs=True
1380 + # dp rank 1: has_unfinished_seqs=False
1381 + # aggregated: has_unfinished_seqs=True
1382 + # so this is an OR operation, i.e. MAX in integers
1383 + torch.distributed.all_reduce(tensor, op=ReduceOp.MAX, group=dp_group)
1384 + aggregated_has_unfinished = bool(tensor.item())
1385 + return aggregated_has_unfinished
1386 +
1387 def compute_hash(self):
1388 """
1389 Provide a hash that uniquely identifies all the configs
@@ -1350,6 +1401,12 @@ class ParallelConfig:
1401 self.world_size = self.pipeline_parallel_size * \
1402 self.tensor_parallel_size
1403
1404 + self.data_parallel_size = envs.VLLM_DP_SIZE
1405 + self.data_parallel_rank = envs.VLLM_DP_RANK
1406 + self.data_parallel_master_ip = envs.VLLM_DP_MASTER_IP
1407 + self.data_parallel_master_port = envs.VLLM_DP_MASTER_PORT
1408 + self.world_size_across_dp = self.world_size * self.data_parallel_size
1409 +
1410 ray_only_devices = ["tpu"]
1411 from vllm.platforms import current_platform
1412 if (current_platform.device_type in ray_only_devices
vllm/distributed/device_communicators/cuda_communicator.py
+2 -2
@@ -16,8 +16,8 @@ class CudaCommunicator(DeviceCommunicatorBase):
16 device_group: Optional[ProcessGroup] = None,
17 unique_name: str = ""):
18 super().__init__(cpu_group, device, device_group, unique_name)
19 - if "pp" in unique_name:
20 - # pipeline parallel does not need custom allreduce
19 + if "tp" not in unique_name:
20 + # only tp uses custom allreduce
21 use_custom_allreduce = False
22 else:
23 from vllm.distributed.parallel_state import (
vllm/distributed/device_communicators/custom_all_reduce.py
+7 -4
@@ -87,6 +87,7 @@ class CustomAllreduce:
87 return
88
89 rank = dist.get_rank(group=self.group)
90 + self.rank = rank
91 world_size = dist.get_world_size(group=self.group)
92 if world_size == 1:
93 # No need to initialize custom allreduce for single GPU case.
@@ -201,8 +202,10 @@ class CustomAllreduce:
202
203 @staticmethod
204 def free_shared_buffer(pointers: List[int],
204 - group: Optional[ProcessGroup] = None) -> None:
205 - rank = dist.get_rank(group=group)
205 + group: Optional[ProcessGroup] = None,
206 + rank: Optional[int] = None) -> None:
207 + if rank is None:
208 + rank = dist.get_rank(group=group)
209 lib = CudaRTLibrary()
210 lib.cudaFree(ctypes.c_void_p(pointers[rank]))
211
@@ -298,8 +301,8 @@ class CustomAllreduce:
301 if not self.disabled and self._ptr:
302 ops.dispose(self._ptr)
303 self._ptr = 0
301 - self.free_shared_buffer(self.meta_ptrs)
302 - self.free_shared_buffer(self.buffer_ptrs)
304 + self.free_shared_buffer(self.meta_ptrs, rank=self.rank)
305 + self.free_shared_buffer(self.buffer_ptrs, rank=self.rank)
306
307 def __del__(self):
308 self.close()
vllm/distributed/parallel_state.py
+62 -14
@@ -750,6 +750,13 @@ get_tensor_model_parallel_group = get_tp_group
750
751 _PP: Optional[GroupCoordinator] = None
752
753 +_DP: Optional[GroupCoordinator] = None
754 +
755 +
756 +def get_dp_group() -> GroupCoordinator:
757 + assert _DP is not None, ("data parallel group is not initialized")
758 + return _DP
759 +
760
761 def get_pp_group() -> GroupCoordinator:
762 assert _PP is not None, (
@@ -811,6 +818,21 @@ def init_distributed_environment(
818 "world_size=%d rank=%d local_rank=%d "
819 "distributed_init_method=%s backend=%s", world_size, rank, local_rank,
820 distributed_init_method, backend)
821 + from vllm.config import get_current_vllm_config
822 + config = get_current_vllm_config()
823 + if config is not None and config.parallel_config.data_parallel_size > 1:
824 + parallel_config = config.parallel_config
825 + # adjust to take into account data parallelism
826 + # offset the rank by the data parallel rank
827 + rank = parallel_config.data_parallel_rank * world_size + rank
828 + # adjust the world size to take into account data parallelism
829 + world_size = parallel_config.world_size_across_dp
830 + ip = parallel_config.data_parallel_master_ip
831 + port = parallel_config.get_next_dp_init_port()
832 + distributed_init_method = f"tcp://{ip}:{port}" # noqa
833 + logger.info(
834 + "Adjusting world_size=%d rank=%d distributed_init_method=%s for DP",
835 + world_size, rank, distributed_init_method)
836 if not torch.distributed.is_initialized():
837 assert distributed_init_method is not None, (
838 "distributed_init_method must be provided when initializing "
@@ -870,20 +892,28 @@ def initialize_model_parallel(
892 # Get world size and rank. Ensure some consistencies.
893 assert torch.distributed.is_initialized()
894 world_size: int = torch.distributed.get_world_size()
895 + rank = torch.distributed.get_rank()
896 backend = backend or torch.distributed.get_backend(
897 get_world_group().device_group)
898
899 + data_parallel_size = 1
900 + from vllm.config import get_current_vllm_config
901 + config = get_current_vllm_config()
902 + if config is not None:
903 + data_parallel_size = config.parallel_config.data_parallel_size
904 +
905 + # the layout order is: DP x PP x TP
906 + # to get group_ranks for each dimension, transpose that dimension to the
907 + # last dimension, then reshape to 2D, then unbind the last dimension
908 + all_ranks = torch.arange(world_size).reshape(
909 + data_parallel_size, pipeline_model_parallel_size,
910 + tensor_model_parallel_size) # noqa
911 +
912 # Build the tensor model-parallel groups.
877 - num_tensor_model_parallel_groups: int = (world_size //
878 - tensor_model_parallel_size)
913 global _TP
914 assert _TP is None, ("tensor model parallel group is already initialized")
881 - group_ranks = []
882 - for i in range(num_tensor_model_parallel_groups):
883 - ranks = list(
884 - range(i * tensor_model_parallel_size,
885 - (i + 1) * tensor_model_parallel_size))
886 - group_ranks.append(ranks)
915 + group_ranks = all_ranks.view(-1, tensor_model_parallel_size).unbind(0)
916 + group_ranks = [x.tolist() for x in group_ranks]
917
918 # message queue broadcaster is only used in tensor model parallel group
919 _TP = init_model_parallel_group(group_ranks,
@@ -893,20 +923,33 @@ def initialize_model_parallel(
923 group_name="tp")
924
925 # Build the pipeline model-parallel groups.
896 - num_pipeline_model_parallel_groups: int = (world_size //
897 - pipeline_model_parallel_size)
926 global _PP
927 assert _PP is None, (
928 "pipeline model parallel group is already initialized")
901 - group_ranks = []
902 - for i in range(num_pipeline_model_parallel_groups):
903 - ranks = list(range(i, world_size, num_pipeline_model_parallel_groups))
904 - group_ranks.append(ranks)
929 + group_ranks = all_ranks.transpose(1, 2).reshape(
930 + -1, pipeline_model_parallel_size).unbind(0)
931 + group_ranks = [x.tolist() for x in group_ranks]
932 _PP = init_model_parallel_group(group_ranks,
933 get_world_group().local_rank,
934 backend,
935 group_name="pp")
936
937 + global _DP
938 + assert _DP is None, ("data parallel group is already initialized")
939 + group_ranks = all_ranks.transpose(0,
940 + 2).reshape(-1,
941 + data_parallel_size).unbind(0)
942 + group_ranks = [x.tolist() for x in group_ranks]
943 + _DP = init_model_parallel_group(group_ranks,
944 + get_world_group().local_rank,
945 + backend,
946 + group_name="dp")
947 +
948 + logger.info(
949 + "rank %s in world size %s is assigned as "
950 + "DP rank %s, PP rank %s, TP rank %s", rank, world_size,
951 + _DP.rank_in_group, _PP.rank_in_group, _TP.rank_in_group)
952 +
953
954 def ensure_kv_transfer_initialized(vllm_config: "VllmConfig") -> None:
955 """
@@ -1011,6 +1054,11 @@ def destroy_model_parallel():
1054 _PP.destroy()
1055 _PP = None
1056
1057 + global _DP
1058 + if _DP:
1059 + _DP.destroy()
1060 + _DP = None
1061 +
1062
1063 def destroy_distributed_environment():
1064 global _WORLD
vllm/distributed/utils.py
+90 -1
@@ -11,7 +11,11 @@ from collections import deque
11 from typing import Any, Deque, Dict, Optional, Sequence, Tuple
12
13 import torch
14 -from torch.distributed import TCPStore
14 +from torch.distributed import ProcessGroup, TCPStore
15 +from torch.distributed.distributed_c10d import (Backend, PrefixStore,
16 + _get_default_timeout,
17 + is_nccl_available)
18 +from torch.distributed.rendezvous import rendezvous
19
20 import vllm.envs as envs
21 from vllm.logger import init_logger
@@ -227,3 +231,88 @@ class StatelessProcessGroup:
231 world_size=world_size,
232 store=store,
233 data_expiration_seconds=data_expiration_seconds)
234 +
235 +
236 +def stateless_init_torch_distributed_process_group(
237 + host: str, port: int, rank: int, world_size: int,
238 + backend: str) -> ProcessGroup:
239 + """
240 + A replacement for `torch.distributed.init_process_group` that does not
241 + pollute the global state. The created ProcessGroup object can be used for
242 + some operations such as `allreduce`, because it does not depend on the
243 + global rank. However, some operations such as `broadcast` cannot be used
244 + because it depends on the global rank.
245 +
246 + # TODO: ask for help from PyTorch team if we need the `broadcast` operation.
247 +
248 + This function is useful when we are not sure about the total number of
249 + processes in the process group. For example, we may have process
250 + 1, 2, ..., 8 who want to communicate, and process 9 might be the same
251 + process as process 1, or it might be a different process; process 10
252 + might be the same process as process 5, or it might be a different process.
253 + In this case, how can we reliably form a communication channel within
254 + process 9 and 10, without affecting the communication channel within
255 + process 1, 2, ..., 8?
256 +
257 + One possible solution is to figure out if process 9 and 10 are the same
258 + as process 1 and 5 beforehand, and then form a communication channel
259 + based on the information, adjusting the ranks and world_size etc. However,
260 + figuring out the information is not always easy, and it will interfere
261 + with the main communication channel.
262 +
263 + Our solution is to always form a communication channel with process 1, 2,
264 + ..., 8, and then use this function to form another communication channel
265 + with process 9 and 10. This way, regardless of whether process 9 and 10
266 + are the same as process 1 and 5, the main communication channel is
267 + always formed with process 1, 2, ..., 8, and the additional communication
268 + channel is formed with process 9 and 10.
269 + """
270 + init_method = f"tcp://{host}:{port}"
271 + backend = Backend(backend) # it is basically string
272 + timeout = _get_default_timeout(backend)
273 +
274 + store, rank, world_size = next(
275 + rendezvous(init_method, rank, world_size, timeout=timeout))
276 + store.set_timeout(timeout)
277 +
278 + group_rank = rank
279 + group_size = world_size
280 +
281 + # Use a PrefixStore to avoid accidental overrides of keys used by
282 + # different systems (e.g. RPC) in case the store is multi-tenant.
283 + prefix_store = PrefixStore(init_method, store)
284 +
285 + pg_options = ProcessGroup.Options(backend=backend, timeout=timeout)
286 +
287 + pg: ProcessGroup = ProcessGroup(
288 + prefix_store,
289 + group_rank,
290 + group_size,
291 + pg_options,
292 + )
293 +
294 + if backend == "gloo":
295 + from torch.distributed.distributed_c10d import ProcessGroupGloo
296 + backend_class = ProcessGroupGloo(prefix_store,
297 + group_rank,
298 + group_size,
299 + timeout=timeout)
300 + backend_type = ProcessGroup.BackendType.GLOO
301 + device = torch.device("cpu")
302 + elif backend == "nccl":
303 + assert is_nccl_available()
304 + from torch.distributed.distributed_c10d import ProcessGroupNCCL
305 +
306 + backend_options = ProcessGroupNCCL.Options()
307 + backend_options._timeout = timeout
308 +
309 + backend_class = ProcessGroupNCCL(prefix_store, group_rank, group_size,
310 + backend_options)
311 + backend_type = ProcessGroup.BackendType.NCCL
312 + device = torch.device("cuda")
313 +
314 + backend_class._set_sequence_number_for_group()
315 +
316 + pg._register_backend(device, backend_type, backend_class)
317 +
318 + return pg
vllm/envs.py
+20
@@ -90,6 +90,10 @@ if TYPE_CHECKING:
90 VLLM_RAY_BUNDLE_INDICES: str = ""
91 VLLM_CUDART_SO_PATH: Optional[str] = None
92 VLLM_USE_HPU_CONTIGUOUS_CACHE_FETCH: bool = True
93 + VLLM_DP_RANK: int = 0
94 + VLLM_DP_SIZE: int = 1
95 + VLLM_DP_MASTER_IP: str = ""
96 + VLLM_DP_MASTER_PORT: int = 0
97
98
99 def get_default_cache_root():
@@ -593,6 +597,22 @@ environment_variables: Dict[str, Callable[[], Any]] = {
597 "VLLM_USE_HPU_CONTIGUOUS_CACHE_FETCH":
598 lambda: os.environ.get("VLLM_CONTIGUOUS_PA", "true").lower() in
599 ("1", "true"),
600 +
601 + # Rank of the process in the data parallel setting
602 + "VLLM_DP_RANK":
603 + lambda: int(os.getenv("VLLM_DP_RANK", "0")),
604 +
605 + # World size of the data parallel setting
606 + "VLLM_DP_SIZE":
607 + lambda: int(os.getenv("VLLM_DP_SIZE", "1")),
608 +
609 + # IP address of the master node in the data parallel setting
610 + "VLLM_DP_MASTER_IP":
611 + lambda: os.getenv("VLLM_DP_MASTER_IP", "127.0.0.1"),
612 +
613 + # Port of the master node in the data parallel setting
614 + "VLLM_DP_MASTER_PORT":
615 + lambda: int(os.getenv("VLLM_DP_MASTER_PORT", "0")),
616 }
617
618 # end-env-vars-definition
vllm/forward_context.py
+31 -3
@@ -4,9 +4,10 @@ import time
4 from collections import defaultdict
5 from contextlib import contextmanager
6 from dataclasses import dataclass
7 -from typing import TYPE_CHECKING, Any, Dict, Optional
7 +from typing import TYPE_CHECKING, Any, Dict, List, Optional
8
9 import torch
10 +import torch.distributed as dist
11
12 import vllm.envs as envs
13 from vllm.config import VllmConfig
@@ -32,6 +33,8 @@ class ForwardContext:
33 attn_metadata: "AttentionMetadata" # set dynamically for each forward pass
34 # TODO: remove after making all virtual_engines share the same kv cache
35 virtual_engine: int # set dynamically for each forward pass
36 + num_tokens_across_dp: Optional[
37 + List[int]] = None # set dynamically for each forward pass
38
39
40 _forward_context: Optional[ForwardContext] = None
@@ -48,7 +51,8 @@ def get_forward_context() -> ForwardContext:
51 @contextmanager
52 def set_forward_context(attn_metadata: Any,
53 vllm_config: VllmConfig,
51 - virtual_engine: int = 0):
54 + virtual_engine: int = 0,
55 + num_tokens: int = 0):
56 """A context manager that stores the current forward context,
57 can be attention metadata, etc.
58 Here we can inject common logic for every model forward pass.
@@ -57,12 +61,36 @@ def set_forward_context(attn_metadata: Any,
61 need_to_track_batchsize = track_batchsize and attn_metadata is not None
62 if need_to_track_batchsize:
63 forward_start_time = time.perf_counter()
64 + num_tokens_across_dp = None
65 + if vllm_config.parallel_config.data_parallel_size > 1:
66 + dp_size = vllm_config.parallel_config.data_parallel_size
67 + dp_rank = vllm_config.parallel_config.data_parallel_rank
68 + if attn_metadata is not None:
69 + if hasattr(attn_metadata, "num_prefill_tokens"):
70 + # for v0 attention backends
71 + batchsize = attn_metadata.num_prefill_tokens + \
72 + attn_metadata.num_decode_tokens
73 + else:
74 + # for v1 attention backends
75 + batchsize = attn_metadata.num_input_tokens
76 + else:
77 + batchsize = num_tokens
78 + num_tokens_across_dp = [0] * dp_size
79 + num_tokens_across_dp[dp_rank] = batchsize
80 + num_tokens_tensor = torch.tensor(num_tokens_across_dp,
81 + device="cpu",
82 + dtype=torch.int32)
83 + from vllm.distributed.parallel_state import get_dp_group
84 + dist.all_reduce(num_tokens_tensor, group=get_dp_group().cpu_group)
85 + num_tokens_across_dp = num_tokens_tensor.tolist()
86 +
87 global _forward_context
88 prev_context = _forward_context
89 _forward_context = ForwardContext(
90 attn_layers=vllm_config.compilation_config.static_forward_context,
91 virtual_engine=virtual_engine,
65 - attn_metadata=attn_metadata)
92 + attn_metadata=attn_metadata,
93 + num_tokens_across_dp=num_tokens_across_dp)
94 try:
95 yield
96 finally:
vllm/utils.py
+18
@@ -501,6 +501,24 @@ def get_open_zmq_ipc_path() -> str:
501
502
503 def get_open_port() -> int:
504 + """
505 + Get an open port for the vLLM process to listen on.
506 + An edge case to handle, is when we run data parallel,
507 + we need to avoid ports that are potentially used by
508 + the data parallel master process.
509 + Right now we reserve 10 ports for the data parallel master
510 + process. Currently it uses 2 ports.
511 + """
512 + if "VLLM_DP_MASTER_PORT" in os.environ:
513 + dp_port = envs.VLLM_DP_MASTER_PORT
514 + while True:
515 + port = _get_open_port()
516 + if port >= dp_port and port < dp_port + 10:
517 + continue
518 + return port
519 + return _get_open_port()
520 +
521 +def _get_open_port() -> int:
522 port = envs.VLLM_PORT
523 if port is not None:
524 while True:
vllm/v1/engine/core.py
+3
@@ -219,6 +219,9 @@ class EngineCore:
219 def wake_up(self):
220 self.model_executor.wake_up()
221
222 + def execute_dummy_batch(self):
223 + self.model_executor.collective_rpc("execute_dummy_batch")
224 +
225 def add_lora(self, lora_request: LoRARequest) -> None:
226 self.model_executor.add_lora(lora_request)
227
vllm/v1/engine/core_client.py
+14
@@ -87,6 +87,12 @@ class EngineCoreClient(ABC):
87 def wake_up(self) -> None:
88 raise NotImplementedError
89
90 + def execute_dummy_batch(self) -> None:
91 + raise NotImplementedError
92 +
93 + async def execute_dummy_batch_async(self) -> None:
94 + raise NotImplementedError
95 +
96 def abort_requests(self, request_ids: List[str]) -> None:
97 raise NotImplementedError
98
@@ -156,6 +162,9 @@ class InprocClient(EngineCoreClient):
162 def wake_up(self) -> None:
163 self.engine_core.wake_up()
164
165 + def execute_dummy_batch(self) -> None:
166 + self.engine_core.execute_dummy_batch()
167 +
168 def add_lora(self, lora_request: LoRARequest) -> None:
169 self.engine_core.add_lora(lora_request)
170
@@ -331,6 +340,8 @@ class SyncMPClient(MPClient):
340 def wake_up(self) -> None:
341 self._call_utility("wake_up")
342
343 + def execute_dummy_batch(self) -> None:
344 + self._call_utility("execute_dummy_batch")
345
346 class AsyncMPClient(MPClient):
347 """Asyncio-compatible client for multi-proc EngineCore."""
@@ -414,5 +425,8 @@ class AsyncMPClient(MPClient):
425 async def wake_up_async(self) -> None:
426 await self._call_utility_async("wake_up")
427
428 + async def execute_dummy_batch_async(self) -> None:
429 + await self._call_utility_async("execute_dummy_batch")
430 +
431 async def add_lora_async(self, lora_request: LoRARequest) -> None:
432 await self._call_utility_async("add_lora", lora_request)
vllm/v1/engine/llm_engine.py
+24 -2
@@ -4,7 +4,7 @@ from typing import Dict, List, Mapping, Optional, Type, Union
4
5 from typing_extensions import TypeVar
6
7 -from vllm.config import VllmConfig
7 +from vllm.config import ParallelConfig, VllmConfig
8 from vllm.engine.arg_utils import EngineArgs
9 from vllm.engine.metrics_types import StatLoggerBase
10 from vllm.envs import VLLM_ENABLE_V1_MULTIPROCESSING
@@ -47,6 +47,13 @@ class LLMEngine:
47 self.model_config = vllm_config.model_config
48 self.cache_config = vllm_config.cache_config
49
50 + # important: init dp group before init the engine_core
51 + self.parallel_config = vllm_config.parallel_config
52 + self.dp_enabled = self.parallel_config.data_parallel_size > 1 # noqa
53 + self.should_execute_dummy_batch = False
54 + if self.dp_enabled:
55 + self.dp_group = self.parallel_config.stateless_init_dp_group()
56 +
57 # Tokenizer (+ ensure liveness if running in another process).
58 self.tokenizer = init_tokenizer_from_configs(
59 model_config=vllm_config.model_config,
@@ -106,7 +113,17 @@ class LLMEngine:
113 return self.output_processor.get_num_unfinished_requests()
114
115 def has_unfinished_requests(self) -> bool:
109 - return self.output_processor.has_unfinished_requests()
116 + has_unfinished = self.output_processor.has_unfinished_requests()
117 + if not self.dp_enabled:
118 + return has_unfinished
119 + return self.has_unfinished_requests_dp(has_unfinished)
120 +
121 + def has_unfinished_requests_dp(self, has_unfinished: bool) -> bool:
122 + aggregated_has_unfinished = ParallelConfig.has_unfinished_dp(
123 + self.dp_group, has_unfinished)
124 + if not has_unfinished and aggregated_has_unfinished:
125 + self.should_execute_dummy_batch = True
126 + return aggregated_has_unfinished
127
128 @classmethod
129 def validate_outputs(cls, outputs, output_type):
@@ -145,6 +162,11 @@ class LLMEngine:
162
163 def step(self) -> List[RequestOutput]:
164
165 + if self.should_execute_dummy_batch:
166 + self.should_execute_dummy_batch = False
167 + self.engine_core.execute_dummy_batch()
168 + return []
169 +
170 # 1) Get EngineCoreOutput from the EngineCore.
171 outputs = self.engine_core.get_output()
172
vllm/v1/executor/multiproc_executor.py
+1 -1
@@ -239,7 +239,7 @@ class WorkerProc:
239 ready_socket.send_string(WorkerProc.READY_STR)
240 ready_socket.send(payload)
241
242 - self.worker.init_device()
242 + wrapper.init_device()
243 self.worker.load_model()
244
245 @staticmethod
vllm/v1/worker/gpu_model_runner.py
+1 -1
@@ -1167,7 +1167,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1167 for k, v in self.intermediate_tensors.items()
1168 })
1169
1170 - with set_forward_context(None, self.vllm_config):
1170 + with set_forward_context(None, self.vllm_config, num_tokens=num_tokens):
1171 hidden_states = model(
1172 input_ids=input_ids,
1173 positions=positions,
vllm/v1/worker/gpu_worker.py
+3
@@ -235,6 +235,9 @@ class Worker(WorkerBase):
235 else:
236 self.profiler.stop()
237
238 + def execute_dummy_batch(self) -> None:
239 + self.model_runner._dummy_run(1)
240 +
241 def add_lora(self, lora_request: LoRARequest) -> bool:
242 return self.model_runner.add_lora(lora_request)
243
vllm/worker/worker_base.py
+5
@@ -567,6 +567,11 @@ class WorkerWrapperBase:
567 self.worker = worker_class(**kwargs)
568 assert self.worker is not None
569
570 + def init_device(self):
571 + with set_current_vllm_config(self.vllm_config):
572 + # To make vLLM config available during device initialization
573 + self.worker.init_device() # type: ignore
574 +
575 def execute_method(self, method: Union[str, bytes], *args, **kwargs):
576 try:
577 target = self if self.worker is None else self.worker