From 57e5f0fe69531e0a0bb86e77281e282a68b1ddfb Mon Sep 17 00:00:00 2001 From: Qing An <82441879+qinganrice@users.noreply.github.com> Date: Sun, 31 May 2026 09:10:14 -0700 Subject: [PATCH] [model, omni] feat: Qwen3-Omni Thinker LoRA for RL training (#3915) Signed-off-by: Qing An Co-authored-by: Samit <285365963@qq.com> --- vllm_omni/entrypoints/async_omni.py | 19 +++++++++++++++++-- .../qwen3_omni/qwen3_omni_moe_thinker.py | 11 ++++++++++- 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index c07bb2085..8a4a1d48e 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -295,6 +295,15 @@ class AsyncOmni(EngineClient, OmniBase): # Start final output dispatcher on the first call to generate() self._final_output_handler() + # Forward bare sampling_params (e.g. from /v1/completions) as the stage-0 entry. + if sampling_params_list is None and sampling_params is not None: + if self.num_stages == 1: + sampling_params_list = [sampling_params] + else: + default = list(self.default_sampling_params_list) + default[0] = sampling_params + sampling_params_list = default + # Expand sampling params for PD disaggregation (user may provide N-1 params) if ( sampling_params_list is not None @@ -919,8 +928,14 @@ class AsyncOmni(EngineClient, OmniBase): for result in results: if isinstance(result, dict) and result.get("todo"): continue - if isinstance(result, list): - merged.update(result) + if isinstance(result, (list, set)): + for item in result: + if isinstance(item, (list, set)): + merged.update(item) + elif isinstance(item, int): + merged.add(item) + elif isinstance(result, int): + merged.add(result) return sorted(merged) async def pin_lora(self, adapter_id: int) -> bool: diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index df26a1faf..17be417f3 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -64,6 +64,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, + SupportsLoRA, SupportsMRoPE, SupportsMultiModal, SupportsPP, @@ -1071,6 +1072,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( nn.Module, SupportsMultiModal, SupportsPP, + SupportsLoRA, SupportsMRoPE, Qwen3OmniMoeConditionalGenerationMixin, SupportsTranscription, @@ -1089,6 +1091,11 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( "k_proj", "v_proj", ], + "attn.qkv": [ + "attn.q", + "attn.k", + "attn.v", + ], "gate_up_proj": [ "gate_proj", "up_proj", @@ -1111,7 +1118,8 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.vllm_config = vllm_config # needed for torch compile forward context - thinker_config: Qwen3OmniMoeThinkerConfig = vllm_config.model_config.hf_config + hf_config = vllm_config.model_config.hf_config + thinker_config: Qwen3OmniMoeThinkerConfig = getattr(hf_config, "thinker_config", None) or hf_config quant_config = vllm_config.quant_config multimodal_config = vllm_config.model_config.multimodal_config self.config = thinker_config @@ -1660,6 +1668,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( self, input_tokens: list[int], mm_features: list[MultiModalFeatureSpec], + **kwargs, ) -> tuple[torch.Tensor, int]: """Compute M-RoPE input positions using mm_features directly.""" seq_len = len(input_tokens)