# -------------------------------------------------------- # YuanVL # Copyright (c) 2024 YuanLabAI # Licensed under The MIT License [see LICENSE for details] # -------------------------------------------------------- import warnings from typing import (Any, Callable, Iterable, List, Literal, Mapping, Optional, Set, Tuple, Type, TypedDict, Union) import torch.utils.checkpoint import transformers import torch from torch import nn from torch.nn import CrossEntropyLoss from transformers import (AutoModel, GenerationConfig, LlamaForCausalLM, LlamaTokenizer) from transformers.modeling_outputs import CausalLMOutputWithPast from transformers.modeling_utils import PreTrainedModel from transformers.utils import ModelOutput, logging from transformer_engine.pytorch import RMSNorm from transformers.activations import ACT2FN from .configuration_yuanvl import YuanVLChatConfig from .conversation import get_conv_template from .modeling_intern_vit import InternVisionModel, has_flash_attn from .modeling_yuanlm2 import YuanForCausalLM from .utils import flatten_bn, merge_multimodal_embeddings logger = logging.get_logger(__name__) class InternVLImagePixelInputs(TypedDict): type: Literal["pixel_values"] data: Union[torch.Tensor, List[torch.Tensor]] """ Shape: `(batch_size, 1 + num_patches, num_channels, height, width)` Note that `num_patches` may be different for each batch, in which case the data is passed as a list instead of a batched tensor. """ patches_per_image: List[int] """ List of number of total patches for each image in the batch. """ class InternVLImageEmbeddingInputs(TypedDict): type: Literal["image_embeds"] data: Any # in vllm vision this is a NestedTensors """ A tensor of shape `(num_images, total_image_feature_size, hidden_size)` or a list of tensors of shape `(total_image_feature_size, hidden_size)` `hidden_size` must match the hidden size of language model backbone. """ InternVLImageInputs = Union[InternVLImagePixelInputs, InternVLImageEmbeddingInputs] def version_cmp(v1, v2, op='eq'): import operator from packaging import version op_func = getattr(operator, op) return op_func(version.parse(v1), version.parse(v2)) class YuanImageMLP(nn.Module): def __init__( self, hidden_size: int, intermediate_size: int, output_size: int, hidden_act: str, ) -> None: super().__init__() #self.up_proj = ColumnParallelLinear(hidden_size, intermediate_size, bias=False,) #self.gate_proj = ColumnParallelLinear(hidden_size, intermediate_size, bias=False,) #self.down_proj = RowParallelLinear(intermediate_size, output_size, bias=False,) self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) self.down_proj = nn.Linear(intermediate_size, output_size, bias=False) if hidden_act != "silu": raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported for now.") self.act_fn = ACT2FN[hidden_act] @torch.compile def swiglu(self, y_1, y_2): return self.act_fn(y_1) * y_2 def forward(self, x): #import pdb x1 = self.up_proj(x) x2 = self.gate_proj(x) x3 = self.swiglu(x1, x2) #x3 = self.act_fn(x1) #x2 = self.gate_proj(x) x = self.down_proj(x3) return x class YuanVLChatModel(PreTrainedModel): config_class = YuanVLChatConfig main_input_name = 'pixel_values' base_model_prefix = 'language_model' _supports_flash_attn_2 = True _no_split_modules = ['InternVisionModel', 'YuanDeocderLayer'] def __init__(self, config: YuanVLChatConfig, vision_model=None, language_model=None, use_flash_attn=True): super().__init__(config) assert version_cmp(transformers.__version__, '4.37.0', 'ge') image_size = config.force_image_size or config.vision_config.image_size patch_size = config.vision_config.patch_size self.patch_size = patch_size self.select_layer = config.select_layer self.template = config.template self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2)) self.downsample_ratio = config.downsample_ratio self.ps_version = config.ps_version use_flash_attn = use_flash_attn if has_flash_attn else False config.vision_config.use_flash_attn = True if use_flash_attn else False config.llm_config._attn_implementation = 'flash_attention_2' if use_flash_attn else 'eager' logger.info(f'num_image_token: {self.num_image_token}') logger.info(f'ps_version: {self.ps_version}') if vision_model is not None: self.vision_model = vision_model else: self.vision_model = InternVisionModel(config.vision_config) if language_model is not None: self.language_model = language_model else: if config.llm_config.architectures[0] == 'YuanForCausalLM': self.language_model = YuanForCausalLM(config.llm_config) else: raise NotImplementedError(f'{config.llm_config.architectures[0]} is not implemented.') self.pixel_unshuffle = torch.nn.PixelUnshuffle(downscale_factor=2) #vit_hidden_size = config.vision_config.hidden_size #llm_hidden_size = config.llm_config.hidden_size #vit_mlp_ffn_hidden_size = config.vit_mlp_ffn_hidden_size #layernorm_epsilon = config.llm_config.layernorm_epsilon layernorm_epsilon = config.llm_config.rms_norm_eps self.imagemlp_input_hiddensize = int(config.vision_config.hidden_size / self.downsample_ratio ** 2) self.imagemlp_ffn_hidden_size = config.llm_config.ffn_hidden_size self.imagemlp = YuanImageMLP(self.imagemlp_input_hiddensize, self.imagemlp_ffn_hidden_size, output_size=config.llm_config.hidden_size, hidden_act="silu") self.imagemlp_layernorm = RMSNorm(config.llm_config.hidden_size, eps=layernorm_epsilon) ''' # modify internvl vision vit_hidden_size = config.vision_config.hidden_size llm_hidden_size = config.llm_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) ) ''' self.img_context_token_id = config.img_context_token_id self.conv_template = get_conv_template(self.template) self.system_message = self.conv_template.system_message def _validate_pixel_values(self, data: Union[torch.Tensor, List[torch.Tensor]] ) -> Union[torch.Tensor, List[torch.Tensor]]: h = w = self.config.vision_config.image_size expected_dims = (3, h, w) def _validate_shape(d: torch.Tensor): actual_dims = tuple(d.shape) if actual_dims != expected_dims: # expected_expr = ("num_patches", *map(str, expected_dims)) expected_expr = (expected_dims) raise ValueError("The expected shape of pixel values in each batch element " f" is {expected_expr}. You supplied {tuple(d.shape)}.") # data的数据类型可以是tensor，也可以是List[tensor] # 从这一段上来看，image tensor的个数为 imbs*num_images for d in data: _validate_shape(d) return data def _parse_and_validate_image_input(self, pixel_values: List[torch.Tensor] = None, image_token_id: torch.Tensor = None, image_embeds: torch.Tensor = None, ) -> Optional[InternVLImagePixelInputs]: # 没有图像数据 if pixel_values is None and image_embeds is None: return None # 传入数据有image_embeds if image_embeds is not None: if not isinstance(image_embeds, torch.Tensor): raise ValueError("Incorrect type of image embeddings. " f"Got type: {type(image_embeds)}") return InternVLImageEmbeddingInputs( type="image_embeds", data=flatten_bn(image_embeds), ) #self.img_context_token_id = image_token_id[0] if pixel_values is not None: if not isinstance(pixel_values, (torch.Tensor, list)): raise ValueError("Incorrect type of pixel values. " f"Got type: {type(pixel_values)}") patches_per_image = [] # bsz/request循环 for request_pixel_values in pixel_values: # 每个request的images循环 patches_per_image.append(request_pixel_values.shape[0]) # We need to flatten (B, N, P) to (B*N*P) # so we call flatten_bn twice. # (total_patches, 3, h, w) return InternVLImagePixelInputs( type="pixel_values", data=self._validate_pixel_values(flatten_bn(pixel_values)), patches_per_image=patches_per_image) raise AssertionError("This line should be unreachable") def _process_image_input( self, image_input: InternVLImageInputs, ) -> Tuple[torch.Tensor] : if image_input["type"] == "image_embeds": return image_input["data"] assert self.vision_model is not None # (total_patches, tokens_per_image, llm_config.hidden_size) image_embeds = self.extract_feature(image_input["data"]) patches_per_image = image_input["patches_per_image"] # Only one image in the current batch # bsz=1的情况，直接返回image_embeds if len(patches_per_image) == 1: # 返回一个tensor，[1, num_patches*256, text_config.hidden_size] image_embeds = image_embeds.view(-1, self.config.llm_config.hidden_size).unsqueeze(1) return image_embeds # NOTE: Image embeddings are split into separate tensors for each image # by the size of each embedding. # feature_size 每个patch 256个token位置 feature_size = image_embeds.shape[1] # (total_image_tokens, llm_config.hidden_size) image_embeds = image_embeds.view(-1, self.config.llm_config.hidden_size) image_feature_sizes = [num_patches * feature_size for num_patches in patches_per_image] # 切分后得到一个Tuple，元组每个元胞表示一个image的image_embed, [num_patches * 256, llm_config.hidden_size] image_embeds = image_embeds.split(image_feature_sizes) return image_embeds def get_multimodal_embeddings(self, pixel_values: Optional[List[torch.Tensor]] = None, image_token_id: Optional[List[torch.Tensor]] = None, image_embeds: Optional[List[torch.Tensor]] = None, image_input: InternVLImageInputs = None, ): image_input = self._parse_and_validate_image_input(pixel_values, image_token_id, image_embeds) if image_input is None: return None # image_input: (total_patches, 3, h, w) vision_embeddings = self._process_image_input(image_input) return vision_embeddings def get_input_embeddings( self, input_ids: torch.Tensor, multimodal_embeddings: Optional[torch.Tensor] ) -> torch.Tensor: # 生成 token_embeddings inputs_embeds = self.language_model.model.get_input_embeddings(input_ids) # 将image embed放到img_context_token_id的位置 if multimodal_embeddings is not None: assert self.img_context_token_id is not None # input_ids: torch.Tensor # inputs_embeds: torch.Tensor # multimodal_embeddings: torch.Tensor # placeholder_token_id: img_context_token_id inputs_embeds = merge_multimodal_embeddings( input_ids, inputs_embeds, multimodal_embeddings, self.img_context_token_id) return inputs_embeds def forward( self, input_ids: torch.LongTensor = None, attention_mask: torch.Tensor = None, position_ids: torch.LongTensor = None, past_key_values: List[torch.FloatTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, pixel_values: Optional[List[torch.Tensor]] = None, image_token_id: Optional[List[torch.Tensor]] = None, image_embeds: Optional[List[torch.Tensor]] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: import pdb pdb.set_trace() if inputs_embeds is None: # (images, patches * token_per_image) vision_embeddings = self.get_multimodal_embeddings(pixel_values, image_token_id, image_embeds) # (tokens, hidden_size) inputs_embeds = self.get_input_embeddings(input_ids, vision_embeddings).permute(1, 0, 2) input_ids = None hidden_states = self.language_model.model(input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, labels, use_cache, output_attentions, output_hidden_states, return_dict) return hidden_states 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))) if self.ps_version == 'v1': warnings.warn("In ps_version 'v1', the height and width have not been swapped back, " 'which results in a transposed image.') else: x = x.permute(0, 2, 1, 3).contiguous() return x # Internvl vision def extract_feature(self, pixel_values): # pixel_values: (imbs * num_image, ic, ih, iw) pixel_values = pixel_values.to(torch.bfloat16) output = self.vision_model(pixel_values=pixel_values) vit_embeds=output[0] # vit_embeds: (imbs * num_images, h*w, vit_dim) vit_embeds = vit_embeds[:, 1:, :] pn, phw, pc = vit_embeds.shape ph = pw = int(phw**0.5) vit_embeds = vit_embeds.view(pn, ph, pw, pc).permute(0, 3, 1, 2) vit_embeds = self.pixel_unshuffle(vit_embeds) pn, pc, ph, pw = vit_embeds.shape vit_embeds = vit_embeds.view(pn, pc, ph * pw).permute(0, 2, 1) num_images, cvs, chs = vit_embeds.shape #_, cvs, chs = vit_embeds.shape #assert self.imagemlp_ffn_hidden_size == chs #vit_embeds = vit_embeds.contiguous().view(imbs, num_image * cvs, chs).permute(1, 0, 2).contiguous() vit_embeds = vit_embeds.reshape(1, -1, vit_embeds.shape[-1]).permute(1, 0, 2) vit_embeds = self.imagemlp(vit_embeds) vit_embeds = self.imagemlp_layernorm(vit_embeds) vit_embeds = vit_embeds.view(num_images, cvs, -1) return vit_embeds @torch.no_grad() def generate( self, pixel_values: Optional[torch.FloatTensor] = None, input_ids: Optional[torch.FloatTensor] = None, attention_mask: Optional[torch.LongTensor] = None, visual_features: Optional[torch.FloatTensor] = None, generation_config: Optional[GenerationConfig] = None, position_ids: Optional[torch.Tensor] = None, output_hidden_states: Optional[bool] = None, ) -> torch.LongTensor: if pixel_values is not None: if visual_features is not None: vit_embeds = visual_features else: vit_embeds = self.get_multimodal_embeddings(pixel_values) inputs_embeds = self.get_input_embeddings(input_ids, vit_embeds) input_ids = None outputs = self.language_model.generate( inputs_embeds=inputs_embeds, attention_mask=attention_mask, generation_config=generation_config, output_hidden_states=output_hidden_states, position_ids=position_ids, max_length=8192, use_cache=True, ) return outputs