[Multimodal] Automatically fallback to ViT DP when TP is unavailable (#49046)
Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Roger Wang <hey@rogerw.io>
Isotr0py committed
Jul 19, 2026 at 05:41 UTC
9243e0124e5c5396213258b2fb6d6401ee8965e9
3 files changed
+17
-5
vllm/model_executor/models/kimi_k25.py
+3
-3
@@ -34,6 +34,7 @@ from vllm.model_executor.models.kimi_k25_vit import (
34
MoonViT3dPretrainedModel,
35
vision_tower_forward,
36
)
37
+from vllm.model_executor.models.vision import is_vit_use_data_parallel
38
from vllm.multimodal import MULTIMODAL_REGISTRY
39
from vllm.multimodal.inputs import (
40
MultiModalFieldConfig,
@@ -336,9 +337,8 @@ class KimiK25ForConditionalGeneration(
337
self.config = config
338
quant_config = vllm_config.quant_config
339
339
- # Check for MoonViT config compatibility
340
- self.use_data_parallel = (
341
- model_config.multimodal_config.mm_encoder_tp_mode == "data"
340
+ self.use_data_parallel = is_vit_use_data_parallel(
341
+ config.vision_config.num_attention_heads
342
)
343
self.hidden_size = config.text_config.hidden_size
344
self.device = current_platform.current_device()
vllm/model_executor/models/kimi_k25_vit.py
+1
-1
@@ -357,7 +357,7 @@ class MoonViTEncoderLayer(nn.Module):
357
attn_bias: bool = False,
358
):
359
super().__init__()
360
- self.use_data_parallel = is_vit_use_data_parallel()
360
+ self.use_data_parallel = is_vit_use_data_parallel(num_heads)
361
362
self.num_heads = num_heads
363
self.hidden_dim = hidden_dim
vllm/model_executor/models/vision.py
+13
-1
@@ -139,11 +139,23 @@ def get_fp8_padded_hidden_size(num_heads: int, head_dim: int) -> int | None:
139
return num_heads * round_up(head_dim, 16)
140
141
142
-def is_vit_use_data_parallel():
142
+def is_vit_use_data_parallel(num_heads: int | None = None) -> bool:
143
"""
144
Get the tensor parallel type for Vision Transformer.
145
"""
146
mm_cfg = get_multimodal_config()
147
+ can_split = (
148
+ num_heads % get_tensor_model_parallel_world_size() == 0
149
+ if num_heads is not None
150
+ else None
151
+ )
152
+ if num_heads is not None and not can_split:
153
+ logger.warning_once(
154
+ "The number of vision attention heads is not divisible by "
155
+ "the tensor parallel size. Falling back to data parallelism "
156
+ "for the vision encoder."
157
+ )
158
+ return True
159
return mm_cfg is not None and mm_cfg.mm_encoder_tp_mode == "data"
160
161