diff --git a/vllm/model_executor/models/kimi_k25.py b/vllm/model_executor/models/kimi_k25.py index e2b74744b4..79b6e272bb 100644 --- a/vllm/model_executor/models/kimi_k25.py +++ b/vllm/model_executor/models/kimi_k25.py @@ -34,6 +34,7 @@ from vllm.model_executor.models.kimi_k25_vit import ( MoonViT3dPretrainedModel, vision_tower_forward, ) +from vllm.model_executor.models.vision import is_vit_use_data_parallel from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFieldConfig, @@ -336,9 +337,8 @@ class KimiK25ForConditionalGeneration( self.config = config quant_config = vllm_config.quant_config - # Check for MoonViT config compatibility - self.use_data_parallel = ( - model_config.multimodal_config.mm_encoder_tp_mode == "data" + self.use_data_parallel = is_vit_use_data_parallel( + config.vision_config.num_attention_heads ) self.hidden_size = config.text_config.hidden_size self.device = current_platform.current_device() diff --git a/vllm/model_executor/models/kimi_k25_vit.py b/vllm/model_executor/models/kimi_k25_vit.py index 29ecb84674..bcb8dd32de 100644 --- a/vllm/model_executor/models/kimi_k25_vit.py +++ b/vllm/model_executor/models/kimi_k25_vit.py @@ -357,7 +357,7 @@ class MoonViTEncoderLayer(nn.Module): attn_bias: bool = False, ): super().__init__() - self.use_data_parallel = is_vit_use_data_parallel() + self.use_data_parallel = is_vit_use_data_parallel(num_heads) self.num_heads = num_heads self.hidden_dim = hidden_dim diff --git a/vllm/model_executor/models/vision.py b/vllm/model_executor/models/vision.py index 0582c125c6..18e994c95d 100644 --- a/vllm/model_executor/models/vision.py +++ b/vllm/model_executor/models/vision.py @@ -139,11 +139,23 @@ def get_fp8_padded_hidden_size(num_heads: int, head_dim: int) -> int | None: return num_heads * round_up(head_dim, 16) -def is_vit_use_data_parallel(): +def is_vit_use_data_parallel(num_heads: int | None = None) -> bool: """ Get the tensor parallel type for Vision Transformer. """ mm_cfg = get_multimodal_config() + can_split = ( + num_heads % get_tensor_model_parallel_world_size() == 0 + if num_heads is not None + else None + ) + if num_heads is not None and not can_split: + logger.warning_once( + "The number of vision attention heads is not divisible by " + "the tensor parallel size. Falling back to data parallelism " + "for the vision encoder." + ) + return True return mm_cfg is not None and mm_cfg.mm_encoder_tp_mode == "data"