From ca5c8b16f67d13fb3369317834695ec0bcfbf66f Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Wed, 26 Nov 2025 11:43:05 +0800 Subject: [PATCH] [VLM] Support InternVL Vision Encoder Data Parallelism (#13925) Co-authored-by: luoyuan.luo --- python/sglang/srt/models/internvl.py | 108 +++++++++++++++++------ python/sglang/srt/multimodal/mm_utils.py | 34 +++++++ test/nightly/test_encoder_dp.py | 1 + 3 files changed, 118 insertions(+), 25 deletions(-) diff --git a/python/sglang/srt/models/internvl.py b/python/sglang/srt/models/internvl.py index 458cd95f2..105f2d33a 100644 --- a/python/sglang/srt/models/internvl.py +++ b/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.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) - self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) + self.act = get_act_fn(config.hidden_act) + self.fc1 = ColumnParallelLinear( + 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.mlp = InternMLP(config) + self.attn = InternAttention( + 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,23 +422,36 @@ class InternVisionModel(PreTrainedModel): hidden_states = self.embeddings(pixel_values) else: raise ValueError(f"wrong pixel_values size: {pixel_values.shape}") - encoder_outputs = self.encoder( - inputs_embeds=hidden_states, - output_hidden_states=output_hidden_states, - return_dict=return_dict, - ) - last_hidden_state = encoder_outputs.last_hidden_state + + 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, + return_dict=return_dict, + ) + last_hidden_state = encoder_outputs.last_hidden_state pooled_output = last_hidden_state[:, 0, :] if not return_dict: return (last_hidden_state, pooled_output) + encoder_outputs[1:] - return BaseModelOutputWithPooling( - last_hidden_state=last_hidden_state, - pooler_output=pooled_output, - hidden_states=encoder_outputs.hidden_states, - attentions=encoder_outputs.attentions, - ) + 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, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + ) class InternVLChatModel(nn.Module): @@ -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 diff --git a/python/sglang/srt/multimodal/mm_utils.py b/python/sglang/srt/multimodal/mm_utils.py index 12ed18934..afce1079b 100644 --- a/python/sglang/srt/multimodal/mm_utils.py +++ b/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, diff --git a/test/nightly/test_encoder_dp.py b/test/nightly/test_encoder_dp.py index cb47634cf..ed8232074 100644 --- a/test/nightly/test_encoder_dp.py +++ b/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), ]