[MiMo-Audio] Bugfix tp lg than 1 (#1688)

Signed-off-by: 齐保元 <qibaoyuan@xiaomi.com>

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

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

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