projects/internvl_matcher/model/internvl.py

import copy
from collections import OrderedDict
from typing import List, Optional, Tuple, Union
from types import MethodType
import torch
import torch.distributed
import torch.nn as nn
from torch.nn import CrossEntropyLoss
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 xtuner.registry import BUILDER
from xtuner.model.modules import dispatch_modules
from transformers import AutoModel, AutoConfig, AutoTokenizer, BitsAndBytesConfig
from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutput, BaseModelOutputWithPooling
from .modules import VisualPromptEncodeModel
from .utils import (LoadWoInit, traverse_dict, make_inputs_require_grad, find_all_linear_names,
                    guess_load_checkpoint, get_peft_model_state_dict)


def vision_model_forward_cache(self,
                               pixel_values: Optional[torch.FloatTensor] = None,
                               visual_prompt_embeds: Optional[torch.FloatTensor] = None,
                               output_hidden_states: Optional[bool] = None,
                               return_dict: Optional[bool] = None,
                               pixel_embeds: Optional[torch.FloatTensor] = None,
                               )->Union[Tuple, BaseModelOutputWithPooling]:
    output_hidden_states = (
        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
    )
    return_dict = return_dict if return_dict is not None else self.config.use_return_dict

    if pixel_values is None and pixel_embeds is None:
        raise ValueError('You have to specify pixel_values or pixel_embeds')
    
    if pixel_embeds is not None:
        hidden_states = torch.cat([
            pixel_embeds[:, :1, :], pixel_embeds[:, 1:, :] + visual_prompt_embeds.flatten(2).transpose(1, 2)], dim=1)
    else:
        if len(pixel_values.shape) == 4:
            _pixel_embeds = self.embeddings(pixel_values)
            hidden_states = torch.cat([
                _pixel_embeds[:, :1, :], _pixel_embeds[:, 1:, :] + visual_prompt_embeds.flatten(2).transpose(1, 2)], dim=1)
        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
    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,
    )


def extract_feature_cache(self, 
                          pixel_values, 
                          visual_prompt_embeds):
    if self.select_layer == -1:
        vit_embeds = self.vision_model(
            pixel_values=pixel_values,
            visual_prompt_embeds=visual_prompt_embeds,
            output_hidden_states=False,
            return_dict=True).last_hidden_state
    else:
        vit_embeds = self.vision_model(
            pixel_values=pixel_values,
            visual_prompt_embeds=visual_prompt_embeds,
            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


class WrapInternVL(BaseModel):
    def __init__(self, 
                 mllm,
                 tokenizer=None,
                 freeze_llm=False,
                 freeze_visual_encoder=False,
                 freeze_connector=False,
                 unfreeze_lm_head=False,
                 llm_lora=None,
                 visual_encoder_lora=None,
                 quantization_vit=False,
                 quantization_llm=False,
                 pretrained_pth=None,
                 use_activation_checkpointing=True,
                 ):
        super().__init__()

        self.freeze_llm = freeze_llm
        self.freeze_visual_encoder = freeze_visual_encoder
        self.freeze_connector = freeze_connector
        self.unfreeze_lm_head = unfreeze_lm_head
        self.use_llm_lora = llm_lora is not None
        self.use_visual_encoder_lora = visual_encoder_lora is not None
        self.quantization_vit = quantization_vit
        self.quantization_llm = quantization_llm
        self.use_activation_checkpointing=use_activation_checkpointing
        if quantization_vit:
            assert visual_encoder_lora is not None
        if quantization_llm:
            assert quantization_llm and llm_lora is not None

        config = AutoConfig.from_pretrained(mllm["pretrained_model_name_or_path"], trust_remote_code=True)
        if config.llm_config.model_type == 'internlm2':
            config.llm_config.attn_implementation = 'flash_attention_2'
        else:
            config.llm_config._attn_implementation = 'flash_attention_2'

        if quantization_vit is False and quantization_llm is False:
            quantization = None
        else:
            llm_int8_skip_modules = ['mlp1']
            if quantization_llm and not quantization_vit:
                llm_int8_skip_modules.append('vision_model')
            
            if quantization_vit and not quantization_llm:
                llm_int8_skip_modules.append('language_model')
            
            quantization_config = dict(
                type=BitsAndBytesConfig,
                llm_int8_skip_modules=llm_int8_skip_modules,
                load_in_4bit=True,
                load_in_8bit=False,
                llm_int8_threshold=6.0,
                llm_int8_has_fp16_weight=False,
                bnb_4bit_compute_dtype=torch.float16,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type='nf4')
            quantization_clazz = quantization_config.pop('type')
            quantization = quantization_clazz(**quantization_config)

        with LoadWoInit():
            traverse_dict(mllm)
            model_clazz = mllm.pop('type')
            mllm.update(dict(quantization_config=quantization, config=config))
            # The weights in internvl2 modules have been loaded inside the calling of AutoModel.from_pretrained()
            self.model = model_clazz(**mllm)
        # self.model.language_model.config.use_cache = False
        dispatch_modules(self.model.language_model)
        
        self.model.vision_model.forward = MethodType(vision_model_forward_cache, self.model.vision_model)
        self.model.extract_feature = MethodType(extract_feature_cache, self.model)
        self.visual_prompt_encoder = VisualPromptEncodeModel(
            in_channels=3, vision_hidden_size=config.vision_config.hidden_size, 
            language_hidden_size=config.llm_config.hidden_size, force_image_size=config.force_image_size,
            patch_size=config.vision_config.patch_size, downsample_ratio=config.downsample_ratio).to(
            self.model.vision_model.dtype)
        
        if tokenizer is not None:
            self.tokenizer = self._build_from_cfg_or_module(tokenizer)
        else:
            self.tokenizer = AutoTokenizer.from_pretrained(mllm["pretrained_model_name_or_path"], trust_remote_code=True)
        img_context_token_id = self.tokenizer.convert_tokens_to_ids('<IMG_CONTEXT>')
        self.model.img_context_token_id = img_context_token_id
        self._add_special_tokens()

        if self.freeze_llm:
            self.model.language_model.requires_grad_(False)
        if self.freeze_visual_encoder:
            self.model.vision_model.requires_grad_(False)
        if self.freeze_connector:
            self.model.mlp1.requires_grad_(False)
        if self.unfreeze_lm_head:
            # self.model.language_model.get_output_embeddings().require_grad = True
            self.model.language_model.get_output_embeddings().requires_grad_(True)
            # for name, param in self.named_parameters():
            #     if 'tok_' in name or 'lm_head' in name:
            #         print("Unfrozen {} !!!".format(name))
            #         param.requires_grad_(True)
            #     if 'output.' in name and 'llm' in name and 'lora' not in name:
            #         print("Unfrozen {} !!!".format(name))
            #         param.requires_grad_(True)
        
        if use_activation_checkpointing:
            # it is necessary when using gradient checkpointing
            if hasattr(self.model.language_model, 'enable_input_require_grads'):
                    self.model.language_model.enable_input_require_grads()
            else:
                self.model.language_model.get_input_embeddings(
                ).register_forward_hook(make_inputs_require_grad)
        
        self.gradient_checkpointing_enable()
        
        if self.use_llm_lora:
            self._prepare_llm_for_lora(llm_lora)
        
        if self.use_visual_encoder_lora:
            self._prepare_visual_encoder_for_lora(visual_encoder_lora)
        
        if pretrained_pth is not None:
            pretrained_state_dict = guess_load_checkpoint(pretrained_pth)
            self.load_state_dict(pretrained_state_dict, strict=False)  # TODO, check whether the internvl2 weights are loaded correctly.
            print(f"Load pretrained weight from {pretrained_pth}")

        self._count = 0
        print_log(self, logger="current")
        print_log('InternVL_V1_5 construction is complete', logger='current')
    
    def _add_special_tokens(self):
        assert hasattr(self, "tokenizer")
        
        mark_tokens = [f'<mark{str(ii).zfill(3)}>' for ii in range(100)]
        added_tokens_num = self.tokenizer.add_tokens(mark_tokens)
        print_log(f'{added_tokens_num} special mark tokens were added successfully.', logger='current')
        
        self.model.language_model.resize_token_embeddings(len(self.tokenizer))
        
        self.mark_token_ids = {mark_token: self.tokenizer(
            mark_token, add_special_tokens=False).input_ids[0] for mark_token in mark_tokens}

        if self.use_activation_checkpointing or self.use_llm_lora or not self.freeze_llm:
            self.model.language_model.enable_input_require_grads()
        self.added_special_token = True
        
        return

    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 _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.model.language_model = prepare_model_for_kbit_training(
            self.model.language_model, use_activation_checkpointing)
        if lora_config.target_modules is None:
            modules = find_all_linear_names(self.model.language_model)
            lora_config.target_modules = modules
        self.model.language_model = get_peft_model(self.model.language_model, lora_config)

    def _prepare_visual_encoder_for_lora(self, lora_config):
        lora_config = self._parse_lora_config(lora_config)
        if lora_config.target_modules is None:
            modules = find_all_linear_names(self.model.vision_model)
            lora_config.target_modules = modules
        self.model.vision_model = get_peft_model(self.model.vision_model, lora_config)

    def gradient_checkpointing_enable(self):
        self.activation_checkpointing_enable()

    def activation_checkpointing_enable(self):
        self.model.language_model.gradient_checkpointing_enable()
    
    def gradient_checkpointing_disable(self):
        self.activation_checkpointing_disable()

    def activation_checkpointing_disable(self):
        self.model.language_model.gradient_checkpointing_disable()
    
    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.model.vision_model, state_dict=state_dict))
        elif not self.freeze_visual_encoder:
            to_return.update({
                k: v
                for k, v in state_dict.items() if 'model.vision_model.' in k
            })
        # Step 2. LLM
        if self.use_llm_lora:
            to_return.update(
                get_peft_model_state_dict(
                    self.model.language_model, state_dict=state_dict))
        elif not self.freeze_llm:
            to_return.update({
                k: v
                for k, v in state_dict.items() if 'model.language_model.'
            })
        # Step 3. Projector
        to_return.update(
            {k: v
             for k, v in state_dict.items() if 'model.mlp1.' in k})

        # prompt related models
        to_return.update(
            {k: v
             for k, v in state_dict.items() if 'visual_prompt_encoder.' in k})

        # embeds and so on
        # vocabulary embedding
        to_return.update(
            {k: v for k, v in state_dict.items() if 'tok_' in k or 'embed_tokens' in k}
        )
        # logit head
        to_return.update(
            {k: v for k, v in state_dict.items() if
             ('output.' in k or 'lm_head' in k) and 'llm' in k and 'lora' not in k}
        )

        return to_return
    
    def init_weights(self):
        pass

    def forward(self, data, data_samples=None, mode='loss'):
        pixel_values = data['pixel_values'].to(self.model.vision_model.dtype)
        visual_prompts = data['visual_prompts'].to(self.model.vision_model.dtype)
        merged_visual_prompts = data['merged_visual_prompts'].to(self.model.vision_model.dtype)
        num_patches = data['num_patches']
        num_vprompts = data['num_vprompts']
        sampled_mark_token_ids = data['sampled_mark_token_ids']

        # print('pixel values: ', pixel_values.shape)
        # print('visual prompts: ', visual_prompts.shape)
        # print('merged visual prompt: ', merged_visual_prompts.shape)
        # print('num patches: ', num_patches)
        # print('num_vprompts: ', num_vprompts)
        # exit(0)
        
        sampled_mark_tokens = [f'<mark{str(ii.item()).zfill(3)}>' for ii in sampled_mark_token_ids]
        sampled_mark_token_ids = torch.tensor(
            [self.mark_token_ids[mark_token] for mark_token in sampled_mark_tokens], 
            dtype=torch.long).to("cuda")
        # print("sampled mark tokens: ", sampled_mark_tokens)
        # print("sampled mark token ids: ", sampled_mark_token_ids)
        mark_embeddings = self.model.language_model.get_input_embeddings()(sampled_mark_token_ids)
     
        visual_prompts_patch_embeds = self.visual_prompt_encoder(
            merged_visual_prompts, visual_prompts, mark_embeddings, num_patches, num_vprompts)

        input_ids = data['input_ids']
        position_ids = data['position_ids']
        attention_mask = data['attention_mask']
        image_flags = data['image_flags']

        labels = data['labels']
        use_cache = False

        outputs = self._llm_forward(
            input_ids=input_ids,
            position_ids=position_ids,
            attention_mask=attention_mask,
            image_flags=image_flags,
            pixel_values=pixel_values,
            labels=labels,
            use_cache=use_cache,
            visual_prompt_embeds=visual_prompts_patch_embeds,
        )
        loss_dict = {'loss': outputs.loss}
        if mode == 'loss':
            return loss_dict
        else:
            raise NotImplementedError

    def _llm_forward(
        self,
        pixel_values: torch.FloatTensor,
        visual_prompt_embeds: torch.FloatTensor,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        image_flags: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[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,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        return_dict = return_dict if return_dict is not None \
            else self.model.config.use_return_dict
        
        image_flags = image_flags.squeeze(-1)
        # We only added the clone code here to avoid the error. Error will be thrown in the below try...except... codes
        input_embeds = self.model.language_model.get_input_embeddings()(input_ids).clone()
        # input_embeds = self.model.language_model.get_input_embeddings()(input_ids)

        vit_embeds = self.model.extract_feature(pixel_values, visual_prompt_embeds)
        # vit_embeds = self.model.extract_feature(pixel_values)
        vit_embeds = vit_embeds[image_flags == 1]
        vit_batch_size = pixel_values.shape[0]

        B, N, C = input_embeds.shape
        input_embeds = input_embeds.reshape(B*N, C)

        if torch.distributed.get_rank() == 0 and self._count % 100 == 0:
            print(f"dynamic ViT batch size: {vit_batch_size}, "
                  f"images per sample: {vit_batch_size}/B, "
                  f"dynamic token length: {N}")
        self._count += 1

        input_ids = input_ids.reshape(B * N)
        selected = (input_ids == self.model.img_context_token_id)
   
        try:
            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C).to(input_embeds.dtype)
        except Exception as e:
            vit_embeds = vit_embeds.reshape(-1, C)
            print(f"warning: {e}, input_embeds[selected].shape="
                  f"{input_embeds[selected].shape}, "
                  f"vit_embeds.shape={vit_embeds.shape}")
            n_token = selected.sum()
            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds[:n_token].to(input_embeds.dtype)
        
        input_embeds = input_embeds.reshape(B, N, C)

        outputs = self.model.language_model(
            inputs_embeds = input_embeds,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        logits = outputs.logits

        loss = None
        if labels is not None:
            # Shit so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.model.language_model.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # Enable model parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)
        
        if not return_dict:
            output = (logits, ) + outputs[1:]
            return (loss, ) + output if loss is not None else output
        
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )