from attrdict import AttrDict from dataclasses import dataclass import logging import gc from einops import rearrange, repeat from typing import Optional, List, Tuple, Callable, Union import torch import torch.nn as nn import torch.nn.functional as F from transformers.utils import ( add_start_docstrings, add_start_docstrings_to_model_forward, ) from transformers.modeling_outputs import ModelOutput from transformers.configuration_utils import PretrainedConfig from transformers import ( AutoConfig, AutoModelForCausalLM, PreTrainedModel ) from transformers.utils import logging from .siglip_vit import VisionTransformer from .configuration_deepseek import DeepseekV2Config from .modeling_deepseek import DeepseekV2ForCausalLM logger = logging.get_logger(__name__) class MlpProjector(nn.Module): def __init__(self, cfg): super().__init__() self.cfg = cfg if cfg.projector_type == "identity": modules = nn.Identity() elif cfg.projector_type == "linear": modules = nn.Linear(cfg.input_dim, cfg.n_embed) elif cfg.projector_type == "mlp_gelu": mlp_depth = cfg.depth modules = [nn.Linear(cfg.input_dim, cfg.n_embed)] for _ in range(1, mlp_depth): modules.append(nn.GELU()) modules.append(nn.Linear(cfg.n_embed, cfg.n_embed)) modules = nn.Sequential(*modules) elif cfg.projector_type == "downsample_mlp_gelu": mlp_depth = cfg.depth mlp_ratio = cfg.mlp_ratio modules = [nn.Linear(cfg.input_dim * cfg.downsample_ratio * cfg.downsample_ratio, cfg.n_embed * mlp_ratio)] for _ in range(1, mlp_depth - 1): modules.append(nn.GELU()) modules.append(nn.Linear(cfg.n_embed * mlp_ratio, cfg.n_embed * mlp_ratio)) modules.append(nn.GELU()) modules.append(nn.Linear(cfg.n_embed * mlp_ratio, cfg.n_embed)) modules = nn.Sequential(*modules) else: raise ValueError(f"Unknown projector type: {cfg.projector_type}") if cfg.token_pooling: self.token_pooling_layer = nn.Linear(cfg.input_dim * 4, cfg.input_dim) self.layers = modules def forward(self, x): if self.cfg.token_pooling: batch_size, wxh, channels = x.shape w = h = int(wxh ** 0.5) x = x.view(batch_size, w, h, channels) x = x.permute(0, 3, 1, 2) # import ipdb; ipdb.set_trace() patches = x.unfold(2, 2, 2).unfold(3, 2, 2) batch_size, channels, h_patches, w_patches, _, _ = patches.size() # 在通道维度上拼接 patches = patches.contiguous().view(batch_size, channels, h_patches * w_patches, -1) # 通过线性层 patches = patches.permute(0, 2, 1, 3).contiguous() patches = patches.view(batch_size, h_patches * w_patches, channels * 4) x = self.token_pooling_layer(patches) elif self.cfg.projector_type == 'downsample_mlp_gelu': bs, hw, input_dim = x.shape h = w = int((hw) ** 0.5) """compute padding""" if h % self.cfg.downsample_ratio: pad = self.cfg.downsample_ratio - h % self.cfg.downsample_ratio else: pad = 0 x = x.reshape(bs, h, w, input_dim) if pad > 0: x = F.pad(x, (0, 0, 0, pad, 0, pad), "constant", 0) """4 to 1 concat""" x = x.permute(0, 3, 1, 2) # B, C, H, W x = F.unfold(x, kernel_size=self.cfg.downsample_ratio, stride=self.cfg.downsample_ratio, padding=0) # B, C*4, HW // 4 x = x.permute(0, 2, 1) return self.layers(x) class VisionEncoderConfig(PretrainedConfig): model_type: str = "vision" model_name: str = "siglip_large_patch16_384" image_size: int = 384 patch_size: int = 16 width: int = 1024 layers: int = 24 heads: int = 16 mlp_ratio: int = 4 global_pool: str = "map" ignore_head: bool = True class_token: bool = False num_classes: int = 0 use_checkpoint: bool = False weight_init: str = "skip" deterministic: bool = False num_recomputing_layers: int = 0 def __init__( self, model_name: str = "siglip_large_patch16_384", image_size: int = 384, patch_size: int = 16, width: int = 1024, layers: int = 24, heads: int = 16, mlp_ratio: int = 4, global_pool: str = "map", ignore_head: bool = True, class_token: bool = False, num_classes: int = 0, use_checkpoint: bool = False, **kwargs ): self.model_name = model_name self.image_size = image_size self.patch_size = patch_size self.width = width self.layers = layers self.heads = heads self.mlp_ratio = mlp_ratio self.global_pool = global_pool self.ignore_head = ignore_head self.class_token = class_token self.num_classes = num_classes self.use_checkpoint = use_checkpoint super().__init__(**kwargs) class MlpProjectorConfig(PretrainedConfig): model_type = "mlp_projector" projector_type: str = "downsample_mlp_gelu" input_dim: int = 1152 n_embed: int = 2048 depth: int = 2 mlp_ratio: int = 1 downsample_ratio: int = 2 token_pooling: bool = False def __init__( self, projector_type: str = "downsample_mlp_gelu", input_dim: int = 1152, n_embed: int = 2048, depth: int = 2, mlp_ratio: int = 1, downsample_ratio: int = 2, **kwargs ): self.projector_type = projector_type self.input_dim = input_dim self.n_embed = n_embed self.depth = depth self.mlp_ratio = mlp_ratio self.downsample_ratio = downsample_ratio super().__init__(**kwargs) @dataclass class DeepSeekVLV2CausalLMOutputWithPast(ModelOutput): """ Base class for DeepSeek-VL2 causal language model (or autoregressive) outputs. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Language modeling loss (for next-token prediction). logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*): The rope index difference between sequence length and multimodal rope. """ loss: Optional[torch.FloatTensor] = None logits: torch.FloatTensor = None past_key_values: Optional[List[torch.FloatTensor]] = None hidden_states: Optional[Tuple[torch.FloatTensor]] = None attentions: Optional[Tuple[torch.FloatTensor]] = None rope_deltas: Optional[torch.LongTensor] = None class DeepseekVLV2Config(PretrainedConfig): model_type = "deepseek_vl_v2" vision_config: VisionEncoderConfig projector_config: MlpProjectorConfig language_config: DeepseekV2Config tile_tag: str = "2D" global_view_pos: str = "head" candidate_resolutions: Tuple[Tuple[int, int]] = ((384, 384),) def __init__( self, tile_tag: str = "tile_tag", global_view_pos: str = "head", candidate_resolutions: Tuple[Tuple[int, int]] = ((384, 384),), **kwargs ): super().__init__(**kwargs) vision_config = kwargs.get("vision_config", {}) self.vision_config = VisionEncoderConfig(**vision_config) projector_config = kwargs.get("projector_config", {}) self.projector_config = MlpProjectorConfig(**projector_config) language_config = kwargs.get("language_config", {}) if isinstance(language_config, DeepseekV2Config): self.language_config = language_config else: self.language_config = DeepseekV2Config(**language_config) self.tile_tag = tile_tag self.global_view_pos = global_view_pos self.candidate_resolutions = candidate_resolutions class DeepseekVLV2PreTrainedModel(PreTrainedModel): config_class = DeepseekVLV2Config base_model_prefix = "deepseek_vl_v2" _no_split_modules = [] _skip_keys_device_placement = "past_key_values" class DeepseekVLV2ForCausalLM(DeepseekVLV2PreTrainedModel): def __init__(self, config: DeepseekVLV2Config): super().__init__(config) self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" # ----------- vision encoder ------------ vision_config = config.vision_config self.vision = VisionTransformer( img_size=vision_config.image_size, patch_size=vision_config.patch_size, embed_dim=vision_config.width, depth=vision_config.layers, num_heads=vision_config.heads, mlp_ratio=vision_config.mlp_ratio, class_token=vision_config.class_token, global_pool=vision_config.global_pool, ignore_head=vision_config.ignore_head, weight_init=vision_config.weight_init, num_classes=0, deterministic=vision_config.deterministic, num_recomputing_layers=vision_config.num_recomputing_layers ) # ----------- vl projector ------------ projector_config = config.projector_config self.projector = MlpProjector(projector_config) # image token format 形式 # FIXME 目前tile tag & global_view_pos的默认取值都是之前的实验策略；后续应当去掉默认取值，改为没有取值就raise error self.tile_tag = config.tile_tag self.global_view_pos = config.global_view_pos # 用于format image token sequence的特殊token embed_std = 1 / torch.sqrt(torch.tensor(projector_config.n_embed, dtype=torch.float32)) if self.tile_tag == "2D": # <|view_separator|>, <|\n|> self.image_newline = nn.Parameter(torch.randn(projector_config.n_embed) * embed_std) # fix the typo: view_seperater self.view_seperator = nn.Parameter(torch.randn(projector_config.n_embed) * embed_std) elif self.tile_tag == "1D": # <|tile_x|>, <|tile_global|> candidate_resolutions = config.candidate_resolutions if len(candidate_resolutions) == 0: raise ValueError( f"len(candidate_resolutions) should be larger than 0, but got {len(candidate_resolutions)}") tile_variants_num = len(candidate_resolutions) self.tile_indicators = nn.Parameter( torch.randn(size=(tile_variants_num + 1, config.aligner.params.n_embed)) * embed_std ) else: raise ValueError(f"tile tag should be either 1D or 2D, but got {self.tile_tag}") # ----------- language model ------------ language_config = config.language_config self.language = DeepseekV2ForCausalLM(language_config) def prepare_inputs_embeds( self, input_ids: torch.LongTensor, images: Optional[torch.FloatTensor] = None, images_seq_mask: Optional[torch.LongTensor] = None, images_spatial_crop: Optional[torch.LongTensor] = None, **ignore_kwargs ): """ Args: input_ids (torch.LongTensor): [b, T] images (torch.FloatTensor): [b, max_n_images, 3, height, width] images_seq_mask (torch.BoolTensor): [b, T] images_spatial_crop (torch.LongTensor): [b, max_n_images, 2] Returns: input_embeds (torch.Tensor): [b, T, D] """ if images is None or images_spatial_crop.sum() == 0: return self.language.get_input_embeddings()(input_ids) bs, max_n_images, _ = images_spatial_crop.shape batch_num_tiles = [0 for _ in range(bs)] total_tiles = [] for idx in range(bs): for jdx in range(max_n_images): num_width_tiles, num_height_tiles = images_spatial_crop[idx, jdx] if num_width_tiles == 0 or num_height_tiles == 0: break batch_num_tiles[idx] += (1 + num_width_tiles * num_height_tiles) total_tiles.append(images[idx, :batch_num_tiles[idx]]) # [batch_all_tiles, 3, height, width] total_tiles = torch.cat(total_tiles, dim=0) assert total_tiles.shape[0] == sum(batch_num_tiles) if total_tiles.shape[0] == 0: return self.language.get_input_embeddings()(input_ids) # [batch_all_tiles, vit_seq_len, c] images_feature = self.vision(total_tiles) # [batch_all_tiles, hw, D] images_embeds = self.projector(images_feature) _, hw, n_dim = images_embeds.shape h = w = int(hw ** 0.5) # put image tokens into the input_embeds, [b, T, D] input_embeds = self.language.get_input_embeddings()(input_ids) # 根据self.tile_tag & self.global_view_pos填充image token sequence tile_index = 0 for idx in range(images_spatial_crop.shape[0]): images_in_this_batch = [] for jdx in range(images_spatial_crop.shape[1]): # extra global & local features num_width_tiles, num_height_tiles = images_spatial_crop[idx, jdx] if num_width_tiles == 0 or num_height_tiles == 0: break num_tiles_in_image = num_width_tiles * num_height_tiles # [hw, D] global_features = images_embeds[tile_index] # [num_height_tiles * num_width_tiles, hw, D] local_features = images_embeds[tile_index + 1: tile_index + 1 + num_tiles_in_image] tile_index += num_tiles_in_image + 1 # format global and local features if self.tile_tag == "2D": # ----------------- global view add newline ----------------- # [hw, D] -> [h, w, D] global_features = global_features.view(h, w, n_dim) # [D] -> [h, 1, D] new_lines_in_global = repeat(self.image_newline, "d -> h 1 d", h=h) # cat([h, w, D], [h, 1, D], dim=1) -> [h, w + 1, D] global_features = torch.cat([global_features, new_lines_in_global], dim=1) # [h, w + 1, D] -> [h * (w + 1), D] global_features = global_features.view(-1, n_dim) # ----------------- local view add newline ----------------- # [num_height_tiles * num_width_tiles, h * w, D] -> [num_height_tiles * h, num_width_tiles * w, D] local_features = rearrange( local_features, "(th tw) (h w) d -> (th h) (tw w) d", th=num_height_tiles, tw=num_width_tiles, h=h, w=w ) # [D] -> [num_height_tiles * h, 1, D] new_lines_in_local = repeat( self.image_newline, "d -> (th h) 1 d", th=num_height_tiles, h=h ) # [num_height_tiles * h, num_width_tiles * w + 1, D] local_features = torch.cat([local_features, new_lines_in_local], dim=1) # [num_height_tiles * h, num_width_tiles * w + 1, D] # --> [(num_height_tiles * h) * (num_width_tiles * w + 1), D] local_features = local_features.view(-1, n_dim) # ----------------- merge global and local tiles ----------------- if self.global_view_pos == "head": global_local_features = torch.cat( [global_features, self.view_seperator[None, :], local_features], dim=0) else: global_local_features = torch.cat( [local_features, self.view_seperator[None, :], global_features], dim=0) else: # abandoned，实际上不会走这个逻辑 global_features = torch.cat( [self.tile_indicators[0:1], global_features], dim=0 ) local_features = torch.cat( [self.tile_indicators[1:num_tiles_in_image + 1].unsqueeze(1), local_features], dim=1 ) local_features = rearrange(local_features, 'crop_num hw d -> (crop_num hw) d') if self.global_view_pos == "head": global_local_features = torch.cat([global_features, local_features], dim=0) else: global_local_features = torch.cat([local_features, global_features], dim=0) images_in_this_batch.append(global_local_features) if len(images_in_this_batch) > 0: images_in_this_batch = torch.cat(images_in_this_batch, dim=0) input_embeds[idx].masked_scatter_(images_seq_mask[idx].unsqueeze(-1), images_in_this_batch) return input_embeds @torch.no_grad() def incremental_prefilling( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, images: Optional[torch.FloatTensor] = None, images_seq_mask: Optional[torch.LongTensor] = None, images_spatial_crop: Optional[torch.LongTensor] = None, chunk_size: int = 1024 ): if inputs_embeds is None: inputs_embeds = self.prepare_inputs_embeds( input_ids=input_ids, images=images, images_seq_mask=images_seq_mask, images_spatial_crop=images_spatial_crop, ) del images del images_seq_mask del images_spatial_crop if attention_mask is not None: attention_mask = attention_mask.to(inputs_embeds.device) self._clear_cuda_cache() bzs, seq_len, _ = inputs_embeds.shape past_key_values = None # remain the last token for the next forward prefilling_len = seq_len - 1 for i in range(0, prefilling_len, chunk_size): chunk_start = i chunk_end = min(i + chunk_size, prefilling_len) chunk_inputs_embeds = inputs_embeds[:, chunk_start: chunk_end] chunk_attention_mask = attention_mask[:, 0: chunk_end] # print(f"start = {chunk_start}, end = {chunk_end}, prefilling_len = {prefilling_len}, seq_len = {seq_len}") # compute position_ids if past_key_values is not None: position_ids = torch.arange( chunk_start, chunk_end, dtype=torch.long, device=inputs_embeds.device ).unsqueeze(0) past_key_values = self._move_past_key_values_to_gpu(past_key_values, inputs_embeds.device) else: position_ids = None # chunk-forward with torch.no_grad(): outputs = self.forward( inputs_embeds=chunk_inputs_embeds, attention_mask=chunk_attention_mask, past_key_values=past_key_values, position_ids=position_ids, use_cache=True, ) # update past_key_values past_key_values = outputs.past_key_values past_key_values = self._move_past_key_values_to_cpu(past_key_values) del outputs, position_ids self._clear_cuda_cache() prefilling_key_values = [] for layer_past in past_key_values: prefilling_key_values.append( ( layer_past[0][:, :, 0: prefilling_len, ...].to(inputs_embeds.device), layer_past[1][:, :, 0: prefilling_len, ...].to(inputs_embeds.device), ) ) return inputs_embeds, prefilling_key_values def forward( self, input_ids: Optional[torch.LongTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, images: Optional[torch.FloatTensor] = None, images_seq_mask: Optional[torch.LongTensor] = None, images_spatial_crop: Optional[torch.LongTensor] = 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, cache_position: Optional[torch.LongTensor] = None, ): output_attentions = ( output_attentions if output_attentions is not None else self.config.output_attentions ) output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache return_dict = ( return_dict if return_dict is not None else self.config.use_return_dict ) if inputs_embeds is None: inputs_embeds = self.prepare_inputs_embeds( input_ids=input_ids, images=images, images_seq_mask=images_seq_mask, images_spatial_crop=images_spatial_crop, ) if attention_mask is not None: attention_mask = attention_mask.to(inputs_embeds.device) # print(inputs_embeds.shape) outputs = self.language.forward( input_ids=None, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, labels=labels, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict, cache_position=cache_position ) return outputs def _clear_cuda_cache(self): """clear CUDA memory cache""" gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() torch.cuda.synchronize() def _move_past_key_values_to_cpu(self, past_key_values): # print(f"past_key_values -> cpu") if past_key_values is None: return None return tuple(tuple(t.cpu() for t in layer) for layer in past_key_values) def _move_past_key_values_to_gpu(self, past_key_values, device="cuda:0"): # print(f"past_key_values -> gpu") if past_key_values is None: return None return tuple(tuple(t.to(device) for t in layer) for layer in past_key_values) def prepare_inputs_for_generation( self, input_ids, past_key_values=None, inputs_embeds=None, images: Optional[torch.FloatTensor] = None, images_seq_mask: Optional[torch.LongTensor] = None, images_spatial_crop: Optional[torch.LongTensor] = None, attention_mask=None, cache_position=None, pixel_values=None, image_sizes=None, num_logits_to_keep=None, **kwargs, ): # Overwritten -- in specific circumstances we don't want to forward image inputs to the model model_inputs = self.language.prepare_inputs_for_generation( input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, attention_mask=attention_mask, cache_position=cache_position, num_logits_to_keep=num_logits_to_keep, **kwargs, ) # If we're in cached decoding stage, pixel values should be None because input ids do not contain special image token anymore # Otherwise we need pixel values to be passed to model cache_position = model_inputs["cache_position"] if cache_position[0] == 0: model_inputs["images"] = images model_inputs["images_seq_mask"] = images_seq_mask model_inputs["images_spatial_crop"] = images_spatial_crop return model_inputs @staticmethod def _reorder_cache(past_key_values, beam_idx): reordered_past = () for layer_past in past_key_values: reordered_past += ( tuple( past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past ), ) return reordered_past AutoConfig.register("vision", VisionEncoderConfig) AutoConfig.register("mlp_projector", MlpProjectorConfig) AutoConfig.register("deepseek_vl_v2", DeepseekVLV2Config) AutoModelForCausalLM.register(DeepseekVLV2Config, DeepseekVLV2ForCausalLM)