# Copyright (c) OpenMMLab. All rights reserved. import math import copy import os.path as osp import warnings from collections import OrderedDict import torch import torch.nn as nn import numpy as np from accelerate import init_empty_weights from mmengine import print_log from mmengine.config import Config, ConfigDict from mmengine.model import BaseModel from peft import get_peft_model, prepare_model_for_kbit_training from transformers import (AddedToken, AutoConfig, CLIPImageProcessor, CLIPVisionModel, LlamaForCausalLM, LlamaTokenizerFast, LlavaConfig, LlavaForConditionalGeneration, LlavaProcessor ) from transformers.integrations import is_deepspeed_zero3_enabled from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask from transformers import PreTrainedModel, PretrainedConfig from xtuner.registry import BUILDER from xtuner.utils import DEFAULT_IMAGE_TOKEN from .modules import ProjectorConfig, ProjectorModel, dispatch_modules from .modules.dispatch import SUPPORT_FLASH1, SUPPORT_FLASH2 from .utils import (LoadWoInit, find_all_linear_names, get_peft_model_state_dict, guess_load_checkpoint, make_inputs_require_grad, prepare_inputs_labels_for_multimodal, traverse_dict) from .torchscale.model.LongNet import make_longnet_from_name from transformers.models.qwen2.modeling_qwen2 import Qwen2DecoderLayer import torch.nn.functional as F from torch.nn.init import trunc_normal_ # ====================== Copied from first file ====================== # def get_abs_pos(abs_pos, tgt_size): """ Interpolates 1D absolute positional embeddings to a target size. This function is modified to handle 1D positional embeddings, which is suitable for sequences of tokens that do not form a square grid. Args: abs_pos (torch.Tensor): The absolute positional embedding tensor of shape (N, C), where N is the original sequence length and C is the embedding dim. tgt_size (int): The target sequence length. Returns: torch.Tensor: The interpolated positional embedding tensor of shape (tgt_size, C). """ src_size = abs_pos.size(0) dtype = abs_pos.dtype if src_size == tgt_size: return abs_pos # For 1D interpolation, input tensor to F.interpolate should be (N, C, L) # We reshape our (L, C) tensor to (1, C, L) interp_input = abs_pos.float().permute(1, 0).unsqueeze(0) # Perform linear interpolation interp_output = F.interpolate( interp_input, size=tgt_size, mode='linear', align_corners=False, ) # Reshape back to (L_new, C) interpolated_pos = interp_output.squeeze(0).permute(1, 0).to(dtype) return interpolated_pos def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False): grid_h = np.arange(grid_size, dtype=np.float32) grid_w = np.arange(grid_size, dtype=np.float32) grid = np.meshgrid(grid_w, grid_h) grid = np.stack(grid, axis=0) grid = grid.reshape([2, 1, grid_size, grid_size]) pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid) if cls_token: pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) return pos_embed def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): assert embed_dim % 2 == 0 emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) emb = np.concatenate([emb_h, emb_w], axis=1) return emb def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): assert embed_dim % 2 == 0 omega = np.arange(embed_dim // 2, dtype=np.float32) omega /= embed_dim / 2. omega = 1. / 10000**omega pos = pos.reshape(-1) out = np.einsum('m,d->md', pos, omega) emb_sin = np.sin(out) emb_cos = np.cos(out) emb = np.concatenate([emb_sin, emb_cos], axis=1) return emb # Step 1: Create a configuration class for the Resampler class ResamplerConfig(PretrainedConfig): """ Configuration class for the Resampler module. """ model_type = "resampler" _auto_class = 'AutoConfig' def __init__( self, grid_size, embed_dim, num_heads, kv_dim=None, norm_layer=nn.LayerNorm, **kwargs ): self.grid_size = grid_size self.embed_dim = embed_dim self.num_heads = num_heads self.kv_dim = kv_dim # self.hidden_act = hidden_act self.norm_layer = norm_layer super().__init__(**kwargs) class Resampler(PreTrainedModel): _auto_class = 'AutoModel' config_class = ResamplerConfig base_model_prefix = 'model' supports_gradient_checkpointing = True def __init__( self, config: ResamplerConfig ): super().__init__(config) self.gradient_checkpointing = False self.num_queries = config.grid_size self.embed_dim = config.embed_dim self.num_heads = config.num_heads kv_dim = config.kv_dim norm_layer = config.norm_layer # REMOVED: Positional embedding initialization self.query = nn.Parameter(torch.zeros(self.num_queries, self.embed_dim)) self.query.data.normal_(mean=0.0, std=0.02) if kv_dim is not None and kv_dim != self.embed_dim: self.kv_proj = nn.Linear(kv_dim, self.embed_dim, bias=False) else: self.kv_proj = nn.Identity() self.attn = nn.MultiheadAttention(self.embed_dim, self.num_heads, batch_first=True) self.ln_q = norm_layer(self.embed_dim) self.ln_kv = norm_layer(self.embed_dim) nn.init.constant_(self.ln_q.bias, 0) nn.init.constant_(self.ln_q.weight, 1.0) nn.init.constant_(self.ln_kv.bias, 0) nn.init.constant_(self.ln_kv.weight, 1.0) def _init_weights(self, m): if isinstance(m, nn.Linear): trunc_normal_(m.weight, std=.02) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) def init_weights(self): self.query.data.normal_(mean=0.0, std=0.02) nn.init.constant_(self.ln_q.bias, 0) nn.init.constant_(self.ln_q.weight, 1.0) nn.init.constant_(self.ln_kv.bias, 0) nn.init.constant_(self.ln_kv.weight, 1.0) def _set_gradient_checkpointing(self, module, value=False): if isinstance(module, Resampler): module.gradient_checkpointing = value def forward(self, x, attn_mask=None, text=None): Q = self.query x = self.kv_proj(x) x = self.ln_kv(x) Q = self.ln_q(Q) # REMOVED: Positional embedding interpolation and addition out, attn = self.attn( Q.unsqueeze(0).expand(x.size(0), Q.size(0), Q.size(1)), x, x, attn_mask=attn_mask ) return out, attn # ====================== End of copied code ====================== # def convert_state_dict_to_hf(state_dict, mapping): new_state_dict = {} for key, value in state_dict.items(): if key.endswith('.inv_freq'): continue for key_to_modify, new_key in mapping.items(): if key_to_modify in key: key = key.replace(key_to_modify, new_key) new_state_dict[key] = value return new_state_dict class LLaVAModel(BaseModel): def __init__(self, llm, freeze_llm=True, visual_select_layer=-2, pretrained_pth=None, projector_depth=2, llm_lora=None, visual_encoder_lora=None, use_activation_checkpointing=True, max_position_embeddings=None, hidden_size=512, train_stage='2', enable_long_net=True, compressor_grid_size=2, # New parameter compressor_embed_dim=512, # New parameter prefusion_layer_num=4): # New parameter for prefusion layers super().__init__() self.enable_long_net = enable_long_net if enable_long_net: print('enable long net') else: print('disable long net') self.freeze_llm = freeze_llm self.freeze_visual_encoder = True if train_stage == '1': print('train_stage == 1') self.freeze_llm = True self.freeze_long_net = False elif train_stage == '2': print('train_stage == 2') self.freeze_llm = True self.freeze_long_net = True with LoadWoInit(): if isinstance(llm, dict): llm = self._dispatch_lm_model_cfg(llm, max_position_embeddings) self.llm = self._build_from_cfg_or_module(llm) self.encoder_name = "LongNet_{}_layers_{}_dim".format(2, 512) self.LongNet_encoder = make_longnet_from_name(self.encoder_name) self.llm.config.use_cache = False dispatch_modules(self.llm) self.projector_depth = projector_depth self.compressor_grid_size = compressor_grid_size self.compressor_embed_dim = compressor_embed_dim self.prefusion_layer_num = prefusion_layer_num # Build compressor compressor_config = ResamplerConfig( grid_size=compressor_grid_size, embed_dim=compressor_embed_dim, kv_dim= None, num_heads= 8, # Default value, can be adjusted ) self.compressor = Resampler( config=compressor_config ) self.compressor.init_weights() projector_config = ProjectorConfig( visual_hidden_size=hidden_size, llm_hidden_size=self.llm.config.hidden_size, depth=self.projector_depth) self.projector = ProjectorModel(projector_config).to( self.llm.dtype) # Build prefusion layers temps = copy.deepcopy(self.llm.model.layers[:prefusion_layer_num]) self.prefusion_layers = nn.ModuleList(temps) del temps # self.prefusion_layers=nn.ModuleList([Qwen2DecoderLayer(self.llm.config,layer_idx=i) for i in range(self.prefusion_layer_num)]) self.prefusion_layers.to(self.llm.dtype) if self.freeze_llm: print('freeze_llm') self.llm.requires_grad_(False) if self.freeze_long_net: print('freeze_long_net') self.LongNet_encoder.requires_grad_(False) # Move to correct dtype and device self.compressor = self.compressor.to(self.llm.dtype) if use_activation_checkpointing: if hasattr(self.llm, 'enable_input_require_grads'): self.llm.enable_input_require_grads() else: self.llm.get_input_embeddings().register_forward_hook( make_inputs_require_grad) self.projector.enable_input_require_grads() # for layer in self.prefusion_layers: # if hasattr(layer, 'enable_input_require_grads'): # layer.enable_input_require_grads() # else: # layer.get_input_embeddings().register_forward_hook( # make_inputs_require_grad) # enable gradient (activation) checkpointing for memory efficiency self.gradient_checkpointing_enable() self.use_llm_lora = None self.use_visual_encoder_lora = None if self.use_llm_lora: self._prepare_llm_for_lora(llm_lora, use_activation_checkpointing) if pretrained_pth is not None: # load the pretrained checkpoint pretrained_state_dict = guess_load_checkpoint(pretrained_pth) self.load_state_dict(pretrained_state_dict, strict=False) print_log(f'Load pretrained weight from {pretrained_pth}', 'current') self.visual_select_layer = visual_select_layer self._is_init = True self.is_first_iter = True def _parse_lora_config(self, lora_config): if isinstance(lora_config, dict) or isinstance( lora_config, Config) or isinstance(lora_config, ConfigDict): lora_config = BUILDER.build(lora_config) return lora_config def _prepare_llm_for_lora(self, lora_config, use_activation_checkpointing=True): lora_config = self._parse_lora_config(lora_config) self.llm = prepare_model_for_kbit_training( self.llm, use_activation_checkpointing) if lora_config.target_modules is None: modules = find_all_linear_names(self.llm) lora_config.target_modules = modules self.llm = get_peft_model(self.llm, lora_config) def _prepare_visual_encoder_for_lora(self, lora_config, use_activation_checkpointing=True): lora_config = self._parse_lora_config(lora_config) if lora_config.target_modules is None: modules = find_all_linear_names(self.visual_encoder) lora_config.target_modules = modules self.visual_encoder = get_peft_model(self.visual_encoder, lora_config) def gradient_checkpointing_enable(self): self.activation_checkpointing_enable() def activation_checkpointing_enable(self): self.llm.gradient_checkpointing_enable() self.compressor.gradient_checkpointing_enable() self.projector.gradient_checkpointing_enable() # for layer in self.prefusion_layers: # layer.gradient_checkpointing_enable() def gradient_checkpointing_disable(self): self.activation_checkpointing_disable() def activation_checkpointing_disable(self): self.llm.gradient_checkpointing_disable() self.compressor.gradient_checkpointing_disable() self.projector.gradient_checkpointing_disable() # for layer in self.prefusion_layers: # layer.gradient_checkpointing_disable() def init_weights(self): pass def state_dict(self, *args, **kwargs): state_dict = super().state_dict(*args, **kwargs) to_return = OrderedDict() # Step 1. visual_encoder if self.use_visual_encoder_lora: to_return.update( get_peft_model_state_dict( self.visual_encoder, state_dict=state_dict)) elif not self.freeze_visual_encoder: to_return.update({ k: v for k, v in state_dict.items() if 'visual_encoder.' in k }) # Step 2. LLM if self.use_llm_lora: to_return.update( get_peft_model_state_dict(self.llm, state_dict=state_dict)) elif not self.freeze_llm: to_return.update( {k: v for k, v in state_dict.items() if 'llm.' in k}) # Step 3. Compressor and Projector to_return.update( {k: v for k, v in state_dict.items() if 'compressor.' in k}) to_return.update( {k: v for k, v in state_dict.items() if 'projector.' in k}) # Step 4. Prefusion layers to_return.update( {k: v for k, v in state_dict.items() if 'prefusion_layers.' in k}) # Step 5. LongNet_encoder to_return.update( {k: v for k, v in state_dict.items() if 'LongNet_encoder.' in k}) return to_return @staticmethod def _prepare_for_long_context_training(cfg, llm_cfg, max_position_embeddings): orig_rope_scaling = getattr(llm_cfg, 'rope_scaling', None) if orig_rope_scaling is None: orig_rope_scaling = {'factor': 1} orig_rope_scaling_factor = orig_rope_scaling[ 'factor'] if 'factor' in orig_rope_scaling.keys() else 1 orig_ctx_len = getattr(llm_cfg, 'max_position_embeddings', None) if orig_ctx_len: orig_ctx_len *= orig_rope_scaling_factor if max_position_embeddings > orig_ctx_len: scaling_factor = float( math.ceil(max_position_embeddings / orig_ctx_len)) llm_cfg.rope_scaling = { 'type': 'linear', 'factor': scaling_factor } # hardcode for internlm2 llm_cfg.attn_implementation = 'flash_attention_2' cfg.config = llm_cfg return cfg, llm_cfg @staticmethod def _prepare_for_flash_attn(cfg, llm_cfg): cls_name = type(llm_cfg).__name__ SUPPORT_SDPA_ATTN = ('LlamaConfig', 'GemmaConfig', 'MistralConfig', 'MixtralConfig', 'Qwen2Config', 'Qwen2MoeConfig', 'Starcoder2Config', 'Starcoder2Config', 'Phi3Config') SUPPORT_FLASH_ATTN2 = ('InternLM2Config', 'LlamaConfig', 'GemmaConfig', 'MistralConfig', 'MixtralConfig', 'Qwen2Config', 'Qwen2MoeConfig', 'Starcoder2Config', 'Starcoder2Config', 'Phi3Config') torch_dtype = torch.bfloat16 if ( torch.cuda.is_available() and torch.cuda.is_bf16_supported()) \ else torch.float16 if getattr(cfg, 'attn_implementation', None) is not None: # Flash Attention 2.0 only supports torch.float16 and # torch.bfloat16 dtypes if cfg.attn_implementation == 'flash_attention_2': cfg.torch_dtype = torch_dtype elif SUPPORT_FLASH2 and cls_name in SUPPORT_FLASH_ATTN2: cfg.torch_dtype = torch_dtype cfg.attn_implementation = 'flash_attention_2' elif SUPPORT_FLASH1 and cls_name in SUPPORT_SDPA_ATTN: cfg.attn_implementation = 'sdpa' return cfg, llm_cfg @staticmethod def _prepare_for_qlora_zero3(cfg): if (not is_deepspeed_zero3_enabled()) or (not hasattr( cfg, 'quantization_config')): return cfg torch_dtype = torch.bfloat16 if ( torch.cuda.is_available() and torch.cuda.is_bf16_supported()) \ else torch.float16 cfg.torch_dtype = torch_dtype quantization_config = cfg.quantization_config quantization_config.bnb_4bit_compute_dtype = torch_dtype quantization_config.bnb_4bit_quant_storage = torch_dtype return cfg def _dispatch_lm_model_cfg(self, cfg, max_position_embeddings=None): cfg = self._prepare_for_qlora_zero3(cfg) pretrained_model_name_or_path = cfg.pretrained_model_name_or_path llm_cfg = AutoConfig.from_pretrained( pretrained_model_name_or_path, trust_remote_code=True) cfg, llm_cfg = self._prepare_for_flash_attn(cfg, llm_cfg) if max_position_embeddings is not None: cfg, llm_cfg = self._prepare_for_long_context_training( cfg, llm_cfg, max_position_embeddings) return cfg def _build_from_cfg_or_module(self, cfg_or_mod): if isinstance(cfg_or_mod, nn.Module): return cfg_or_mod elif isinstance(cfg_or_mod, dict): traverse_dict(cfg_or_mod) return BUILDER.build(cfg_or_mod) else: raise NotImplementedError def forward(self, data, data_samples=None, mode='loss'): if self.is_first_iter: self.to(data['input_ids'].device) self.is_first_iter = False if 'pixel_values' in data: feat_to_proj = data['pixel_values'].to(self.llm.dtype) if self.enable_long_net: long_net_output = self.LongNet_encoder(src_tokens=None, token_embeddings=feat_to_proj.permute(1, 0, 2))["encoder_out"] feat_to_proj = long_net_output.permute(1, 0, 2) # Apply compressor compressed_features, _ = self.compressor(feat_to_proj) # print_log('compressed_features shape: {}'.format(compressed_features.shape), 'current') # Apply projector pixel_values = self.projector(compressed_features) projected_global_image_features = self.projector(feat_to_proj) # Apply prefusion layers if any if self.prefusion_layer_num > 0: # print_log('Applying prefusion layers', 'current') input_ids = data['input_ids'] text_embeddings = self.llm.get_input_embeddings()(input_ids.clamp(min=0)).detach() padding_mask=(input_ids <= 0) # attention_mask = data['attention_mask'] # position_ids = data.get('position_ids', None) x = torch.cat([projected_global_image_features, pixel_values, text_embeddings], dim=1) mask=torch.cat((torch.zeros((padding_mask.size(0),projected_global_image_features.size(1)+pixel_values.size(1)),device=padding_mask.device).bool(),padding_mask),dim=1) # Prepare attention mask for prefusion layers if getattr(self.llm, "_use_flash_attention_2", False) or \ getattr(self.llm.config, "_attn_implementation", "") == "flash_attention_2": attention_mask = (~mask).int() else: attention_mask =_prepare_4d_causal_attention_mask(~mask, (x.size(0), x.size(1)), x, 0) position_ids = (~mask).int().long().cumsum(-1) - 1 position_ids.masked_fill_((~mask).int() == 0, 1) # Apply prefusion layers for layer in self.prefusion_layers: x = layer( x, attention_mask=attention_mask, position_ids=position_ids, use_cache=False, )[0] # data['inputs_embeds'] = x fusion_text_features=x[:,-1*input_ids.size(1):,:] pixel_values=x[:,-1*input_ids.size(1)-1*pixel_values.size(1):-1*input_ids.size(1),:] fusion_text_features=fusion_text_features*(~padding_mask).unsqueeze(-1).int()+ text_embeddings*padding_mask.unsqueeze(-1) data['text_features'] = fusion_text_features # print_log('text_features shape: {}'.format(fusion_text_features.shape), 'current') data['pixel_values'] = pixel_values data = prepare_inputs_labels_for_multimodal(llm=self.llm, **data) if mode == 'loss': return self.compute_loss(data, data_samples) elif mode == 'predict': return self.predict(data, data_samples) elif mode == 'tensor': return self._forward(data, data_samples) else: raise NotImplementedError def _forward(self, data, data_samples=None): outputs = self.llm(**data) return outputs def predict(self, data, data_samples=None): outputs = self.llm(**data) logits_dict = [{'logits': logits} for logits in outputs.logits] return logits_dict def compute_loss(self, data, data_samples=None): outputs = self.llm(**data) loss_dict = {'loss': outputs.loss} return loss_dict def __getattr__(self, name: str): try: return super().__getattr__(name) except AttributeError: return getattr(self.llm, name) def to_hf(self, cfg, save_dir, fp32=False, save_pretrained_kwargs={}, save_format='xtuner', **kwargs): if save_format == 'xtuner': self.to_xtuner_llava(cfg, save_dir, fp32, save_pretrained_kwargs) elif save_format == 'huggingface': self.to_huggingface_llava(cfg, save_dir, fp32, save_pretrained_kwargs) elif save_format == 'official': self.to_official_llava(cfg, save_dir, fp32, save_pretrained_kwargs) else: raise NotImplementedError def to_xtuner_llava(self, cfg, save_dir, fp32=False, save_pretrained_kwargs={}): self.llm.config.use_cache = True if not fp32: print_log('Convert LLM to float16', 'current') self.llm.half() if self.use_llm_lora: llm_path = osp.join(save_dir, 'llm_adapter') print_log(f'Saving LLM adapter to {llm_path}', 'current') self.llm.save_pretrained(llm_path, **save_pretrained_kwargs) elif not self.freeze_llm: llm_path = save_dir print_log(f'Saving LLM tokenizer to {llm_path}', 'current') tokenizer = BUILDER.build(cfg.tokenizer) tokenizer.save_pretrained(llm_path, **save_pretrained_kwargs) print_log(f'Saving LLM to {llm_path}', 'current') self.llm.save_pretrained(llm_path, **save_pretrained_kwargs) self.llm.config.use_cache = False # Visual Encoder if self.use_visual_encoder_lora: visual_encoder_path = osp.join(save_dir, 'visual_encoder_adapter') print_log( f'Saving visual_encoder adapter to {visual_encoder_path}', 'current') self.visual_encoder.save_pretrained(visual_encoder_path, **save_pretrained_kwargs) elif not self.freeze_visual_encoder: visual_encoder_path = osp.join(save_dir, 'visual_encoder') print_log( 'Saving visual_encoder image_processor to' f'{visual_encoder_path}', 'current') image_processor = BUILDER.build(cfg.image_processor) image_processor.save_pretrained(visual_encoder_path, **save_pretrained_kwargs) print_log(f'Saving visual_encoder to {visual_encoder_path}', 'current') self.visual_encoder.save_pretrained(visual_encoder_path, **save_pretrained_kwargs) # projector projector_path = osp.join(save_dir, 'projector') print_log(f'Saving projector to {projector_path}', 'current') self.projector.save_pretrained(projector_path, **save_pretrained_kwargs) # compressor compressor_path = osp.join(save_dir, 'compressor') print_log(f'Saving compressor to {compressor_path}', 'current') self.compressor.save_pretrained(compressor_path, **save_pretrained_kwargs) # Prefusion layers if self.prefusion_layer_num > 0: prefusion_path = osp.join(save_dir, 'prefusion_layers') print_log(f'Saving prefusion layers to {prefusion_path}', 'current') torch.save(self.prefusion_layers.state_dict(), osp.join(prefusion_path, 'prefusion_layers.bin')) # LongNet_encoder LongNet_encoder_path = osp.join(save_dir, 'LongNet_encoder') print_log(f'Saving LongNet_encoder to {LongNet_encoder_path}', 'current') self.LongNet_encoder.save_pretrained(LongNet_encoder_path, **save_pretrained_kwargs) def to_huggingface_llava(self, cfg, save_dir, fp32=False, save_pretrained_kwargs={}): LLM_MAPPING = { 'model': 'language_model.model', 'lm_head': 'language_model.lm_head', } VIT_MAPPING = { 'vision_model': 'vision_tower.vision_model', } PROJECTOR_MAPPING = { 'model.0': 'multi_modal_projector.linear_1', 'model.2': 'multi_modal_projector.linear_2', } COMPRESSOR_MAPPING = { 'compressor': 'compressor' } PREFUSION_MAPPING = { 'prefusion_layers': 'prefusion_layers' } LONGNET_MAPPING = { 'layers.0': 'LongNet_encoder.layers.0', 'layers.1': 'LongNet_encoder.layers.1', 'layer_norm': 'LongNet_encoder.layer_norm' } assert getattr(self.llm, 'hf_quantizer', None) is None, \ 'This conversion format does not support quantized LLM.' # get state_dict llm = self.llm if self.use_llm_lora: llm = self.llm.merge_and_unload() llm.config.use_cache = True if not fp32: print_log('Convert LLM to float16', 'current') llm.half() assert isinstance(llm, LlamaForCausalLM), \ 'This conversion format only supports LlamaForCausalLM.' llm_state_dict = llm.state_dict() llm_state_dict = convert_state_dict_to_hf(llm_state_dict, LLM_MAPPING) need_visual_encoder = (not self.freeze_visual_encoder or self.use_visual_encoder_lora) visual_encoder = self.visual_encoder if self.use_visual_encoder_lora: visual_encoder = self.visual_encoder.merge_and_unload() assert isinstance(visual_encoder, CLIPVisionModel),\ 'This conversion format only supports CLIPVisionModel.' if need_visual_encoder: visual_encoder_state_dict = visual_encoder.state_dict() visual_encoder_state_dict = convert_state_dict_to_hf( visual_encoder_state_dict, VIT_MAPPING) else: visual_encoder_state_dict = {} compressor_state_dict = self.compressor.state_dict() compressor_state_dict = convert_state_dict_to_hf( compressor_state_dict, COMPRESSOR_MAPPING) projector_state_dict = self.projector.state_dict() projector_state_dict = convert_state_dict_to_hf( projector_state_dict, PROJECTOR_MAPPING) prefusion_state_dict = self.prefusion_layers.state_dict() prefusion_state_dict = convert_state_dict_to_hf( prefusion_state_dict, PREFUSION_MAPPING) LongNet_encoder_state_dict = self.LongNet_encoder.state_dict() LongNet_encoder_state_dict = convert_state_dict_to_hf( LongNet_encoder_state_dict, LONGNET_MAPPING) state_dict = { **projector_state_dict, **compressor_state_dict, **prefusion_state_dict, **llm_state_dict, **visual_encoder_state_dict, **LongNet_encoder_state_dict } # init model text_config = llm.config vision_config = visual_encoder.config config = LlavaConfig( text_config=text_config, vision_config=vision_config, attn_implementation='eager') with init_empty_weights(): with warnings.catch_warnings(): warnings.filterwarnings( 'ignore', message='.*non-meta.*', category=UserWarning) model = LlavaForConditionalGeneration(config) model.load_state_dict(state_dict, strict=True, assign=True) # processor cfg.tokenizer.type = LlamaTokenizerFast.from_pretrained tokenizer = BUILDER.build(cfg.tokenizer) tokenizer.add_tokens( AddedToken(DEFAULT_IMAGE_TOKEN, special=True, normalized=False), special_tokens=True) tokenizer.add_special_tokens({'pad_token': ''}) image_processor = BUILDER.build(cfg.image_processor) assert isinstance(image_processor, CLIPImageProcessor),\ 'This conversion format only supports CLIPImageProcessor.' processor = LlavaProcessor( tokenizer=tokenizer, image_processor=image_processor) # Pad to 64 for performance reasons pad_shape = 64 pre_expansion_embeddings = \ model.language_model.model.embed_tokens.weight.data mu = torch.mean(pre_expansion_embeddings, dim=0).float() n = pre_expansion_embeddings.size()[0] sigma = ((pre_expansion_embeddings - mu).T @ (pre_expansion_embeddings - mu)) / n dist = torch.distributions.multivariate_normal.MultivariateNormal( mu, covariance_matrix=1e-5 * sigma) # We add an image token so we need to resize the model ori_vocab_size = config.text_config.vocab_size tokenizer_vocab_size = tokenizer.encode('')[-1] added_token = tokenizer_vocab_size - ori_vocab_size if added_token > 0: model.resize_token_embeddings(ori_vocab_size + added_token, pad_shape) model.language_model.model.embed_tokens.weight.data[ ori_vocab_size:] = torch.stack( tuple( dist.sample() for _ in range(model.language_model.model.embed_tokens. weight.data[ori_vocab_size:].shape[0])), dim=0, ) model.language_model.lm_head.weight.data[ ori_vocab_size:] = torch.stack( tuple(dist.sample() for _ in range(model.language_model.lm_head.weight. data[ori_vocab_size:].shape[0])), dim=0, ) model.config.image_token_index = tokenizer.encode( DEFAULT_IMAGE_TOKEN)[-1] model.config.pad_token_id = tokenizer.encode('')[-1] # save print_log(f'Saving to {save_dir}', 'current') model.save_pretrained(save_dir, **save_pretrained_kwargs) processor.save_pretrained(save_dir, **save_pretrained_kwargs) def to_official_llava(self, cfg, save_dir, fp32=False, save_pretrained_kwargs={}): VIT_MAPPING = { 'vision_model': 'model.vision_tower.vision_tower.vision_model', } PROJECTOR_MAPPING = { 'model.0': 'model.mm_projector.0', 'model.2': 'model.mm_projector.2', } COMPRESSOR_MAPPING = { 'compressor': 'model.compressor' } PREFUSION_MAPPING = { 'prefusion_layers': 'model.prefusion_layers' } LONGNET_MAPPING = { 'layers.0': 'LongNet_encoder.layers.0', 'layers.1': 'LongNet_encoder.layers.1', 'layer_norm': 'LongNet_encoder.layer_norm' } try: from llava.model import LlavaConfig, LlavaLlamaForCausalLM except ImportError: raise ImportError( 'Please install llava with ' '`pip install git+https://github.com/haotian-liu/LLaVA.git ' '--no-deps`.') assert getattr(self.llm, 'hf_quantizer', None) is None, \ 'This conversion format does not support quantized LLM.' # get state_dict llm = self.llm if self.use_llm_lora: llm = self.llm.merge_and_unload() llm.config.use_cache = True if not fp32: print_log('Convert LLM to float16', 'current') llm.half() assert isinstance(llm, LlamaForCausalLM), \ 'This conversion format only supports LlamaForCausalLM.' llm_state_dict = llm.state_dict() need_visual_encoder = (not self.freeze_visual_encoder or self.use_visual_encoder_lora) visual_encoder = self.visual_encoder if self.use_visual_encoder_lora: visual_encoder = self.visual_encoder.merge_and_unload() assert isinstance(visual_encoder, CLIPVisionModel),\ 'This conversion format only supports CLIPVisionModel.' if need_visual_encoder: visual_encoder_state_dict = visual_encoder.state_dict() visual_encoder_state_dict = convert_state_dict_to_hf( visual_encoder_state_dict, VIT_MAPPING) else: visual_encoder_state_dict = {} compressor_state_dict = self.compressor.state_dict() compressor_state_dict = convert_state_dict_to_hf( compressor_state_dict, COMPRESSOR_MAPPING) projector_state_dict = self.projector.state_dict() projector_state_dict = convert_state_dict_to_hf( projector_state_dict, PROJECTOR_MAPPING) prefusion_state_dict = self.prefusion_layers.state_dict() prefusion_state_dict = convert_state_dict_to_hf( prefusion_state_dict, PREFUSION_MAPPING) LongNet_encoder_state_dict = self.LongNet_encoder.state_dict() LongNet_encoder_state_dict = convert_state_dict_to_hf( LongNet_encoder_state_dict, LONGNET_MAPPING) state_dict = { **projector_state_dict, **compressor_state_dict, **prefusion_state_dict, **llm_state_dict, **visual_encoder_state_dict, **LongNet_encoder_state_dict } # init model tokenizer = BUILDER.build(cfg.tokenizer) image_processor = BUILDER.build(cfg.image_processor) assert isinstance(image_processor, CLIPImageProcessor),\ 'This conversion format only supports CLIPImageProcessor.' llava_config_dict = llm.config.__dict__.copy() llava_config_dict.update( dict( image_aspect_ratio='pad', mm_hidden_size=visual_encoder.config.hidden_size, mm_projector_type=f'mlp{self.projector_depth}x_gelu', mm_use_im_patch_token=False, mm_use_im_start_end=False, mm_vision_select_feature='patch', mm_vision_select_layer=self.visual_select_layer, mm_vision_tower=visual_encoder.config.name_or_path, unfreeze_mm_vision_tower=need_visual_encoder, model_type='llava', use_cache=True, use_mm_proj=True)) llava_config = LlavaConfig(**llava_config_dict) with init_empty_weights(): with warnings.catch_warnings(): warnings.filterwarnings( 'ignore', message='.*non-meta.*', category=UserWarning) model = LlavaLlamaForCausalLM(llava_config) model.load_state_dict(state_dict, strict=True, assign=True) # save print_log(f'Saving to {save_dir}', 'current') model.save_pretrained(save_dir, **save_pretrained_kwargs) image_processor.save_pretrained(save_dir, **save_pretrained_kwargs) tokenizer.save_pretrained(save_dir, **save_pretrained_kwargs)