diff --git a/examples/offline_inference/qwen_2_5_omni/run.sh b/examples/offline_inference/qwen_2_5_omni/run.sh index bee28e2c4..e27d21143 100644 --- a/examples/offline_inference/qwen_2_5_omni/run.sh +++ b/examples/offline_inference/qwen_2_5_omni/run.sh @@ -6,4 +6,4 @@ python end2end.py --model Qwen/Qwen2.5-Omni-7B \ --output-wav output_audio \ --prompt_type text \ --init-sleep-seconds 0 \ - --pt-prompts top100.pt + --pt-prompts .pt diff --git a/vllm_omni/__init__.py b/vllm_omni/__init__.py index 6860b462c..afeb09c1b 100644 --- a/vllm_omni/__init__.py +++ b/vllm_omni/__init__.py @@ -16,11 +16,11 @@ from . import patch # noqa: F401 from .config import OmniModelConfig # Main entry points -from .entrypoints.omni_llm import BaseOmniLLM +from .entrypoints.omni_llm import OmniLLM __all__ = [ # Main components - "BaseOmniLLM", + "OmniLLM", # "AsyncOmniLLM", # TODO: add AsyncOmniLLM back when it is implemented # Configuration "OmniModelConfig", diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 63dab9fde..c52e2e969 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -1,6 +1,7 @@ from typing import Any, Optional, Sequence, Union import logging import multiprocessing as mp +from concurrent.futures import ThreadPoolExecutor, as_completed import os import queue import sys @@ -43,74 +44,7 @@ from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) -class BaseOmniLLM: - def __init__( - self, model: str, stage_configs=None, log_stats: bool = False, **kwargs - ): - if stage_configs is None: - self.initialize_stage_configs(model) - else: - self.stage_configs = stage_configs - - self.stage_list = [] - self.initialize_stages(model) - - def initialize_stage_configs(self, model: str): - self.stage_configs = load_stage_configs_from_model(model) - - def initialize_stages(self, model: str): - for stage_config in self.stage_configs: - stage = OmniStage(stage_config) - stage_llm = OmniStageLLM(model=model, **stage_config.engine_args) - stage.set_engine(stage_llm) - self.stage_list.append(stage) - - def generate( - self, - prompts: Union[PromptType, Sequence[PromptType]], - sampling_params_list: Optional[ - Union[SamplingParams, Sequence[SamplingParams]] - ] = None, - ) -> list[OmniRequestOutput]: - """Generate text outputs for the given prompts.""" - final_outputs: list[OmniRequestOutput] = [] - if len(sampling_params_list) != len(self.stage_list): - raise ValueError( - f"Expected {len(self.stage_list)} sampling params, " - f"got {len(sampling_params_list)}" - ) - for stage_id, stage in enumerate(self.stage_list): - if stage_id > 0: - engine_inputs = stage.process_engine_inputs(self.stage_list, prompts) - else: - engine_inputs = prompts - engine_outputs = self._run_generation( - stage, sampling_params_list[stage_id], engine_inputs - ) - stage.set_engine_outputs(engine_outputs) - if hasattr(stage, "final_output") and stage.final_output: - final_outputs.append( - OmniRequestOutput( - stage_id=stage_id, - final_output_type=stage.final_output_type, - request_output=engine_outputs, - ) - ) - return final_outputs - - def _run_generation( - self, - stage: OmniStage, - sampling_params: SamplingParams, - prompts: Union[PromptType, Sequence[PromptType]], - ): - engine_outputs = [] - for ro in stage.engine.generate(prompts, sampling_params): - engine_outputs.append(ro) - return engine_outputs - - -class OmniLLM(BaseOmniLLM): +class OmniLLM: """Multi-process pipelined OmniLLM. - Per-stage process with copy-based IPC (Queues) @@ -135,16 +69,7 @@ class OmniLLM(BaseOmniLLM): else: self.stage_configs = stage_configs - self.stage_list: list[OmniStage] = [] - for stage_config in self.stage_configs: - # Only construct lightweight OmniStage without setting engine in parent - stage = OmniStage(stage_config) - self.stage_list.append(stage) - logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list)) - self._ctx = mp.get_context("spawn") - self._stage_in_queues: list[mp.Queue] = [] - self._stage_out_queues: list[mp.Queue] = [] # Optional file handler for orchestrator self._log_file = log_file if self._log_file: @@ -152,13 +77,34 @@ class OmniLLM(BaseOmniLLM): configure_orchestrator_logger(logger, self._log_file) self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file) + self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout) + def _initialize_stages(self, model: str, init_sleep_seconds: int, shm_threshold_bytes: int, init_timeout: int) -> None: + self.stage_list: list[OmniStage] = [] + # Build OmniStage instances in parallel, preserve original order + def _build_stage(idx_cfg: tuple[int, Any]) -> tuple[int, OmniStage]: + idx, cfg = idx_cfg + return idx, OmniStage(cfg) + + with ThreadPoolExecutor(max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))) as executor: + futures = [executor.submit(_build_stage, (idx, cfg)) for idx, cfg in enumerate(self.stage_configs)] + results: list[tuple[int, OmniStage]] = [] + for fut in as_completed(futures): + results.append(fut.result()) + results.sort(key=lambda x: x[0]) + self.stage_list = [st for _, st in results] + logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list)) + + self._ctx = mp.get_context("spawn") + self._stage_in_queues: list[mp.Queue] = [] + self._stage_out_queues: list[mp.Queue] = [] self._init_sleep_seconds = max(0, int(init_sleep_seconds)) self._shm_threshold_bytes = max(0, int(shm_threshold_bytes)) self._start_stage_processes(model) # Wait for all stages to report readiness before seeding self._stages_ready: set[int] = set() self._wait_for_stages_ready(timeout=init_timeout) + def _start_stage_processes(self, model: str) -> None: for stage_id, stage in enumerate(self.stage_list): @@ -400,7 +346,7 @@ class OmniLLM(BaseOmniLLM): # Provide actionable suggestions before shutdown try: suggestions = [ - "Verify GPU/device assignment in config (mrs.devices) is correct.", + "Verify GPU/device assignment in config (runtime.devices) is correct.", "Check GPU/host memory availability; reduce model or batch size if needed.", "Check model weights path and network reachability (if loading remotely).", "Increase initialization wait time (init_sleep_seconds or call-site timeout).", diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 06fb12041..6ec93878e 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -99,11 +99,11 @@ class OmniStage: ctx = ctx or mp.get_context("spawn") # Prepare lightweight dict config for worker engine_args = _to_dict(self.engine_args) - mrs_cfg = _to_dict(getattr(self.stage_config, "mrs", {})) + runtime_cfg = _to_dict(getattr(self.stage_config, "runtime", {})) stage_payload: Dict[str, Any] = { "stage_id": self.stage_id, "engine_args": engine_args, - "mrs": mrs_cfg, + "runtime": runtime_cfg, "shm_threshold_bytes": self._shm_threshold_bytes, } self._proc = ctx.Process( @@ -197,7 +197,7 @@ def _stage_worker( stage_id = stage_payload["stage_id"] engine_args = stage_payload.get("engine_args", {}) - mrs_cfg = stage_payload.get("mrs", {}) + runtime_cfg = stage_payload.get("runtime", {}) shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536)) # Per-stage file logger (optional) @@ -228,7 +228,7 @@ def _stage_worker( # Device mapping try: - set_stage_gpu_devices(stage_id, mrs_cfg.get("devices")) + set_stage_gpu_devices(stage_id, runtime_cfg.get("devices")) except Exception as e: _logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e) @@ -250,7 +250,7 @@ def _stage_worker( _logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id) break - max_batch_size = int(mrs_cfg.get("max_batch_size", 1) or 1) + max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1) batch_tasks: List[Dict[str, Any]] = [task] if max_batch_size > 1: while len(batch_tasks) < max_batch_size: diff --git a/vllm_omni/model_executor/models/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni.py index d7415b805..b8e3a2d9a 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni.py @@ -269,13 +269,13 @@ class Qwen2_5OmniForConditionalGeneration( # 2) Talker (if codec not provided) if self.model_stage == "talker": - # Mixed-mode support: in一个step内同时支持 Prefill*n 与 Decode*n。 - # 规则： - # - 使用特殊token将 Prefill 片段包裹为 [BOS][PAD...][EOS] - # - Decode 片段则为“单一非特殊token” - # - 若提供了 additional_information（可为按请求分割的 list 或拼接后的 tensor + 形状列表）， - # 则对 Prefill 片段按请求重建 thinker→talker 的输入嵌入； - # - 对 Decode 片段，若提供了 per-request 的 decode 辅助嵌入（可选），则进行叠加；否则保留原嵌入。 + # Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported. + # Rules: + # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS] + # - Decode segments consist of a single non-special token. + # - If additional_information is provided (can be a list split by request or a concatenated tensor plus a list of shapes), + # then for each request, reconstruct the thinker→talker input embeddings for the Prefill segments; + # - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), add them; otherwise, keep the original embedding. if input_ids is None and additional_information is None: input_ids = torch.zeros( diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml index 2f69b3080..620da3a89 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml @@ -1,7 +1,7 @@ # stage config for running qwen2.5-omni with architecture of OmniLLM. stage_args: - stage_id: 0 - mrs: + runtime: process: true # Run this stage in a separate process devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device) max_batch_size: 1 @@ -18,7 +18,7 @@ stage_args: final_output: true final_output_type: text - stage_id: 1 - mrs: + runtime: process: true devices: "1" max_batch_size: 3 @@ -35,7 +35,7 @@ stage_args: engine_input_source: [0] custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker - stage_id: 2 - mrs: + runtime: process: true devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU max_batch_size: 1 @@ -53,8 +53,8 @@ stage_args: final_output: true final_output_type: audio -# Top-level MRS config (concise): default windows and stage edges -mrs: +# Top-level runtime config (concise): default windows and stage edges +runtime: enabled: true defaults: window_size: -1 # Simplified: trigger downstream only after full upstream completion