From 32fd084ae25afa68a8f5e53dc79502a631674839 Mon Sep 17 00:00:00 2001 From: Baoyuan Qi Date: Fri, 6 Mar 2026 14:45:08 +0800 Subject: [PATCH] [MiMo-Audio] Bugfix tp lg than 1 (#1688) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: 齐保元 --- vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py | 3 +++ vllm_omni/model_executor/stage_configs/mimo_audio.yaml | 2 ++ .../model_executor/stage_configs/mimo_audio_async_chunk.yaml | 2 ++ 3 files changed, 7 insertions(+) diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index 3b9f94f23..c1b71e076 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py @@ -556,12 +556,14 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor config.hidden_size, bias=False, return_bias=False, + gather_output=True, ) self.hidden_states_downcast = ColumnParallelLinear( config.hidden_size, self.local_config.hidden_size, bias=False, return_bias=False, + gather_output=True, ) self.lm_head = ColumnParallelLinear( @@ -569,6 +571,7 @@ class MiMoAudioLLMForConditionalGeneration(nn.Module, SupportsMultiModal, Suppor config.vocab_size, bias=False, return_bias=False, + gather_output=True, ) # Re-encode the sum of multi-layer RVQ embeddings to obtain true Audio Code Embeddings diff --git a/vllm_omni/model_executor/stage_configs/mimo_audio.yaml b/vllm_omni/model_executor/stage_configs/mimo_audio.yaml index 552ba2ab5..b824fcb41 100644 --- a/vllm_omni/model_executor/stage_configs/mimo_audio.yaml +++ b/vllm_omni/model_executor/stage_configs/mimo_audio.yaml @@ -15,6 +15,7 @@ stage_args: model_arch: MiMoAudioForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + tensor_parallel_size: 1 # Change to desired TP size for multi-GPU inference (e.g., 4 for 4 GPUs) gpu_memory_utilization: 0.3 enforce_eager: true # need to discuss trust_remote_code: true @@ -45,6 +46,7 @@ stage_args: model_arch: MiMoAudioForConditionalGeneration worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + tensor_parallel_size: 1 # Change to desired TP size for multi-GPU inference (e.g., 4 for 4 GPUs) gpu_memory_utilization: 0.2 enforce_eager: true trust_remote_code: true diff --git a/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml index a60fc566a..7177aa809 100644 --- a/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml +++ b/vllm_omni/model_executor/stage_configs/mimo_audio_async_chunk.yaml @@ -15,6 +15,7 @@ stage_args: model_arch: MiMoAudioForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + tensor_parallel_size: 1 # Change to desired TP size for multi-GPU inference (e.g., 4 for 4 GPUs) gpu_memory_utilization: 0.3 enforce_eager: true # need to discuss trust_remote_code: true @@ -46,6 +47,7 @@ stage_args: model_arch: MiMoAudioForConditionalGeneration worker_type: generation scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + tensor_parallel_size: 1 # Change to desired TP size for multi-GPU inference (e.g., 4 for 4 GPUs) gpu_memory_utilization: 0.2 enforce_eager: true trust_remote_code: true