anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Debug] Clean code in Qwen 3 Omni and add warning for talker temperature. (#688)

Browse Source 
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: root <root@hk01dgx028.cm.cluster>
This commit is contained in:
Zhou Taichangandroot authored and GitHub committed 2026-01-12 23:13:45 +13:00
1 parent 8c7f704c23
commit 2e13bddf87
5 files changed
+86 -91
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py
+1 -1
View File
@@ -771,7 +771,7 @@ class Qwen2_5OmniForConditionalGeneration(
) # for decode
return output_token_ids, processed_output_token_embeds
def compute_logits(self, hidden_states: torch.Tensor | OmniOutput) -> torch.Tensor | None:
def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, **kwargs: object) -> torch.Tensor | None:
# Handle OmniOutput type
if isinstance(hidden_states, OmniOutput):
hidden_states = hidden_states.text_hidden_states
vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py
+35 -72
View File
@@ -8,13 +8,6 @@ from functools import cached_property
import torch
import torch.nn as nn
from transformers.generation.logits_process import (
LogitsProcessorList,
RepetitionPenaltyLogitsProcessor,
SuppressTokensLogitsProcessor,
TemperatureLogitsWarper,
TopKLogitsWarper,
)
from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
Qwen3OmniMoeCode2WavConfig,
Qwen3OmniMoeConfig,
@@ -132,6 +125,11 @@ class Qwen3OmniMoeForConditionalGeneration(
self.model = self.thinker
self.talker = None
self.code2wav = None
self.tts_tokens = torch.tensor(
[[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]],
device=self._module_device(self.thinker),
dtype=torch.long,
)
elif self.model_stage == "talker":
self.has_preprocess = True
self.has_postprocess = True
@@ -154,15 +152,6 @@ class Qwen3OmniMoeForConditionalGeneration(
self.model = self.talker
self.code2wav = None
self.logits_processors = LogitsProcessorList(
[
RepetitionPenaltyLogitsProcessor(penalty=1.05),
TemperatureLogitsWarper(temperature=0.9),
SuppressTokensLogitsProcessor(suppress_tokens=self._get_talker_suppressed_tokens()),
TopKLogitsWarper(top_k=50),
]
)
# for CI: Initialize special tokens embeddings early to avoid AttributeError when loading dummy weights
self._init_special_tokens_embeddings()
@@ -188,12 +177,6 @@ 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":
self.tts_tokens = torch.tensor(
[[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]],
device=self._module_device(self.thinker),
dtype=torch.long,
)
# ==================== Device utilities ====================
@@ -208,31 +191,6 @@ class Qwen3OmniMoeForConditionalGeneration(
return buf.device
return torch.device("cpu")
def move_submodules_to_devices(
self,
*,
thinker_device: str | torch.device | None = None,
talker_device: str | torch.device | None = None,
code_predictor_device: str | torch.device | None = None,
code2wav_device: str | torch.device | None = None,
) -> None:
"""
Optionally move thinker/talker/code2wav to different devices.
Example:
model.move_submodules_to_devices(
thinker_device='cuda:0',
talker_device='cuda:1',
code2wav_device='cuda:2',
)
"""
if thinker_device is not None and self.thinker is not None:
self.thinker.to(thinker_device)
if talker_device is not None and self.talker is not None:
self.talker.to(talker_device)
if code2wav_device is not None and self.code2wav is not None:
self.code2wav.to(code2wav_device)
@cached_property
def sampler(self):
"""Get sampler from active model."""
@@ -372,6 +330,7 @@ class Qwen3OmniMoeForConditionalGeneration(
# ========== Stage 2.1: Talker ==========
elif self.model_stage == "talker":
if input_ids is None:
# special case for profile run
input_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device)
# Ensure we have base embeddings when only ids are provided
@@ -1017,6 +976,29 @@ class Qwen3OmniMoeForConditionalGeneration(
# ==================== Logits and Sampling ====================
def _warn_talker_sampling_temperature(self, sampling_metadata: SamplingMetadata):
warning_parts = []
if sampling_metadata.temperature is None:
warning_parts.append(
"Temperature is set to None, as all requests are greedy. "
"This is equivalent to setting temperature to 0.0."
"Please consider setting a higher temperature i.e. 0.4."
)
else:
warning_parts.append(
"Temperature is set to: "
f"{sampling_metadata.temperature}, where temperature as 0.0 may "
"cause repetitive output. Please consider setting a higher "
"temperature i.e. 0.4."
)
warning_parts.append(
"This warning will be shown only once, for the first request where "
"temperature is 0.0. Later requests will not show this warning but "
"still be affected by the temperature."
)
warning_info = "\n".join(warning_parts)
logger.warning_once(warning_info)
def compute_logits(
self,
hidden_states: torch.Tensor | OmniOutput,
@@ -1024,35 +1006,16 @@ class Qwen3OmniMoeForConditionalGeneration(
) -> torch.Tensor | None:
"""Compute logits from hidden states."""
# Handle OmniOutput type
from vllm.v1.sample.logits_processor import LogitsProcessors
if isinstance(hidden_states, OmniOutput):
hidden_states = hidden_states.text_hidden_states
num_reqs = hidden_states.shape[0]
top_k = hidden_states.size(1) - 1
def dummy_tensors(v):
return torch.full((num_reqs,), v, device=hidden_states.device)
if (
getattr(self, "model_stage", None) == "talker"
and sampling_metadata is not None
and (sampling_metadata.temperature is None or (sampling_metadata.temperature <= 0).any())
):
self._warn_talker_sampling_temperature(sampling_metadata)
if sampling_metadata is None:
sampling_metadata = SamplingMetadata(
temperature=dummy_tensors(0.5),
all_greedy=False,
all_random=False,
top_p=dummy_tensors(0.9),
top_k=dummy_tensors(top_k),
generators={},
max_num_logprobs=None,
no_penalties=True,
prompt_token_ids=None,
frequency_penalties=dummy_tensors(0.1),
presence_penalties=dummy_tensors(0.1),
repetition_penalties=dummy_tensors(0.1),
output_token_ids=[[] for _ in range(num_reqs)],
allowed_token_ids_mask=None,
bad_words_token_ids={},
logitsprocs=LogitsProcessors(),
)
# Use active model for logits computation
logits = self.model.compute_logits(hidden_states) # V, d
# Talker: suppress tokens by setting their probability to ~1e-9 (finite very small),
vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py
+1 -1
View File
@@ -219,6 +219,6 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
str(device),
)
except Exception:
pass
logger.error("Error logging model load summary")
return loaded
vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py
+43 -15
View File
@@ -24,7 +24,6 @@ from vllm.model_executor.layers.linear import ReplicatedLinear
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.models.interfaces import (
MultiModalEmbeddings,
SupportsMultiModal,
SupportsPP,
)
from vllm.model_executor.models.qwen2_5_omni_thinker import (
@@ -68,7 +67,7 @@ Qwen3OmniMoeThinkerDummyInputsBuilder = Qwen2_5OmniThinkerDummyInputsBuilder
)
class Qwen3OmniMoeTalkerForConditionalGeneration(
nn.Module,
SupportsMultiModal,
# SupportsMultiModal,
SupportsPP,
Qwen3OmniMoeConditionalGenerationMixin,
):
@@ -430,9 +429,16 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
if not mm_input_by_modality:
return []
logger.warning(
"\n\n\n"
"THIS FUNCTION RETURNS DUMMY MULTIMODAL EMBEDDINGS FOR PROFILE RUN, "
"SHOULD NOT BE CALLED IN INFERENCE."
"\n\n\n"
)
# The result multimodal_embeddings is tuple of tensors, with each
# tensor correspoending to a multimodal data item (image or video).
multimodal_embeddings: tuple[torch.Tensor, ...] = ()
dummy_multimodal_embeddings: tuple[torch.Tensor, ...] = ()
# NOTE: It is important to iterate over the keys in this dictionary
# to preserve the order of the modalities.
@@ -441,22 +447,44 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
multimodal_input = mm_input_by_modality[modality]
if modality == "image":
image_embeddings = self._process_image_input(multimodal_input)
image_embeddings = self.hidden_projection(image_embeddings)
multimodal_embeddings += tuple(image_embeddings)
dummy_image_embeddings = ()
for image_embed in image_embeddings:
dummy_image_embeddings += (
torch.zeros(
image_embed.shape[0],
self.config.text_config.hidden_size,
device=image_embed.device,
dtype=torch.bfloat16,
),
)
dummy_multimodal_embeddings += tuple(image_embeddings)
if modality == "video":
video_embeddings = self._process_video_input(multimodal_input)
video_video_embeddings_project = ()
dummy_video_video_embeddings = ()
for video_embed in video_embeddings:
proj = nn.Linear(8192, 2048).to(device=video_embed.device, dtype=torch.bfloat16)
video_embed = proj(video_embed)
video_embed_project = self.hidden_projection(video_embed)
video_video_embeddings_project += (video_embed_project,)
multimodal_embeddings += tuple(video_video_embeddings_project)
dummy_video_video_embeddings += (
torch.zeros(
video_embed.shape[0],
self.config.text_config.hidden_size,
device=video_embed.device,
dtype=torch.bfloat16,
),
)
dummy_multimodal_embeddings += tuple(dummy_video_video_embeddings)
if modality == "audio":
audio_embeddings = self._process_audio_input(multimodal_input)
audio_embeddings = self.hidden_projection(audio_embeddings)
multimodal_embeddings += tuple(audio_embeddings)
return multimodal_embeddings
dummy_audio_embeddings = ()
for audio_embed in audio_embeddings:
dummy_audio_embeddings += (
torch.zeros(
audio_embed.shape[0],
self.config.text_config.hidden_size,
device=audio_embed.device,
dtype=torch.bfloat16,
),
)
dummy_multimodal_embeddings += tuple(dummy_audio_embeddings)
return dummy_multimodal_embeddings
def embed_input_ids(
self,
@@ -497,7 +525,7 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
str(device),
)
except Exception:
pass
logger.error("Error logging model load summary")
multi_model_weights = set()
for name, param in self.visual.named_parameters():
vllm_omni/worker/gpu_ar_model_runner.py
+6 -2
View File
@@ -223,7 +223,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
return output
sample_hidden_states = hidden_states[logits_indices]
logits = self.model.compute_logits(sample_hidden_states)
logits = self.model.compute_logits(
sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
)
else:
assert not self.is_pooling_model
@@ -239,7 +241,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
logits = None
else:
sample_hidden_states = hidden_states[logits_indices]
logits = self.model.compute_logits(sample_hidden_states)
logits = self.model.compute_logits(
sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata
)
model_output_broadcast_data: dict[str, Any] = {}
if logits is not None: