[V1] AsyncLLM data parallel WIP

njhill · njhill · commit 9ca44ce6bd34 · 2025-02-26T12:32:05.000-08:00
Signed-off-by: Nick Hill &lt;nhill@redhat.com&gt;
diff --git a/vllm/config.py b/vllm/config.py
@@ -37,7 +37,8 @@
 from vllm.transformers_utils.s3_utils import S3Model
 from vllm.transformers_utils.utils import is_s3
 from vllm.utils import (GiB_bytes, LayerBlockType, cuda_device_count_stateless,
-                        get_cpu_memory, random_uuid, resolve_obj_by_qualname)
+                        get_cpu_memory, get_open_port, random_uuid,
+                        resolve_obj_by_qualname)
 
 if TYPE_CHECKING:
     from ray.util.placement_group import PlacementGroup
@@ -1423,10 +1424,19 @@ def __post_init__(self) -> None:
         self.world_size = self.pipeline_parallel_size * \
             self.tensor_parallel_size
 
-        self.data_parallel_size = envs.VLLM_DP_SIZE
-        self.data_parallel_rank = envs.VLLM_DP_RANK
-        self.data_parallel_master_ip = envs.VLLM_DP_MASTER_IP
-        self.data_parallel_master_port = envs.VLLM_DP_MASTER_PORT
+        if self.data_parallel_size > 1:
+            import os
+            if os.getenv("VLLM_ENABLE_V1_MULTIPROCESSING", "1") != "1":
+                raise ValueError(
+                    "VLLM_ENABLE_V1_MULTIPROCESSING can't be disabled when "
+                    "using data parallel.")
+            self.data_parallel_master_port = get_open_port()
+        else:
+            self.data_parallel_size = envs.VLLM_DP_SIZE
+            self.data_parallel_rank = envs.VLLM_DP_RANK
+            self.data_parallel_master_ip = envs.VLLM_DP_MASTER_IP
+            self.data_parallel_master_port = envs.VLLM_DP_MASTER_PORT
+
         self.world_size_across_dp = self.world_size * self.data_parallel_size
 
         if self.distributed_executor_backend == "external_launcher":
diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py
@@ -113,6 +113,7 @@ class EngineArgs:
     # number of P/D disaggregation (or other disaggregation) workers
     pipeline_parallel_size: int = 1
     tensor_parallel_size: int = 1
+    data_parallel_size: int = 1
     max_parallel_loading_workers: Optional[int] = None
     block_size: Optional[int] = None
     enable_prefix_caching: Optional[bool] = None
@@ -430,6 +431,11 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
                             type=int,
                             default=EngineArgs.tensor_parallel_size,
                             help='Number of tensor parallel replicas.')
+        parser.add_argument('--data-parallel-size',
+                            '-dp',
+                            type=int,
+                            default=EngineArgs.data_parallel_size,
+                            help='Number of data parallel replicas.')
         parser.add_argument(
             '--max-parallel-loading-workers',
             type=int,
@@ -1170,6 +1176,7 @@ def create_engine_config(self,
         parallel_config = ParallelConfig(
             pipeline_parallel_size=self.pipeline_parallel_size,
             tensor_parallel_size=self.tensor_parallel_size,
+            data_parallel_size=self.data_parallel_size,
             max_parallel_loading_workers=self.max_parallel_loading_workers,
             disable_custom_all_reduce=self.disable_custom_all_reduce,
             tokenizer_pool_config=TokenizerPoolConfig.create_config(
diff --git a/vllm/v1/core/scheduler.py b/vllm/v1/core/scheduler.py
@@ -20,6 +20,9 @@
 
 logger = init_logger(__name__)
 
+# Used to trigger dummy requests whose outputs should be ignored.
+DUMMY_REQ_ID = "__DUMMY_REQ_ID"
+
 
 class Scheduler:
 
@@ -483,6 +486,7 @@ def update_from_output(
 
         new_running: List[Request] = []
         outputs: List[EngineCoreOutput] = []
+        finished_requests: List[str] = []
 
         # NOTE(woosuk): As len(self.running) can be up to 1K or more, the below
         # loop can be a performance bottleneck. We should do our best to avoid
@@ -564,17 +568,21 @@ def update_from_output(
                     new_logprobs = logprobs.slice(req_index, req_index + 1)
 
             # Transmit partial if chunked prefill & prompt logprobs is enabled
-            if new_token_ids or prompt_logprobs_tensors is not None:
+            if (new_token_ids or prompt_logprobs_tensors is not None) \
+                    and req_id != DUMMY_REQ_ID:
                 # Add EngineCoreOutput for this Request.
+                finish_reason = request.get_finished_reason()
                 outputs.append(
                     EngineCoreOutput(
                         request_id=req_id,
                         new_token_ids=new_token_ids,
-                        finish_reason=request.get_finished_reason(),
+                        finish_reason=finish_reason,
                         new_logprobs=new_logprobs,
                         new_prompt_logprobs_tensors=prompt_logprobs_tensors,
                         stop_reason=request.stop_reason,
                         events=request.take_events()))
+                if finish_reason:
+                    finished_requests.append(req_id)
 
             self.scheduled_req_ids.remove(request.request_id)
             if not stopped:
@@ -583,6 +591,7 @@ def update_from_output(
         self.running = new_running
         return EngineCoreOutputs(
             outputs=outputs,
+            finished_requests=finished_requests,
             scheduler_stats=self.make_stats(),
         )
 
@@ -653,7 +662,7 @@ def get_num_unfinished_requests(self) -> int:
         return len(self.waiting) + len(self.running)
 
     def has_unfinished_requests(self) -> bool:
-        return self.get_num_unfinished_requests() > 0
+        return len(self.running) > 0 or len(self.waiting) > 0
 
     def get_num_unscheduled_requests(self) -> int:
         """Number of requests that are not being processed by the executor."""
diff --git a/vllm/v1/engine/__init__.py b/vllm/v1/engine/__init__.py
@@ -133,6 +133,7 @@ class EngineCoreOutputs(
     timestamp: float = 0.0
 
     utility_output: Optional[UtilityOutput] = None
+    finished_requests: List[str] = []
 
     def __post_init__(self):
         if self.timestamp == 0.0:
diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py
@@ -268,7 +268,10 @@ def __init__(
         ready_pipe.send({"status": "READY"})
 
     @staticmethod
-    def run_engine_core(*args, **kwargs):
+    def run_engine_core(*args,
+                        vllm_config: VllmConfig,
+                        dp_rank: int = 0,
+                        **kwargs):
         """Launch EngineCore busy loop in background process."""
 
         # Signal handler used for graceful termination.
@@ -289,6 +292,9 @@ def signal_handler(signum, frame):
         signal.signal(signal.SIGTERM, signal_handler)
         signal.signal(signal.SIGINT, signal_handler)
 
+        # Set data parallel rank for this engine process.
+        vllm_config.parallel_config.data_parallel_rank = dp_rank
+
         parent_process = psutil.Process().parent()
         engine_core = None
         try:
@@ -313,11 +319,17 @@ def run_busy_loop(self):
         step_fn = (self.step
                    if self.batch_queue is None else self.step_with_batch_queue)
 
+        dp_idle_mode = False
+
         # Loop until process is sent a SIGINT or SIGTERM
         while True:
             # 1) Poll the input queue until there is work to do.
             if not self.scheduler.has_unfinished_requests():
                 while True:
+                    if dp_idle_mode and self.input_queue.empty():
+                        # TODO if time has passed here, break to log stats
+                        self.execute_dummy_batch()
+                        continue
                     try:
                         req = self.input_queue.get(timeout=POLLING_TIMEOUT_S)
                         self._handle_client_request(*req)
@@ -327,14 +339,16 @@ def run_busy_loop(self):
                         # Break out the loop so we can log_stats in step().
                         if self.log_stats:
                             break
-                    except BaseException:
-                        raise
 
             # 2) Handle any new client requests.
             while not self.input_queue.empty():
                 req = self.input_queue.get_nowait()
                 self._handle_client_request(*req)
 
+            if self.scheduler.has_unfinished_requests():
+                # TODO client to reset this in coordinated way
+                dp_idle_mode = True
+
             # 3) Step the engine core.
             outputs = step_fn()
 
diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py
