[VLM] Support InternVL Vision Encoder Data Parallelism (#13925)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

python/sglang/srt/models/internvl.py

@@ -7,11 +7,16 @@ import torch
 import torch.nn.functional as F
 from torch import nn
 from transformers import PretrainedConfig, PreTrainedModel
-from transformers.activations import ACT2FN
 from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
 
+from sglang.srt.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
+from sglang.srt.layers.activation import get_act_fn
 from sglang.srt.layers.attention import vision_utils
 from sglang.srt.layers.attention.vision import SingletonCache, VisionAttention
+from sglang.srt.layers.linear import ColumnParallelLinear, RowParallelLinear
 from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
 from sglang.srt.layers.quantization.base_config import QuantizationConfig
 from sglang.srt.managers.mm_utils import MultiModalityDataPaddingPatternTokenPairs
@@ -28,6 +33,8 @@ from sglang.srt.models.internlm2 import InternLM2ForCausalLM
 from sglang.srt.models.qwen2 import Qwen2ForCausalLM
 from sglang.srt.models.qwen3 import Qwen3ForCausalLM
 from sglang.srt.models.qwen3_moe import Qwen3MoeForCausalLM
+from sglang.srt.multimodal.mm_utils import run_dp_sharded_vision_model
+from sglang.srt.server_args import get_global_server_args
 from sglang.utils import logger
 
 
@@ -36,6 +43,7 @@ class InternAttention(nn.Module):
         self,
         config,
         quant_config: QuantizationConfig = None,
+        use_data_parallel: bool = False,
     ):
         super().__init__()
         self.config = config
@@ -57,6 +65,7 @@ class InternAttention(nn.Module):
             qk_normalization=getattr(config, "qk_normalization", False)
             or getattr(config, "use_qk_norm", False),
             flatten_batch=False,
+            use_data_parallel=use_data_parallel,
         )
 
         self.proj_drop = nn.Dropout(config.dropout)
@@ -160,17 +169,39 @@ class InternRMSNorm(nn.Module):
 
 
 class InternMLP(nn.Module):
-    def __init__(self, config: PretrainedConfig):
+    def __init__(
+        self,
+        config: PretrainedConfig,
+        use_data_parallel: bool = False,
+    ):
         super().__init__()
+        self.tp_size = (
+            1 if use_data_parallel else get_tensor_model_parallel_world_size()
+        )
+        self.tp_rank = 0 if use_data_parallel else get_tensor_model_parallel_rank()
         self.config = config
-        self.act = ACT2FN[config.hidden_act]
+        self.act = get_act_fn(config.hidden_act)
-        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc1 = ColumnParallelLinear(
-        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+            config.hidden_size,
+            config.intermediate_size,
+            bias=True,
+            quant_config=None,
+            tp_size=self.tp_size,
+            tp_rank=self.tp_rank,
+        )
+        self.fc2 = RowParallelLinear(
+            config.intermediate_size,
+            config.hidden_size,
+            bias=True,
+            quant_config=None,
+            tp_size=self.tp_size,
+            tp_rank=self.tp_rank,
+        )
 
     def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        hidden_states = self.fc1(hidden_states)
+        hidden_states, _ = self.fc1(hidden_states)
         hidden_states = self.act(hidden_states)
-        hidden_states = self.fc2(hidden_states)
+        hidden_states, _ = self.fc2(hidden_states)
         return hidden_states
 
 
@@ -187,13 +218,18 @@ class InternVisionEncoderLayer(nn.Module):
         config: PretrainedConfig,
         drop_path_rate: float,
         quant_config: QuantizationConfig = None,
+        use_data_parallel: bool = False,
     ):
         super().__init__()
         self.embed_dim = config.hidden_size
         self.intermediate_size = config.intermediate_size
         self.norm_type = config.norm_type
-        self.attn = InternAttention(config=config, quant_config=quant_config)
+        self.attn = InternAttention(
-        self.mlp = InternMLP(config)
+            config=config,
+            quant_config=quant_config,
+            use_data_parallel=use_data_parallel,
+        )
+        self.mlp = InternMLP(config, use_data_parallel)
         self.norm1 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps)
         self.norm2 = NORM2FN[self.norm_type](self.embed_dim, eps=config.layer_norm_eps)
 
@@ -248,6 +284,7 @@ class InternVisionEncoder(nn.Module):
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        use_data_parallel: bool = False,
     ):
         super().__init__()
         self.config = config
@@ -258,7 +295,9 @@ class InternVisionEncoder(nn.Module):
         ]
         self.layers = nn.ModuleList(
             [
-                InternVisionEncoderLayer(config, dpr[idx], quant_config)
+                InternVisionEncoderLayer(
+                    config, dpr[idx], quant_config, use_data_parallel
+                )
                 for idx in range(config.num_hidden_layers)
             ]
         )
@@ -319,14 +358,16 @@ class InternVisionModel(PreTrainedModel):
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        use_data_parallel: bool = False,
     ):
         super().__init__(config)
-        self.config = config
 
+        self.config = config
+        self.use_data_parallel = use_data_parallel
         self.embeddings = InternVisionEmbeddings(
             config,
         )
-        self.encoder = InternVisionEncoder(config, quant_config)
+        self.encoder = InternVisionEncoder(config, quant_config, use_data_parallel)
 
     def resize_pos_embeddings(self, old_size, new_size, patch_size):
         pos_emb = self.embeddings.position_embedding
@@ -381,6 +422,11 @@ class InternVisionModel(PreTrainedModel):
                 hidden_states = self.embeddings(pixel_values)
             else:
                 raise ValueError(f"wrong pixel_values size: {pixel_values.shape}")
+
+        if self.use_data_parallel:
+            encoder_outputs = run_dp_sharded_vision_model(hidden_states, self.encoder)
+            last_hidden_state = encoder_outputs
+        else:
             encoder_outputs = self.encoder(
                 inputs_embeds=hidden_states,
                 output_hidden_states=output_hidden_states,
@@ -392,6 +438,14 @@ class InternVisionModel(PreTrainedModel):
         if not return_dict:
             return (last_hidden_state, pooled_output) + encoder_outputs[1:]
 
+        if self.use_data_parallel:
+            return BaseModelOutputWithPooling(
+                last_hidden_state=last_hidden_state,
+                pooler_output=pooled_output,
+                hidden_states=None,
+                attentions=None,
+            )
+        else:
             return BaseModelOutputWithPooling(
                 last_hidden_state=last_hidden_state,
                 pooler_output=pooled_output,
@@ -409,6 +463,7 @@ class InternVLChatModel(nn.Module):
     ) -> None:
         super().__init__()
         self.config = config
+        self.use_data_parallel = get_global_server_args().mm_enable_dp_encoder
         self.quant_config = quant_config
         vision_utils.update_vit_attn_dummy_heads_config(self.config)
         image_size = config.force_image_size or config.vision_config.image_size
@@ -430,7 +485,10 @@ class InternVLChatModel(nn.Module):
         logger.info(f"num_image_token: {self.num_image_token}")
         logger.info(f"ps_version: {self.ps_version}")
 
-        self.vision_model = InternVisionModel(config.vision_config)
+        self.vision_model = InternVisionModel(
+            config.vision_config,
+            use_data_parallel=self.use_data_parallel,
+        )
         if config.llm_config.architectures[0] == "Qwen2ForCausalLM":
             self.language_model = Qwen2ForCausalLM(
                 config=config.llm_config, quant_config=quant_config

python/sglang/srt/multimodal/mm_utils.py

@@ -427,6 +427,40 @@ def get_dp_encoder_lb_assignment(
     return (shuffle_indices, gpu_sample_counts, gpu_loads)
 
 
+# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/vision.py
+def run_dp_sharded_vision_model(
+    image_input: torch.Tensor, vision_model: torch.nn.Module
+) -> torch.Tensor:
+    """Run a vision model with data parallelism (DP) sharding. The function
+    will shard the input image tensor on the first dimension and run the vision
+    model
+
+    Args:
+        image_input (torch.Tensor): Image input tensor.
+        vision_model (torch.nn.Module): Vision model.
+    Returns:
+        torch.Tensor: Output image embeddings
+    """
+
+    num_chunks = image_input.shape[0]
+    mp_world_size = get_tensor_model_parallel_world_size()
+    num_chunks_per_rank = (num_chunks + mp_world_size - 1) // mp_world_size
+    num_padded_chunks = num_chunks_per_rank * mp_world_size - num_chunks
+    pad = (0,) * (2 * (image_input.dim() - 1)) + (0, num_padded_chunks)
+    image_input_padded = torch.nn.functional.pad(image_input, pad)
+    rank = get_tensor_model_parallel_rank()
+    image_input_per_rank = image_input_padded[
+        rank * num_chunks_per_rank : (rank + 1) * num_chunks_per_rank, ...
+    ]
+
+    vision_embeddings = vision_model(image_input_per_rank)
+    # Ensure tensor is contiguous before all_gather
+    vision_embeddings = vision_embeddings.last_hidden_state.contiguous()
+    vision_embeddings = tensor_model_parallel_all_gather(vision_embeddings, dim=0)
+    vision_embeddings = vision_embeddings[:num_chunks, ...]
+    return vision_embeddings
+
+
 # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/vision.py
 def run_dp_sharded_mrope_vision_model(
     vision_model: torch.nn.Module,

test/nightly/test_encoder_dp.py

@@ -19,6 +19,7 @@ from sglang.test.test_utils import (
 
 MODELS = [
     SimpleNamespace(model="Qwen/Qwen2.5-VL-72B-Instruct", mmmu_accuracy=0.55),
+    SimpleNamespace(model="OpenGVLab/InternVL2_5-8B", mmmu_accuracy=0.52),
 ]