diff --git a/docs/user_guide/examples/online_serving/qwen3_omni.md b/docs/user_guide/examples/online_serving/qwen3_omni.md index d4e13fc89..4e01ef4c2 100644 --- a/docs/user_guide/examples/online_serving/qwen3_omni.md +++ b/docs/user_guide/examples/online_serving/qwen3_omni.md @@ -22,6 +22,10 @@ both stages on vLLM's Triton unquantized MoE path and avoids the unstable FlashI some GPUs. Asynchronous chunk streaming is **enabled by default** within the bundled configuration. +For Thinker pipeline parallelism, configure `pipeline_parallel_size` and `devices` +on the Thinker stage and set `async_scheduling: false`. Keep Talker and Code2Wav +at `pipeline_parallel_size: 1`. Audio chunk streaming can remain enabled. + To explicitly utilize a custom deployment YAML, specify the configuration path: ```bash vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \ diff --git a/tests/model_executor/models/qwen3_omni/fixtures/thinker_pp2.yaml b/tests/model_executor/models/qwen3_omni/fixtures/thinker_pp2.yaml new file mode 100644 index 000000000..56720b49c --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/fixtures/thinker_pp2.yaml @@ -0,0 +1,66 @@ +async_chunk: true +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + initial_codec_chunk_frames: 4 + codec_chunk_frames: 25 + codec_left_context_frames: 25 +stages: +- stage_id: 0 + max_num_batched_tokens: 4096 + max_num_seqs: 2 + gpu_memory_utilization: 0.65 + moe_backend: triton + trust_remote_code: true + enable_prefix_caching: false + devices: 0,1 + default_sampling_params: + temperature: 0.0 + max_tokens: 128 + enforce_eager: true + max_model_len: 4096 + async_scheduling: false + tensor_parallel_size: 1 + pipeline_parallel_size: 2 +- stage_id: 1 + max_num_batched_tokens: 4096 + max_num_seqs: 2 + gpu_memory_utilization: 0.2 + moe_backend: triton + trust_remote_code: true + enable_prefix_caching: false + devices: '1' + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_k: 50 + max_tokens: 1024 + repetition_penalty: 1.05 + enforce_eager: true + max_model_len: 4096 + async_scheduling: false + tensor_parallel_size: 1 + pipeline_parallel_size: 1 +- stage_id: 2 + max_num_batched_tokens: 4096 + max_num_seqs: 2 + gpu_memory_utilization: 0.05 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + enable_chunked_prefill: false + async_scheduling: false + devices: '1' + input_connectors: + from_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + repetition_penalty: 1.1 + max_model_len: 4096 + tensor_parallel_size: 1 + pipeline_parallel_size: 1 diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py new file mode 100644 index 000000000..de8d02e3e --- /dev/null +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_pp_capture.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Required Thinker snapshots must survive consecutive pipeline ranks. + +Actual forward methods with deterministic decoder doubles test the contract +without model weights. Distributed serving remains a separate integration gate. +""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from vllm.model_executor.models.utils import make_empty_intermediate_tensors_factory +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.qwen3_omni import qwen3_omni as combined_module +from vllm_omni.model_executor.models.qwen3_omni import qwen3_omni_moe_thinker as thinker_module + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.parallel] + + +class _Layer(nn.Module): + def forward(self, positions, hidden_states, residual): + total = hidden_states if residual is None else hidden_states + residual + return total * 0.01 + 0.125, total + + +class _Norm(nn.Module): + def forward(self, hidden_states, residual): + return hidden_states + residual, None + + +def _forward(monkeypatch, start, end, embeddings, intermediate=None, *, capture=True): + rank = SimpleNamespace(is_first_rank=start == 0, is_last_rank=end == 48) + monkeypatch.setattr(thinker_module, "get_pp_group", lambda: rank) + model = object.__new__(thinker_module.Qwen3MoeLLMModel) + nn.Module.__init__(model) + model.start_layer = start + model.end_layer = end + model.layers = nn.ModuleList([_Layer() for _ in range(48)]) + model.norm = _Norm() + return model.forward( + input_ids=None, + positions=torch.arange(embeddings.shape[0]), + inputs_embeds=embeddings, + intermediate_tensors=intermediate, + capture_layer_indices=[0, 24] if capture else None, + return_hidden_states=capture, + ) + + +@pytest.mark.parametrize("tokens", [1, 4]) +@pytest.mark.parametrize("boundaries", [(0, 12, 48), (0, 24, 48), (0, 36, 48), (0, 16, 32, 48)]) +def test_required_captures_match_unsplit_forward(monkeypatch, tokens, boundaries): + embeddings = torch.arange(tokens * 8, dtype=torch.float32).reshape(tokens, 8) / 8 + reference, captures = _forward(monkeypatch, 0, 48, embeddings) + intermediate = None + for start, end in zip(boundaries, boundaries[1:]): + result = _forward(monkeypatch, start, end, embeddings, intermediate) + if end < 48: + assert isinstance(result, IntermediateTensors) + assert all(isinstance(value, torch.Tensor) for value in result.tensors.values()) + intermediate = result + actual, actual_captures = result + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + actual_layers = actual_captures["hidden_states"]["layers"] + assert set(actual_layers) == {0, 24} + for key in (0, 24): + torch.testing.assert_close(actual_layers[key], captures["hidden_states"]["layers"][key], rtol=0, atol=0) + # A later receive into the rank's persistent buffers cannot rewrite snapshots. + assert intermediate is not None + for value in intermediate.tensors.values(): + value.zero_() + for key in (0, 24): + torch.testing.assert_close(actual_layers[key], captures["hidden_states"]["layers"][key], rtol=0, atol=0) + + +def test_capture_disabled_keeps_original_intermediate_schema(monkeypatch): + embeddings = torch.ones(4, 8) + first = _forward(monkeypatch, 0, 24, embeddings, capture=False) + assert set(first.tensors) == {"hidden_states", "residual"} + actual, captures = _forward(monkeypatch, 24, 48, embeddings, first, capture=False) + reference, _ = _forward(monkeypatch, 0, 48, embeddings, capture=False) + assert captures is None + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + + +@pytest.mark.parametrize( + "start,expected", [(0, set()), (16, {"capture_0"}), (24, {"capture_0"}), (32, {"capture_0", "capture_24"})] +) +@pytest.mark.parametrize("staged", [False, True]) +def test_factory_declares_only_captures_from_earlier_ranks(monkeypatch, start, expected, staged): + thinker_config = SimpleNamespace(text_config=SimpleNamespace(hidden_size=8)) + config = SimpleNamespace( + thinker_config=thinker_config, + talker_config=SimpleNamespace(accept_hidden_layer=24), + code2wav_config=SimpleNamespace(), + tts_bos_token_id=1, + tts_eos_token_id=2, + tts_pad_token_id=3, + ) + model_config = SimpleNamespace(hf_config=config, multimodal_config=None) + if staged: + model_config.model_stage = "thinker" + vllm_config = SimpleNamespace(model_config=model_config, with_hf_config=lambda *args, **kwargs: None) + thinker = nn.Module() + thinker.weight = nn.Parameter(torch.ones(1)) + thinker.language_model = SimpleNamespace(model=SimpleNamespace(start_layer=start)) + thinker.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(["hidden_states", "residual"], 8) + monkeypatch.setattr(combined_module, "init_vllm_registered_model", lambda **kwargs: thinker) + model = combined_module.Qwen3OmniMoeForConditionalGeneration(vllm_config=vllm_config) + buffers = model.make_empty_intermediate_tensors(4, dtype=torch.float32, device="cpu") + assert set(buffers.tensors) == {"hidden_states", "residual"} | (expected if staged else set()) + assert all(value.shape == (4, 8) for value in buffers.tensors.values()) diff --git a/tests/worker/test_pp_intermediate_output.py b/tests/worker/test_pp_intermediate_output.py new file mode 100644 index 000000000..a9cf02bba --- /dev/null +++ b/tests/worker/test_pp_intermediate_output.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Non-final pipeline outputs must reach vLLM's existing PP sender unchanged.""" + +from types import SimpleNamespace + +import pytest +import torch +from vllm.sequence import IntermediateTensors +from vllm.v1.worker.gpu_model_runner import GPUModelRunner + +from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.parallel] + + +def test_nonfinal_forward_bypasses_final_output_conversion(monkeypatch): + intermediate = IntermediateTensors({"hidden_states": torch.ones(4, 8), "residual": torch.ones(4, 8)}) + monkeypatch.setattr(GPUModelRunner, "_model_forward", lambda *args, **kwargs: intermediate) + + def final_only_converter(*args, **kwargs): + pytest.fail("A non-final rank must not convert intermediate tensors to OmniOutput") + + runner = object.__new__(OmniGPUModelRunner) + runner.model = SimpleNamespace(make_omni_output=final_only_converter, have_multimodal_outputs=True) + runner._build_model_kwargs_extra = lambda: {} + result = runner._model_forward() + assert result is intermediate + hidden, multimodal = runner.extract_multimodal_outputs(result) + assert hidden is intermediate + assert multimodal == {} + + +def test_final_tensor_still_uses_model_converter(monkeypatch): + tensor = torch.ones(4, 8) + monkeypatch.setattr(GPUModelRunner, "_model_forward", lambda *args, **kwargs: tensor) + converted = object() + calls = [] + + def convert(value, **kwargs): + calls.append(value) + return converted + + runner = object.__new__(OmniGPUModelRunner) + runner.model = SimpleNamespace(make_omni_output=convert) + runner._build_model_kwargs_extra = lambda: {} + assert runner._model_forward() is converted + assert calls == [tensor] diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index a5856718d..28accbc51 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -34,6 +34,7 @@ from vllm.model_executor.models.qwen3_omni_moe_thinker import ( from vllm.model_executor.models.utils import ( WeightsMapper, init_vllm_registered_model, + make_empty_intermediate_tensors_factory, maybe_prefix, ) from vllm.multimodal import MULTIMODAL_REGISTRY @@ -53,6 +54,7 @@ from vllm_omni.model_executor.models.qwen3_omni.quantization import ( apply_outer_quant_config_mapping, ) from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( + PP_CAPTURE_PREFIX, Qwen3OmniMoeThinkerDummyInputsBuilder, Qwen3OmniMoeThinkerForConditionalGeneration, Qwen3OmniMoeThinkerMultiModalProcessor, @@ -274,6 +276,18 @@ class Qwen3OmniMoeForConditionalGeneration( self.make_empty_intermediate_tensors = ( self.thinker.make_empty_intermediate_tensors if self.model_stage == "thinker" else lambda: None ) + if self.model_stage == "thinker": + capture_indices = self._thinker_capture_layer_indices() + start_layer = self.thinker.language_model.model.start_layer + incoming_captures = [f"{PP_CAPTURE_PREFIX}{index}" for index in capture_indices if index < start_layer] + if incoming_captures: + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual", *incoming_captures], thinker_config.text_config.hidden_size + ) + + def _thinker_capture_layer_indices(self) -> list[int]: + accept_layer = getattr(self.talker_config, "accept_hidden_layer", None) + return [0, int(accept_layer)] if self.is_staged_run and accept_layer is not None else [] @classmethod async def buffer_realtime_audio( @@ -441,11 +455,11 @@ class Qwen3OmniMoeForConditionalGeneration( # Only staged runs have a talker stage to consume the capture; in a # plain vLLM run capturing would waste a clone per layer per step # and make the thinker return a tuple stock vLLM cannot handle. - accept_layer = getattr(self.talker_config, "accept_hidden_layer", None) + capture_indices = self._thinker_capture_layer_indices() capture_kwargs = {} - if accept_layer is not None and self.is_staged_run: + if capture_indices: capture_kwargs = { - "capture_layer_indices": [0, int(accept_layer)], + "capture_layer_indices": capture_indices, "return_hidden_states": True, } 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 b2a7d7abe..8d0389d63 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 @@ -118,6 +118,7 @@ from vllm.multimodal.utils import set_mm_embedding_modality from vllm.sequence import IntermediateTensors from vllm.transformers_utils.processor import cached_processor_from_config +from vllm_omni.data_entry_keys import OmniPayload from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( Qwen2_5OmniConditionalGenerationMixin, Qwen2_5OmniThinkerMultiModalDataParser, @@ -140,6 +141,8 @@ except (ImportError, ModuleNotFoundError): logger = init_logger(__name__) +PP_CAPTURE_PREFIX = "capture_" + _THINKER_ARCHITECTURE = "Qwen3OmniMoeThinkerForConditionalGeneration" @@ -568,7 +571,15 @@ class Qwen3MoeLLMModel(_Qwen3MoeLLMModel): hidden_states = intermediate_tensors["hidden_states"] residual = intermediate_tensors["residual"] capture_set = set(capture_layer_indices) if capture_layer_indices else None - captured_hidden_states: dict[str, torch.Tensor] | None = {} if return_hidden_states else None + captured_hidden_states: OmniPayload | None = {} if return_hidden_states else None + + if captured_hidden_states is not None and capture_set and intermediate_tensors is not None: + for layer_idx in capture_set: + if layer_idx < self.start_layer: + hs = captured_hidden_states.setdefault("hidden_states", {}) + layers = hs.setdefault("layers", {}) + # Receive buffers are reused on the next step; retain an independent snapshot. + layers[layer_idx] = intermediate_tensors[f"{PP_CAPTURE_PREFIX}{layer_idx}"].clone() for layer_idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]): layer_idx = layer_idx + self.start_layer @@ -589,7 +600,15 @@ class Qwen3MoeLLMModel(_Qwen3MoeLLMModel): hidden_states = hidden_states + deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"] if not get_pp_group().is_last_rank: - return IntermediateTensors({"hidden_states": hidden_states, "residual": residual}) + tensors = {"hidden_states": hidden_states, "residual": residual} + if captured_hidden_states: + tensors.update( + { + f"{PP_CAPTURE_PREFIX}{index}": value + for index, value in captured_hidden_states["hidden_states"]["layers"].items() + } + ) + return IntermediateTensors(tensors) hidden_states, _ = self.norm(hidden_states, residual) if captured_hidden_states is not None: return hidden_states, captured_hidden_states diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 8dd9f5f58..d3ee8de02 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -34,6 +34,7 @@ from vllm.v1.worker.gpu_model_runner import GPUModelRunner, PerLayerAttnMetadata from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm_omni.core.prefix_cache import OmniTensorPrefixCache +from vllm_omni.data_entry_keys import OmniPayload from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEmbedding as MRotaryEmbedding from vllm_omni.model_executor.models.model_local_kv import collect_model_local_kv_specs @@ -892,8 +893,8 @@ class OmniGPUModelRunner(GPUModelRunner): @torch.inference_mode() def extract_multimodal_outputs( - self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput - ) -> tuple[Any, Any]: + self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput | IntermediateTensors + ) -> tuple[torch.Tensor | IntermediateTensors, OmniPayload | None]: if ( hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs @@ -902,7 +903,7 @@ class OmniGPUModelRunner(GPUModelRunner): text_hidden_states = hidden_states.text_hidden_states multimodal_outputs = hidden_states.multimodal_outputs - elif isinstance(hidden_states, torch.Tensor): + elif isinstance(hidden_states, (torch.Tensor, IntermediateTensors)): text_hidden_states = hidden_states multimodal_outputs = {} elif isinstance(hidden_states, list) or isinstance(hidden_states, tuple): @@ -1225,6 +1226,9 @@ class OmniGPUModelRunner(GPUModelRunner): else: hidden_states = outputs hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) + if isinstance(hidden_states, IntermediateTensors): + # Profiling/graph capture needs a tensor to select dummy sampler rows. + hidden_states = hidden_states["hidden_states"] if self.speculative_config and ( self.speculative_config.use_eagle() or self.speculative_config.uses_draft_model() @@ -2072,7 +2076,7 @@ class OmniGPUModelRunner(GPUModelRunner): **model_kwargs, **model_kwargs_extra, ) - if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"): + if not isinstance(model_output, (OmniOutput, IntermediateTensors)) and hasattr(self.model, "make_omni_output"): model_output = self.model.make_omni_output(model_output, **model_kwargs, **model_kwargs_extra) # Cache model output so later sample_tokens can consume multimodal results. self._omni_last_model_output = model_output