diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py index 027aa41fd..37733e290 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py @@ -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 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 6ff339f58..7675caa63 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -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), diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py index 869d50ae5..f430ecb16 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py @@ -219,6 +219,6 @@ class Qwen3OmniMoeCode2Wav(nn.Module): str(device), ) except Exception: - pass + logger.error("Error logging model load summary") return loaded diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index 1bfb0d3a3..4e8730eab 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -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(): diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 77b32c681..2d2b7ef8e 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -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: