[WIP] Two batch overlap

examples/basic-ub.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import os
+
+from vllm import LLM, SamplingParams
+
+# Sample prompts.
+prompts = [
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
+]
+# Configure logging level for vllm (optional, uses VLLM_LOGGING_LEVEL env var).
+logging_level = os.getenv("VLLM_LOGGING_LEVEL", "").upper()
+if logging_level:
+    logging.basicConfig(level=getattr(logging, logging_level, logging.INFO))
+
+# Create a sampling params object, optionally limiting output tokens via MAX_TOKENS env var.
+param_kwargs = {"temperature": 0.8, "top_p": 0.95}
+max_tokens_env = os.getenv("MAX_TOKENS")
+if max_tokens_env is not None:
+    try:
+        param_kwargs["max_tokens"] = int(max_tokens_env)
+    except ValueError:
+        raise ValueError(f"Invalid MAX_TOKENS value: {max_tokens_env}")
+sampling_params = SamplingParams(**param_kwargs)
+
+
+def main():
+    # Create an LLM.
+    model = "deepseek-ai/DeepSeek-V2-Lite"
+    # model = "facebook/opt-125m"
+    llm = LLM(model=model, 
+              enforce_eager=True,
+              compilation_config=2,
+              ###############
+              trust_remote_code=True,
+              max_model_len=1024,
+              #load_format="dummy",
+              ###############
+              #tensor_parallel_size=1,
+              data_parallel_size=2,
+              enable_expert_parallel=True,
+              ###############
+              #enable_microbatching=True, 
+    )
+    # Generate texts from the prompts.
+    # The output is a list of RequestOutput objects
+    # that contain the prompt, generated text, and other information.
+    outputs = llm.generate(prompts, sampling_params)
+    # Print the outputs.
+    print("\nGenerated Outputs:\n" + "-" * 60)
+    for output in outputs:
+        prompt = output.prompt
+        generated_text = output.outputs[0].text
+        print(f"Prompt:    {prompt!r}")
+        print(f"Output:    {generated_text!r}")
+        print("-" * 60)
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/data_parallel.py

@@ -32,66 +32,43 @@
 import os
 from time import sleep
 
-from vllm import LLM, SamplingParams
-from vllm.utils import get_open_port
+from vllm import LLM, EngineArgs, SamplingParams
+from vllm.utils import FlexibleArgumentParser, get_open_port
 
 
 def parse_args():
-    import argparse
-
-    parser = argparse.ArgumentParser(description="Data Parallel Inference")
-    parser.add_argument(
-        "--model",
-        type=str,
-        default="ibm-research/PowerMoE-3b",
-        help="Model name or path",
-    )
-    parser.add_argument("--dp-size", type=int, default=2, help="Data parallel size")
-    parser.add_argument("--tp-size", type=int, default=2, help="Tensor parallel size")
-    parser.add_argument(
-        "--node-size", type=int, default=1, help="Total number of nodes"
-    )
-    parser.add_argument(
-        "--node-rank", type=int, default=0, help="Rank of the current node"
-    )
-    parser.add_argument(
-        "--master-addr", type=str, default="", help="Master node IP address"
-    )
-    parser.add_argument("--master-port", type=int, default=0, help="Master node port")
-    parser.add_argument(
-        "--enforce-eager", action="store_true", help="Enforce eager mode execution."
-    )
-    parser.add_argument(
-        "--trust-remote-code", action="store_true", help="Trust remote code."
-    )
-    parser.add_argument(
-        "--max-num-seqs",
-        type=int,
-        default=64,
-        help=("Maximum number of sequences to be processed in a single iteration."),
-    )
-    parser.add_argument(
-        "--gpu-memory-utilization",
-        type=float,
-        default=0.8,
-        help=("Fraction of GPU memory vLLM is allowed to allocate (0.0, 1.0]."),
-    )
+    parser = FlexibleArgumentParser()
+    EngineArgs.add_cli_args(parser)
+    parser.set_defaults(model="ibm-research/PowerMoE-3b")
+    parser.add_argument("--dp-size",
+                        type=int,
+                        default=2,
+                        help="Data parallel size")
+    parser.add_argument("--tp-size",
+                        type=int,
+                        default=2,
+                        help="Tensor parallel size")
+    parser.add_argument("--node-size",
+                        type=int,
+                        default=1,
+                        help="Total number of nodes")
+    parser.add_argument("--node-rank",
+                        type=int,
+                        default=0,
+                        help="Rank of the current node")
+    parser.add_argument("--master-addr",
+                        type=str,
+                        default="",
+                        help="Master node IP address")
+    parser.add_argument("--master-port",
+                        type=int,
+                        default=0,
+                        help="Master node port")
     return parser.parse_args()
 
 
-def main(
-    model,
-    dp_size,
-    local_dp_rank,
-    global_dp_rank,
-    dp_master_ip,
-    dp_master_port,
-    GPUs_per_dp_rank,
-    enforce_eager,
-    trust_remote_code,
-    max_num_seqs,
-    gpu_memory_utilization,
-):
+def main(args, dp_size, local_dp_rank, global_dp_rank, dp_master_ip,
+         dp_master_port, GPUs_per_dp_rank):
     os.environ["VLLM_DP_RANK"] = str(global_dp_rank)
     os.environ["VLLM_DP_RANK_LOCAL"] = str(local_dp_rank)
     os.environ["VLLM_DP_SIZE"] = str(dp_size)
@@ -107,7 +84,11 @@ def main(
         "The president of the United States is",
         "The capital of France is",
         "The future of AI is",
-    ] * 100
+    ] * 10
+    # import random
+    # import string
+    # prompts = [''.join(random.choices(string.ascii_letters, k=128)) for _ in range(2048)]
+
 
     # with DP, each rank should process different prompts.
     # usually all the DP ranks process a full dataset,
@@ -131,18 +112,18 @@ def start(rank):
     # sampling params. here we set different max_tokens for different
     # ranks for demonstration.
     sampling_params = SamplingParams(
-        temperature=0.8, top_p=0.95, max_tokens=[16, 20][global_dp_rank % 2]
+        temperature=0.8, top_p=0.95, max_tokens=[20, 16][global_dp_rank % 2]
     )
 
+    # Fixed params
+    args.pop("tensor_parallel_size")
+    args.pop("enable_expert_parallel")
+
     # Create an LLM.
     llm = LLM(
-        model=model,
         tensor_parallel_size=GPUs_per_dp_rank,
-        enforce_eager=enforce_eager,
         enable_expert_parallel=True,
-        trust_remote_code=trust_remote_code,
-        max_num_seqs=max_num_seqs,
-        gpu_memory_utilization=gpu_memory_utilization,
+        **args,
     )
     outputs = llm.generate(prompts, sampling_params)
     # Print the outputs.
@@ -162,19 +143,22 @@ def start(rank):
 
 
 if __name__ == "__main__":
-    args = parse_args()
 
-    dp_size = args.dp_size
-    tp_size = args.tp_size
-    node_size = args.node_size
-    node_rank = args.node_rank
+    args = vars(parse_args())
+
+    dp_size = args.pop("dp_size")
+    tp_size = args.pop("tp_size")
+    node_size = args.pop("node_size")
+    node_rank = args.pop("node_rank")
 
     if node_size == 1:
         dp_master_ip = "127.0.0.1"
         dp_master_port = get_open_port()
+        args.pop("master_addr")
+        args.pop("master_port")
     else:
-        dp_master_ip = args.master_addr
-        dp_master_port = args.master_port
+        dp_master_ip = args.pop("master_addr")
+        dp_master_port = args.pop("master_port")
 
     assert dp_size % node_size == 0, "dp_size should be divisible by node_size"
     dp_per_node = dp_size // node_size
@@ -183,29 +167,22 @@ def start(rank):
 
     procs = []
     for local_dp_rank, global_dp_rank in enumerate(
-        range(node_rank * dp_per_node, (node_rank + 1) * dp_per_node)
-    ):
-        proc = Process(
-            target=main,
-            args=(
-                args.model,
-                dp_size,
-                local_dp_rank,
-                global_dp_rank,
-                dp_master_ip,
-                dp_master_port,
-                tp_size,
-                args.enforce_eager,
-                args.trust_remote_code,
-                args.max_num_seqs,
-                args.gpu_memory_utilization,
-            ),
-        )
+            range(node_rank * dp_per_node, (node_rank + 1) * dp_per_node)):
+        proc = Process(target=main,
+                       args=(
+                           args,
+                           dp_size,
+                           local_dp_rank,
+                           global_dp_rank,
+                           dp_master_ip,
+                           dp_master_port,
+                           tp_size,
+                       ))
         proc.start()
         procs.append(proc)
     exit_code = 0
     for proc in procs:
-        proc.join(timeout=300)
+        proc.join(timeout=1200)
         if proc.exitcode is None:
             print(f"Killing process {proc.pid} that didn't stop within 5 minutes.")
             proc.kill()

vllm/compilation/cuda_piecewise_backend.py

@@ -106,6 +106,7 @@ def check_for_ending_compilation(self):
             end_monitoring_torch_compile(self.vllm_config)
 
     def __call__(self, *args) -> Any:
+        # logger.info("CUDA BACKEND CALL")
         if not self.first_run_finished:
             self.first_run_finished = True
             self.check_for_ending_compilation()

vllm/compilation/decorators.py

@@ -157,6 +157,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = '', **kwargs):
             vllm_config.compilation_config.level in [
             CompilationLevel.NO_COMPILATION, CompilationLevel.DYNAMO_AS_IS
         ] or not supports_dynamo()
+        self.do_not_compile = True
         if self.do_not_compile:
             return
         compilation_counter.num_models_seen += 1
@@ -170,6 +171,7 @@ def __call__(self, *args, **kwargs):
         # e.g. TPU has the compilation logic in model runner, so we don't
         # need to compile the model inside.
         if self.do_not_compile or torch.compiler.is_compiling():
+            # logger.info("SKIPPING COMPILATION")
             return self.forward(*args, **kwargs)
 
         # the first compilation needs to have dynamic shapes marked

vllm/config.py

@@ -1816,6 +1816,18 @@ class ParallelConfig:
     disable_custom_all_reduce: bool = False
     """Disable the custom all-reduce kernel and fall back to NCCL."""
 
+    enable_microbatching: bool = False
+    """Enable microbatching for the model executor."""
+
+    always_microbatch_if_enabled: bool = True
+    """Always microbatch if microbatching is enabled. Easier to sync between
+       dp workers."""
+
+    microbatching_token_threshold: int = 4
+    """The threshold for microbatching. If the number of tokens in the
+    request is greater than this threshold, microbatching will be used.
+    Otherwise, the request will be processed in a single batch."""
+
     tokenizer_pool_config: Optional[TokenizerPoolConfig] = None
     """This parameter is deprecated and will be removed in a future release.
     Please remove it from your configs"""
@@ -4564,6 +4576,20 @@ def __post_init__(self):
                         "cascade attention. Disabling cascade attention.")
             self.model_config.disable_cascade_attn = True
 
+        if self.parallel_config.enable_microbatching and \
+            self.compilation_config.level >= CompilationLevel.PIECEWISE:
+            # Microbatching is not supported with piecewise compilation yet.
+            #  More specifically piecewise cuda-graphs
+            logger.warning_once(
+                "Piecewise compilation is not supported with "
+                "microbatching. Disabling piecewiseching compilation.")
+            self.compilation_config.level = CompilationLevel.NO_COMPILATION
+            if not self.model_config.enforce_eager:
+                self.compilation_config.full_cuda_graph = True
+                logger.warning_once(
+                    "Enabling fullcudagraphs for microbatching"
+                )
+
         disable_chunked_prefill_reasons: list[str] = []
 
         if self.model_config and self.model_config.pooler_config: