chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2025-11-26 11:43:05 +08:00
committed by GitHub
parent f33e5d1ef1
commit ca5c8b16f6
3 changed files with 118 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

108
python/sglang/srt/models/internvl.py
View File

@@ -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

34
python/sglang/srt/multimodal/mm_utils.py
View File

@@ -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,