modeling_ovis.py

import os
from datetime import datetime
from importlib import import_module
from typing import List, Union, Callable, Optional

import deepspeed
import torch
from torch import Tensor, LongTensor, IntTensor
from torch.nn import init
from transformers import PreTrainedModel, AutoConfig, AutoModel, AutoTokenizer, AutoModelForCausalLM
from transformers.integrations.deepspeed import is_deepspeed_zero3_enabled, deepspeed_config

from .visual_tokenizer import ClipVisualTokenizer
from .configuration_ovis import OvisConfig
from .conversation_formatter import ConversationFormatter
from .utils import IGNORE_INDEX, IMAGE_TOKEN_INDEX, BEGIN_LINE, END_LINE, rank0_print

class VisualEmbedding(torch.nn.Embedding):
    def forward(self, input: Tensor) -> Tensor:
        if any((isinstance(input, LongTensor), isinstance(input, IntTensor))):
            return super().forward(input)
        return torch.matmul(input, self.weight)

    def reset_parameters(self, mean=0., std=1.) -> None:
        init.normal_(self.weight, mean=mean, std=std)
        self._fill_padding_idx_with_zero()


class OvisPreTrainedModel(PreTrainedModel):
    config_class = OvisConfig
    base_model_prefix = "ovis"


class Ovis(OvisPreTrainedModel):

    def __init__(self, config: OvisConfig, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        if kwargs.get('train_from_scratch'):
            self.llm = AutoModelForCausalLM.from_pretrained(kwargs['llm_name_or_path'], token=None)  # add token for gated model
            self.generation_config = self.llm.generation_config
            self.config.llm_config = self.llm.config
            self.config.hidden_size = self.llm.config.hidden_size  # for deepspeed auto configuration
            self.text_tokenizer = AutoTokenizer.from_pretrained(kwargs['llm_name_or_path'], token=None)  # add token for gated model
            if self.text_tokenizer.pad_token_id is None and kwargs.get('pad_token_id') is not None:
                self.text_tokenizer.pad_token_id = kwargs['pad_token_id']
            if kwargs.get('visual_tokenizer_pretrained_path'):
                self.visual_tokenizer = AutoModel.from_pretrained(kwargs['visual_tokenizer_pretrained_path'],
                                                                  image_processor_name_or_path=kwargs[
                                                                      'visual_tokenizer_pretrained_path'])
            else:
                self.visual_tokenizer = AutoModel.from_config(self.config.visual_tokenizer_config,
                                                              train_from_scratch=True)
            self.config.visual_tokenizer_config = self.visual_tokenizer.config
        else:
            self.llm = AutoModelForCausalLM.from_config(self.config.llm_config)
            assert self.config.hidden_size == self.llm.config.hidden_size, "hidden size mismatch"
            self.text_tokenizer = AutoTokenizer.from_pretrained(self.config.name_or_path)
            self.visual_tokenizer = AutoModel.from_config(self.config.visual_tokenizer_config,
                                                          image_processor_name_or_path=self.config.name_or_path)

        # initialize vte
        if is_deepspeed_zero3_enabled():
            with deepspeed.zero.Init(config_dict_or_path=deepspeed_config()):
                self.vte = VisualEmbedding(self.config.visual_tokenizer_config.vocab_size, self.config.hidden_size)
        else:
            self.visual_tokenizer.to(device=self.llm.device)
            self.vte = VisualEmbedding(self.config.visual_tokenizer_config.vocab_size, self.config.hidden_size,
                                       device=self.visual_tokenizer.device, dtype=self.visual_tokenizer.dtype)

        def _merge_modules(modules_list: tuple):
            merged_modules = []
            for modules in modules_list:
                merged_modules.extend(modules if modules else [])
            return merged_modules

        self._no_split_modules = _merge_modules((self.llm._no_split_modules, self.visual_tokenizer._no_split_modules))
        self._skip_keys_device_placement = self.llm._skip_keys_device_placement
        self._keep_in_fp32_modules = _merge_modules(
            (self.llm._keep_in_fp32_modules, self.visual_tokenizer._keep_in_fp32_modules))
        self.is_parallelizable = all((self.llm.is_parallelizable, self.visual_tokenizer.is_parallelizable))
        self.supports_gradient_checkpointing = all(
            (self.llm.supports_gradient_checkpointing, self.visual_tokenizer.supports_gradient_checkpointing))
        self._supports_flash_attn_2 = all(
            (self.llm._supports_flash_attn_2, self.visual_tokenizer._supports_flash_attn_2))
        self._supports_sdpa = all((self.llm._supports_sdpa, self.visual_tokenizer._supports_sdpa))
        self._supports_cache_class = all((self.llm._supports_cache_class, self.visual_tokenizer._supports_cache_class))

    def get_text_tokenizer(self):
        return self.text_tokenizer

    def get_visual_tokenizer(self):
        return self.visual_tokenizer

    def re_init_vte(self, mean, std):
        vte = self.get_vte()
        rank0_print(BEGIN_LINE)
        rank0_print(f'[{datetime.now()}] Before re-initialization of vte: ')
        with deepspeed.zero.GatheredParameters([vte.weight]):
            rank0_print(f'vte.weight: {vte.weight}')
        with deepspeed.zero.GatheredParameters([vte.weight], modifier_rank=0):
            if not is_deepspeed_zero3_enabled() or deepspeed.comm.get_rank() == 0:
                vte.reset_parameters(mean, std)
        rank0_print(f'[{datetime.now()}] After re-initialization of vte:')
        with deepspeed.zero.GatheredParameters([vte.weight]):
            rank0_print(f'vte.weight: {vte.weight}')
        rank0_print(END_LINE)

    def get_monitor_tensors(self):
        monitor_tensors = dict(
            wte=self.get_wte().weight,
            lm_head=self.get_lm_head().weight,
            vte=self.get_vte().weight
        )
        monitor_tensors.update(
            {f'visual_tokenizer_{k}': v for k, v in self.get_visual_tokenizer().get_monitor_tensors().items()})
        return monitor_tensors

    def get_lm_head(self):
        return self.get_llm().get_output_embeddings()

    def get_llm(self):
        return self.llm

    def get_vte(self):
        return self.vte

    def get_wte(self):
        return self.llm.get_input_embeddings()

    def get_conversation_formatter(self) -> ConversationFormatter:
        if getattr(self, 'conversation_formatter', None) is None:
            self.conversation_formatter = getattr(import_module(".conversation_formatter", __package__),
                                                  self.config.conversation_formatter_class)(self.text_tokenizer)
        return self.conversation_formatter

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, pixel_values: List[Optional[torch.Tensor]],
                labels: Optional[torch.Tensor] = None, **kwargs):
        _, inputs_embeds, labels, attention_mask = self.merge_multimodal(
            text_input_ids=input_ids,
            text_attention_masks=attention_mask,
            text_labels=labels,
            pixel_values=pixel_values,
            with_kv_cache=kwargs.get('past_key_values') is not None)
        return self.llm(inputs_embeds=inputs_embeds, labels=labels, attention_mask=attention_mask, **kwargs)

    def merge_multimodal(self, text_input_ids: torch.Tensor, text_attention_masks: torch.Tensor,
                         text_labels: Optional[torch.Tensor], pixel_values: List[Optional[torch.Tensor]],
                         with_kv_cache: bool = False):
        if with_kv_cache:
            return None, self.get_wte()(text_input_ids), text_labels, text_attention_masks
        if self.training:
            # When training, to be compatible with deepspeed zero, each sample has to include pixel_value tensor.
            # For text-only sample, one can simply use a full zero tensor as pixel_value, which will be ignored
            # (see below in this function); so, the gradient will not be affected.
            num_images = [x.shape[0] for x in pixel_values]
            visual_tokens = self.visual_tokenizer(torch.cat([x for x in pixel_values], dim=0))
            visual_embeds = torch.split(self.get_vte()(visual_tokens).to(dtype=self.dtype),
                                        split_size_or_sections=num_images, dim=0)
            visual_input_ids = torch.split(torch.argmax(visual_tokens, dim=-1),
                                           split_size_or_sections=num_images, dim=0)
            visual_labels = [torch.full(x.shape, IGNORE_INDEX, dtype=torch.long) for x in visual_input_ids]
        else:
            # When inference, sample can include only text with `None` pixel_value
            num_images = [x.shape[0] if x is not None else 0 for x in pixel_values]
            if sum(num_images) > 0:
                visual_tokens = self.visual_tokenizer(torch.cat([x for x in pixel_values if x is not None], dim=0))
                visual_embeds = torch.split(self.get_vte()(visual_tokens).to(dtype=self.dtype),
                                            split_size_or_sections=num_images, dim=0)
                visual_input_ids = torch.split(torch.argmax(visual_tokens, dim=-1),
                                               split_size_or_sections=num_images, dim=0)
                visual_labels = [torch.full(x.shape, IGNORE_INDEX, dtype=torch.long) for x in visual_input_ids]
            else:
                # just placeholders
                visual_embeds = [None] * len(num_images)
                visual_input_ids = [None] * len(num_images)
                visual_labels = [None] * len(num_images)
            # just placeholders
            text_labels = torch.full(text_input_ids.shape, IGNORE_INDEX, dtype=torch.long, device=text_input_ids.device)

        input_embeds = []
        attention_masks = []
        labels = []
        for text_input_id, text_label, text_attention_mask, visual_embed, visual_input_id, visual_label in zip(
                text_input_ids, text_labels, text_attention_masks, visual_embeds, visual_input_ids, visual_labels
        ):
            image_token_mask = torch.eq(text_input_id, IMAGE_TOKEN_INDEX)
            text_embed = self.get_wte()(torch.masked_fill(text_input_id, image_token_mask, 0))
            image_token_positions = torch.where(image_token_mask)[0].tolist()
            if len(image_token_positions) > 0:
                input_embed_parts = []
                attention_mask_parts = []
                label_parts = []
                prev_image_token_position = -1
                for index, image_token_position in enumerate(image_token_positions):
                    input_embed_parts.append(
                        text_embed[prev_image_token_position + 1:image_token_position, :])
                    label_parts.append(
                        text_label[prev_image_token_position + 1:image_token_position])
                    attention_mask_parts.append(
                        text_attention_mask[prev_image_token_position + 1:image_token_position])
                    input_embed_parts.append(visual_embed[index])
                    attention_mask_parts.append(
                        torch.ones_like(visual_label[index], device=text_label.device, dtype=torch.bool))
                    label_parts.append(visual_label[index].to(device=text_label.device))
                    prev_image_token_position = image_token_position
                if prev_image_token_position + 1 < text_input_id.shape[0]:
                    input_embed_parts.append(
                        text_embed[prev_image_token_position + 1:, :])
                    attention_mask_parts.append(
                        text_attention_mask[prev_image_token_position + 1:])
                    label_parts.append(
                        text_label[prev_image_token_position + 1:])
                input_embed = torch.cat(input_embed_parts, dim=0)
                attention_mask = torch.cat(attention_mask_parts, dim=0)
                label = torch.cat(label_parts, dim=0)
            else:
                input_embed = text_embed
                attention_mask = text_attention_mask
                label = text_label
                if self.training:
                    # Make visual_embed involved in the backward graph, to be compatible with deepspeed zero and ddp.
                    input_embed += torch.sum(visual_embed * 0.0)
            input_embeds.append(input_embed)
            attention_masks.append(attention_mask)
            labels.append(label)

        batch_input_embeds = torch.nn.utils.rnn.pad_sequence(input_embeds, batch_first=True, padding_value=0.0)[:,
                             :self.config.multimodal_max_length, :]
        batch_attention_mask = torch.nn.utils.rnn.pad_sequence(attention_masks, batch_first=True, padding_value=False)[
                               :,
                               :self.config.multimodal_max_length]
        batch_labels = torch.nn.utils.rnn.pad_sequence(labels, batch_first=True, padding_value=IGNORE_INDEX)[:,
                       :self.config.multimodal_max_length]

        return visual_input_ids, batch_input_embeds, batch_labels, batch_attention_mask

    def save_pretrained(
            self,
            save_directory: Union[str, os.PathLike],
            is_main_process: bool = True,
            state_dict: Optional[dict] = None,
            save_function: Callable = torch.save,
            push_to_hub: bool = False,
            max_shard_size: Union[int, str] = "5GB",
            safe_serialization: bool = True,
            variant: Optional[str] = None,
            token: Optional[Union[str, bool]] = None,
            save_peft_format: bool = True,
            **kwargs,
    ):
        super().save_pretrained(save_directory,
                                is_main_process=is_main_process,
                                state_dict=state_dict,
                                save_function=save_function,
                                safe_serialization=safe_serialization)
        self.get_text_tokenizer().save_pretrained(save_directory)
        self.get_visual_tokenizer().get_image_processor().save_pretrained(save_directory)

        # uncomment the following will additionally save a separate visual tokenizer
        # visual_tokenizer_directory = os.path.join(save_directory, 'visual_tokenizer')
        # self.get_visual_tokenizer().save_pretrained(visual_tokenizer_directory,
        #                                             is_main_process=is_main_process,
        #                                             state_dict=None,
        #                                             save_function=save_function,
        #                                             safe_serialization=safe_serialization)
        # self.get_visual_tokenizer().get_image_processor().save_pretrained(visual_tokenizer_directory)

    # TODO: support batch generation
    def prepare_inputs_for_generation(
            self, input_ids, pixel_values, attention_mask, past_key_values=None, inputs_embeds=None, **kwargs):
        if past_key_values is not None:
            input_ids = input_ids[:, -1:]
            attention_mask = attention_mask[:, -1:]

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        model_inputs.update(
            {
                "past_key_values": past_key_values,
                "use_cache": kwargs.get("use_cache"),
                "attention_mask": attention_mask,
                "pixel_values": pixel_values
            }
        )
        return model_inputs