from typing import Iterable, List, Optional, Tuple import torch from torch import nn from transformers import PretrainedConfig from sglang.srt.layers.attention import vision_utils 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, general_mm_embed_routine, ) from sglang.srt.managers.schedule_batch import ( Modality, MultimodalDataItem, MultimodalInputs, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.internvl import InternVisionModel from sglang.srt.models.qwen2 import Qwen2ForCausalLM from sglang.srt.models.qwen3 import Qwen3ForCausalLM from sglang.srt.models.qwen3_moe import Qwen3MoeForCausalLM from sglang.utils import logger class InternS1ForConditionalGeneration(nn.Module): def __init__( self, config: PretrainedConfig, quant_config: Optional[QuantizationConfig] = None, use_flash_attn=True, ) -> None: super().__init__() self.config = config self.quant_config = quant_config vision_utils.update_vit_attn_dummy_heads_config(self.config) image_size = ( getattr(config, "force_image_size", None) or config.vision_config.image_size ) patch_size = config.vision_config.patch_size if isinstance(image_size, list): image_size = image_size[0] if isinstance(patch_size, list): patch_size = patch_size[0] self.patch_size = patch_size self.select_layer = config.vision_feature_layer self.num_image_token = int( (image_size // patch_size) ** 2 * (config.downsample_ratio**2) ) self.downsample_ratio = config.downsample_ratio config.vision_config.use_flash_attn = True if use_flash_attn else False config.text_config._attn_implementation = ( "flash_attention_2" if use_flash_attn else "eager" ) logger.info(f"num_image_token: {self.num_image_token}") self.vision_model = InternVisionModel(config.vision_config) if config.text_config.architectures[0] == "Qwen2ForCausalLM": self.language_model = Qwen2ForCausalLM( config=config.text_config, quant_config=quant_config ) elif config.text_config.architectures[0] == "Qwen3MoeForCausalLM": self.language_model = Qwen3MoeForCausalLM( config=config.text_config, quant_config=quant_config ) elif config.text_config.architectures[0] == "Qwen3ForCausalLM": self.language_model = Qwen3ForCausalLM( config=config.text_config, quant_config=quant_config ) else: raise NotImplementedError( f"{config.text_config.architectures[0]} is not implemented." ) vit_hidden_size = config.vision_config.hidden_size llm_hidden_size = config.text_config.hidden_size self.mlp1 = nn.Sequential( nn.LayerNorm(vit_hidden_size * int(1 / self.downsample_ratio) ** 2), nn.Linear( vit_hidden_size * int(1 / self.downsample_ratio) ** 2, llm_hidden_size ), nn.GELU(), nn.Linear(llm_hidden_size, llm_hidden_size), ) def pixel_shuffle(self, x, scale_factor=0.5): n, w, h, c = x.size() # N, W, H, C --> N, W, H * scale, C // scale x = x.view(n, w, int(h * scale_factor), int(c / scale_factor)) # N, W, H * scale, C // scale --> N, H * scale, W, C // scale x = x.permute(0, 2, 1, 3).contiguous() # N, H * scale, W, C // scale --> N, H * scale, W * scale, C // (scale ** 2) x = x.view( n, int(h * scale_factor), int(w * scale_factor), int(c / (scale_factor * scale_factor)), ) x = x.permute(0, 2, 1, 3).contiguous() return x def extract_feature(self, pixel_values): if self.select_layer == -1: vit_embeds = self.vision_model( pixel_values=pixel_values, output_hidden_states=False, return_dict=True ).last_hidden_state else: vit_embeds = self.vision_model( pixel_values=pixel_values, output_hidden_states=True, return_dict=True ).hidden_states[self.select_layer] vit_embeds = vit_embeds[:, 1:, :] h = w = int(vit_embeds.shape[1] ** 0.5) vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1) vit_embeds = self.pixel_shuffle(vit_embeds, scale_factor=self.downsample_ratio) vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], -1, vit_embeds.shape[-1]) vit_embeds = self.mlp1(vit_embeds) return vit_embeds def get_image_feature(self, items: List[MultimodalDataItem]): """ Projects the last hidden state from the vision model into language model space. Returns: image_features (`torch.Tensor`): Image feature tensor of shape `(num_images, image_length, embed_dim)`). """ pixel_values = torch.cat([item.feature for item in items]) image_features = self.extract_feature(pixel_values) return image_features @torch.no_grad() def forward( self, input_ids: torch.Tensor, positions: torch.Tensor, forward_batch: ForwardBatch, input_embeds: torch.Tensor = None, ) -> torch.Tensor: hs = general_mm_embed_routine( input_ids=input_ids, forward_batch=forward_batch, language_model=self.language_model, data_embedding_funcs={ Modality.IMAGE: self.get_image_feature, }, positions=positions, ) return hs def pad_input_ids(self, input_ids: List[int], mm_inputs: MultimodalInputs): # Get all special token IDs im_start_id: int = mm_inputs.im_start_id im_end_id: int = mm_inputs.im_end_id media_token_pairs = [(im_start_id, im_end_id)] helper = MultiModalityDataPaddingPatternTokenPairs(media_token_pairs) return helper.pad_input_tokens(input_ids, mm_inputs) def _mapping_interns1_name(self, name): names_map = { "lm_head.weight": "language_model.lm_head.weight", "model.multi_modal_projector.layer_norm.bias": "mlp1.0.bias", "model.multi_modal_projector.layer_norm.weight": "mlp1.0.weight", "model.multi_modal_projector.linear_1.bias": "mlp1.1.bias", "model.multi_modal_projector.linear_1.weight": "mlp1.1.weight", "model.multi_modal_projector.linear_2.bias": "mlp1.3.bias", "model.multi_modal_projector.linear_2.weight": "mlp1.3.weight", "model.vision_tower.embeddings.cls_token": "vision_model.embeddings.class_embedding", "model.vision_tower.embeddings.patch_embeddings.projection.bias": "vision_model.embeddings.patch_embedding.bias", "model.vision_tower.embeddings.patch_embeddings.projection.weight": "vision_model.embeddings.patch_embedding.weight", "model.vision_tower.embeddings.position_embeddings": "vision_model.embeddings.position_embedding", } if name in names_map: name = names_map[name] elif name.startswith("model.language_model."): name = "language_model.model." + name[len("model.language_model.") :] elif name.startswith("model.vision_tower."): name = "vision_model." + name[len("model.vision_tower.") :] if name.startswith("vision_model.encoder.layer"): name = name.replace(r".layer.", r".layers.") name = name.replace(r".attention.", r".attn.attn.") name = name.replace(r".projection_layer.", r".proj.") name = name.replace(r".lambda_1", r".ls1") name = name.replace(r".lambda_2", r".ls2") name = name.replace(r".layernorm_before.", r".norm1.") name = name.replace(r".layernorm_after.", r".norm2.") return name def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): stacked_params_mapping = [ # (param_name, shard_name, shard_id) ("qkv_proj", "q_proj", "q"), ("qkv_proj", "k_proj", "k"), ("qkv_proj", "v_proj", "v"), ("gate_up_proj", "gate_proj", 0), ("gate_up_proj", "up_proj", 1), ] expert_params_mapping = [] if "Qwen3MoeForCausalLM" in self.config.text_config.architectures: expert_params_mapping = FusedMoE.make_expert_params_mapping( ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", num_experts=self.config.num_experts, ) params_dict = dict(self.named_parameters()) for name, loaded_weight in weights: if "rotary_emb.inv_freq" in name: continue name = self._mapping_interns1_name(name) if "vision_model" in name: loaded_weight = vision_utils.pad_vit_attn_dummy_heads( self.config, name, loaded_weight ) for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue # We have mlp.experts[0].gate_proj in the checkpoint. # Since we handle the experts below in expert_params_mapping, # we need to skip here BEFORE we update the name, otherwise # name will be updated to mlp.experts[0].gate_up_proj, which # will then be updated below in expert_params_mapping # for mlp.experts[0].gate_gate_up_proj, which breaks load. if "mlp.experts" in name: continue name = name.replace(weight_name, param_name) # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, shard_id) break else: for mapping in expert_params_mapping: param_name, weight_name, expert_id, shard_id = mapping if weight_name not in name: continue name = name.replace(weight_name, param_name) param = params_dict[name] weight_loader = param.weight_loader weight_loader( param, loaded_weight, name, shard_id=shard_id, expert_id=expert_id, ) break else: # Skip loading extra bias for GPTQ models. if name.endswith(".bias") and name not in params_dict: continue param = params_dict[name] weight_loader = getattr( param, "weight_loader", default_weight_loader ) weight_loader(param, loaded_weight) EntryClass = InternS1ForConditionalGeneration