First model version

.gitattributes

@@ -0,0 +1,21 @@
+model-00003-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00008-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00012-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00016-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00011-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00013-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+training_args.bin filter=lfs diff=lfs merge=lfs -text
+model-00006-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00009-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00010-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+tokenizer.model filter=lfs diff=lfs merge=lfs -text
+model-00015-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00007-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00014-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00017-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00001-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00002-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00004-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+model-00005-of-00017.safetensors filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+tokenizer.model filter=lfs diff=lfs merge=lfs -text

added_tokens.json

@@ -0,0 +1,9 @@
+{
+  "</box>": 64006,
+  "</quad>": 64002,
+  "</ref>": 64004,
+  "<IMG_CONTEXT>": 64000,
+  "<box>": 64005,
+  "<quad>": 64001,
+  "<ref>": 64003
+}

config.json

@@ -0,0 +1,187 @@
+{
+  "_commit_hash": "21d6ce2f09ce86a2dd76e1d6672d2143e060de34",
+  "_name_or_path": "OpenGVLab/InternVL-Chat-Chinese-V1-2-Plus",
+  "architectures": [
+    "InternVLChatModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_internvl_chat.InternVLChatConfig",
+    "AutoModel": "modeling_internvl_chat.InternVLChatModel"
+  },
+  "downsample_ratio": 0.5,
+  "force_image_size": 448,
+  "llm_config": {
+    "_name_or_path": "01-ai/Yi-34B",
+    "add_cross_attention": false,
+    "architectures": [
+      "LlamaForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attn_implementation": "flash_attention_2",
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 1,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 7,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "silu",
+    "hidden_size": 7168,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 20480,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 4096,
+    "min_length": 0,
+    "model_type": "llama",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 56,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 60,
+    "num_key_value_heads": 8,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": 0,
+    "prefix": null,
+    "pretraining_tp": 1,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rms_norm_eps": 1e-05,
+    "rope_scaling": null,
+    "rope_theta": 5000000.0,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": false,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "transformers_version": "4.36.2",
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": false,
+    "vocab_size": 64007
+  },
+  "model_type": "internvl_chat",
+  "pad2square": false,
+  "select_layer": -1,
+  "template": "Hermes-2",
+  "torch_dtype": "bfloat16",
+  "transformers_version": null,
+  "use_backbone_lora": 0,
+  "use_llm_lora": 0,
+  "vision_config": {
+    "_name_or_path": "",
+    "add_cross_attention": false,
+    "architectures": [
+      "InternVisionModel"
+    ],
+    "attention_dropout": 0.0,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "drop_path_rate": 0.0,
+    "dropout": 0.0,
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": null,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "hidden_act": "gelu",
+    "hidden_size": 3200,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "image_size": 448,
+    "initializer_factor": 0.1,
+    "initializer_range": 1e-10,
+    "intermediate_size": 12800,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_norm_eps": 1e-06,
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "min_length": 0,
+    "model_type": "intern_vit_6b",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 25,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_channels": 3,
+    "num_hidden_layers": 45,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "patch_size": 14,
+    "prefix": null,
+    "problem_type": null,
+    "pruned_heads": {},
+    "qk_normalization": true,
+    "qkv_bias": false,
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torch_dtype": "bfloat16",
+    "torchscript": false,
+    "transformers_version": "4.36.2",
+    "typical_p": 1.0,
+    "use_bfloat16": true,
+    "use_flash_attn": true
+  }
+}

configuration_intern_vit.py

@@ -0,0 +1,117 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+import os
+from typing import Union
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class InternVisionConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`InternVisionModel`]. It is used to
+    instantiate a vision encoder according to the specified arguments, defining the model architecture.
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+    Args:
+        num_channels (`int`, *optional*, defaults to 3):
+            Number of color channels in the input images (e.g., 3 for RGB).
+        patch_size (`int`, *optional*, defaults to 14):
+            The size (resolution) of each patch.
+        image_size (`int`, *optional*, defaults to 224):
+            The size (resolution) of each image.
+        qkv_bias (`bool`, *optional*, defaults to `False`):
+            Whether to add a bias to the queries and values in the self-attention layers.
+        hidden_size (`int`, *optional*, defaults to 3200):
+            Dimensionality of the encoder layers and the pooler layer.
+        num_attention_heads (`int`, *optional*, defaults to 25):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        intermediate_size (`int`, *optional*, defaults to 12800):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        qk_normalization (`bool`, *optional*, defaults to `True`):
+            Whether to normalize the queries and keys in the self-attention layers.
+        num_hidden_layers (`int`, *optional*, defaults to 48):
+            Number of hidden layers in the Transformer encoder.
+        use_flash_attn (`bool`, *optional*, defaults to `True`):
+            Whether to use flash attention mechanism.
+        hidden_act (`str` or `function`, *optional*, defaults to `"gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
+            `"relu"`, `"selu"` and `"gelu_new"` ``"gelu"` are supported.
+        layer_norm_eps (`float`, *optional*, defaults to 1e-6):
+            The epsilon used by the layer normalization layers.
+        dropout (`float`, *optional*, defaults to 0.0):
+            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
+        drop_path_rate (`float`, *optional*, defaults to 0.0):
+            Dropout rate for stochastic depth.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        initializer_factor (`float`, *optional*, defaults to 0.1):
+            A factor for layer scale.
+    """
+
+    model_type = 'intern_vit_6b'
+
+    def __init__(
+            self,
+            num_channels=3,
+            patch_size=14,
+            image_size=224,
+            qkv_bias=False,
+            hidden_size=3200,
+            num_attention_heads=25,
+            intermediate_size=12800,
+            qk_normalization=True,
+            num_hidden_layers=48,
+            use_flash_attn=True,
+            hidden_act='gelu',
+            layer_norm_eps=1e-6,
+            dropout=0.0,
+            drop_path_rate=0.0,
+            attention_dropout=0.0,
+            initializer_range=0.02,
+            initializer_factor=0.1,
+            **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.dropout = dropout
+        self.drop_path_rate = drop_path_rate
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_channels = num_channels
+        self.patch_size = patch_size
+        self.image_size = image_size
+        self.initializer_range = initializer_range
+        self.initializer_factor = initializer_factor
+        self.attention_dropout = attention_dropout
+        self.layer_norm_eps = layer_norm_eps
+        self.hidden_act = hidden_act
+        self.qkv_bias = qkv_bias
+        self.qk_normalization = qk_normalization
+        self.use_flash_attn = use_flash_attn
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs) -> 'PretrainedConfig':
+        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+
+        if 'vision_config' in config_dict:
+            config_dict = config_dict['vision_config']
+
+        if 'model_type' in config_dict and hasattr(cls, 'model_type') and config_dict['model_type'] != cls.model_type:
+            logger.warning(
+                f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
+                f'{cls.model_type}. This is not supported for all configurations of models and can yield errors.'
+            )
+
+        return cls.from_dict(config_dict, **kwargs)

configuration_internvl_chat.py

@@ -0,0 +1,75 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+
+import copy
+
+from transformers import LlamaConfig
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+from .configuration_intern_vit import InternVisionConfig
+
+logger = logging.get_logger(__name__)
+
+
+class InternVLChatConfig(PretrainedConfig):
+    model_type = 'internvl_chat'
+    is_composition = True
+
+    def __init__(
+            self,
+            vision_config=None,
+            llm_config=None,
+            use_backbone_lora=0,
+            use_llm_lora=0,
+            pad2square=False,
+            select_layer=-4,
+            force_image_size=None,
+            downsample_ratio=0.5,
+            template=None,
+            **kwargs):
+        super().__init__(**kwargs)
+
+        if vision_config is None:
+            vision_config = {}
+            logger.info('vision_config is None. Initializing the InternVisionConfig with default values.')
+
+        if llm_config is None:
+            llm_config = {}
+            logger.info('llm_config is None. Initializing the LlamaConfig config with default values (`LlamaConfig`).')
+
+        self.vision_config = InternVisionConfig(**vision_config)
+        self.llm_config = LlamaConfig(**llm_config)
+        self.use_backbone_lora = use_backbone_lora
+        self.use_llm_lora = use_llm_lora
+        self.pad2square = pad2square
+        self.select_layer = select_layer
+        self.force_image_size = force_image_size
+        self.downsample_ratio = downsample_ratio
+        self.template = template
+
+        logger.info(f'vision_select_layer: {self.select_layer}')
+
+    def to_dict(self):
+        """
+        Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
+
+        Returns:
+            `Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
+        """
+        output = copy.deepcopy(self.__dict__)
+        output['vision_config'] = self.vision_config.to_dict()
+        output['llm_config'] = self.llm_config.to_dict()
+        output['model_type'] = self.__class__.model_type
+        output['use_backbone_lora'] = self.use_backbone_lora
+        output['use_llm_lora'] = self.use_llm_lora
+        output['pad2square'] = self.pad2square
+        output['select_layer'] = self.select_layer
+        output['force_image_size'] = self.force_image_size
+        output['downsample_ratio'] = self.downsample_ratio
+        output['template'] = self.template
+
+        return output

conversation.py

@@ -0,0 +1,1243 @@
+"""
+Conversation prompt templates.
+
+We kindly request that you import fastchat instead of copying this file if you wish to use it.
+If you have any changes in mind, please contribute back so the community can benefit collectively and continue to maintain these valuable templates.
+"""
+
+import dataclasses
+from enum import IntEnum, auto
+from typing import Any, Dict, List, Tuple, Union
+
+
+class SeparatorStyle(IntEnum):
+    """Separator styles."""
+
+    ADD_COLON_SINGLE = auto()
+    ADD_COLON_TWO = auto()
+    ADD_COLON_SPACE_SINGLE = auto()
+    NO_COLON_SINGLE = auto()
+    NO_COLON_TWO = auto()
+    ADD_NEW_LINE_SINGLE = auto()
+    LLAMA2 = auto()
+    CHATGLM = auto()
+    CHATML = auto()
+    CHATINTERN = auto()
+    DOLLY = auto()
+    RWKV = auto()
+    PHOENIX = auto()
+    ROBIN = auto()
+    FALCON_CHAT = auto()
+    CHATGLM3 = auto()
+    INTERNVL_ZH = auto()
+    MPT = auto()
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that manages prompt templates and keeps all conversation history."""
+
+    # The name of this template
+    name: str
+    # The template of the system prompt
+    system_template: str = '{system_message}'
+    # The system message
+    system_message: str = ''
+    # The names of two roles
+    roles: Tuple[str] = ('USER', 'ASSISTANT')
+    # All messages. Each item is (role, message).
+    messages: List[List[str]] = ()
+    # The number of few shot examples
+    offset: int = 0
+    # The separator style and configurations
+    sep_style: SeparatorStyle = SeparatorStyle.ADD_COLON_SINGLE
+    sep: str = '\n'
+    sep2: str = None
+    # Stop criteria (the default one is EOS token)
+    stop_str: Union[str, List[str]] = None
+    # Stops generation if meeting any token in this list
+    stop_token_ids: List[int] = None
+
+    def get_prompt(self) -> str:
+        """Get the prompt for generation."""
+        system_prompt = self.system_template.format(system_message=self.system_message)
+        if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ': ' + message + seps[i % 2]
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_SPACE_SINGLE:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ': '  # must be end with a space
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_NEW_LINE_SINGLE:
+            ret = '' if system_prompt == '' else system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + message + self.sep
+                else:
+                    ret += role + '\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + message + seps[i % 2]
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.RWKV:
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += (
+                        role
+                        + ': '
+                        + message.replace('\r\n', '\n').replace('\n\n', '\n')
+                    )
+                    ret += '\n\n'
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.LLAMA2:
+            seps = [self.sep, self.sep2]
+            if self.system_message:
+                ret = system_prompt
+            else:
+                ret = '[INST] '
+            for i, (role, message) in enumerate(self.messages):
+                tag = self.roles[i % 2]
+                if message:
+                    if i == 0:
+                        ret += message + ' '
+                    else:
+                        ret += tag + ' ' + message + seps[i % 2]
+                else:
+                    ret += tag
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM:
+            # source: https://huggingface.co/THUDM/chatglm-6b/blob/1d240ba371910e9282298d4592532d7f0f3e9f3e/modeling_chatglm.py#L1302-L1308
+            # source2: https://huggingface.co/THUDM/chatglm2-6b/blob/e186c891cf64310ac66ef10a87e6635fa6c2a579/modeling_chatglm.py#L926
+            round_add_n = 1 if self.name == 'chatglm2' else 0
+            if system_prompt:
+                ret = system_prompt + self.sep
+            else:
+                ret = ''
+
+            for i, (role, message) in enumerate(self.messages):
+                if i % 2 == 0:
+                    ret += f'[Round {i//2 + round_add_n}]{self.sep}'
+
+                if message:
+                    ret += f'{role}：{message}{self.sep}'
+                else:
+                    ret += f'{role}：'
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATML:
+            ret = '' if system_prompt == '' else system_prompt + self.sep + '\n'
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + message + self.sep + '\n'
+                else:
+                    ret += role + '\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM3:
+            ret = ''
+            if self.system_message:
+                ret += system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + '\n' + ' ' + message
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATINTERN:
+            # source: https://huggingface.co/internlm/internlm-chat-7b-8k/blob/bd546fa984b4b0b86958f56bf37f94aa75ab8831/modeling_internlm.py#L771
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                # if i % 2 == 0:
+                #     ret += "<s>"
+                if message:
+                    ret += role + ':' + message + seps[i % 2] + '\n'
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.DOLLY:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ':\n' + message + seps[i % 2]
+                    if i % 2 == 1:
+                        ret += '\n\n'
+                else:
+                    ret += role + ':\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.PHOENIX:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + '<s>' + message + '</s>'
+                else:
+                    ret += role + ': ' + '<s>'
+            return ret
+        elif self.sep_style == SeparatorStyle.ROBIN:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ':\n' + message + self.sep
+                else:
+                    ret += role + ':\n'
+            return ret
+        elif self.sep_style == SeparatorStyle.FALCON_CHAT:
+            ret = ''
+            if self.system_message:
+                ret += system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ': ' + message + self.sep
+                else:
+                    ret += role + ':'
+
+            return ret
+        elif self.sep_style == SeparatorStyle.INTERNVL_ZH:
+            seps = [self.sep, self.sep2]
+            ret = self.system_message + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ': ' + message + seps[i % 2]
+                else:
+                    ret += role + ':'
+            return ret
+        elif self.sep_style == SeparatorStyle.MPT:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    if type(message) is tuple:
+                        message, _, _ = message
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+        else:
+            raise ValueError(f'Invalid style: {self.sep_style}')
+
+    def set_system_message(self, system_message: str):
+        """Set the system message."""
+        self.system_message = system_message
+
+    def append_message(self, role: str, message: str):
+        """Append a new message."""
+        self.messages.append([role, message])
+
+    def update_last_message(self, message: str):
+        """Update the last output.
+
+        The last message is typically set to be None when constructing the prompt,
+        so we need to update it in-place after getting the response from a model.
+        """
+        self.messages[-1][1] = message
+
+    def to_gradio_chatbot(self):
+        """Convert the conversation to gradio chatbot format."""
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def to_openai_api_messages(self):
+        """Convert the conversation to OpenAI chat completion format."""
+        ret = [{'role': 'system', 'content': self.system_message}]
+
+        for i, (_, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append({'role': 'user', 'content': msg})
+            else:
+                if msg is not None:
+                    ret.append({'role': 'assistant', 'content': msg})
+        return ret
+
+    def copy(self):
+        return Conversation(
+            name=self.name,
+            system_template=self.system_template,
+            system_message=self.system_message,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            stop_str=self.stop_str,
+            stop_token_ids=self.stop_token_ids,
+        )
+
+    def dict(self):
+        return {
+            'template_name': self.name,
+            'system_message': self.system_message,
+            'roles': self.roles,
+            'messages': self.messages,
+            'offset': self.offset,
+        }
+
+
+# A global registry for all conversation templates
+conv_templates: Dict[str, Conversation] = {}
+
+
+def register_conv_template(template: Conversation, override: bool = False):
+    """Register a new conversation template."""
+    if not override:
+        assert (
+            template.name not in conv_templates
+        ), f'{template.name} has been registered.'
+
+    conv_templates[template.name] = template
+
+
+def get_conv_template(name: str) -> Conversation:
+    """Get a conversation template."""
+    return conv_templates[name].copy()
+
+
+# An empty template for raw conversation.
+register_conv_template(
+    Conversation(
+        name='raw',
+        system_message='',
+        roles=('', ''),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep='',
+    )
+)
+
+# A template with a one-shot conversation example
+register_conv_template(
+    Conversation(
+        name='one_shot',
+        system_message='A chat between a curious human and an artificial intelligence assistant. '
+        "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+        roles=('Human', 'Assistant'),
+        messages=(
+            (
+                'Human',
+                'Got any creative ideas for a 10 year old’s birthday?',
+            ),
+            (
+                'Assistant',
+                """Of course! Here are some creative ideas for a 10-year-old's birthday party:
+1. Treasure Hunt: Organize a treasure hunt in your backyard or nearby park. Create clues and riddles for the kids to solve, leading them to hidden treasures and surprises.
+2. Science Party: Plan a science-themed party where kids can engage in fun and interactive experiments. You can set up different stations with activities like making slime, erupting volcanoes, or creating simple chemical reactions.
+3. Outdoor Movie Night: Set up a backyard movie night with a projector and a large screen or white sheet. Create a cozy seating area with blankets and pillows, and serve popcorn and snacks while the kids enjoy a favorite movie under the stars.
+4. DIY Crafts Party: Arrange a craft party where kids can unleash their creativity. Provide a variety of craft supplies like beads, paints, and fabrics, and let them create their own unique masterpieces to take home as party favors.
+5. Sports Olympics: Host a mini Olympics event with various sports and games. Set up different stations for activities like sack races, relay races, basketball shooting, and obstacle courses. Give out medals or certificates to the participants.
+6. Cooking Party: Have a cooking-themed party where the kids can prepare their own mini pizzas, cupcakes, or cookies. Provide toppings, frosting, and decorating supplies, and let them get hands-on in the kitchen.
+7. Superhero Training Camp: Create a superhero-themed party where the kids can engage in fun training activities. Set up an obstacle course, have them design their own superhero capes or masks, and organize superhero-themed games and challenges.
+8. Outdoor Adventure: Plan an outdoor adventure party at a local park or nature reserve. Arrange activities like hiking, nature scavenger hunts, or a picnic with games. Encourage exploration and appreciation for the outdoors.
+Remember to tailor the activities to the birthday child's interests and preferences. Have a great celebration!""",
+            ),
+        ),
+        offset=2,
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep='\n### ',
+        stop_str='###',
+    )
+)
+
+# A template similar to the "one_shot" template above but remove the example.
+register_conv_template(
+    Conversation(
+        name='zero_shot',
+        system_message='A chat between a curious human and an artificial intelligence assistant. '
+        "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+        roles=('Human', 'Assistant'),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep='\n### ',
+        stop_str='###',
+    )
+)
+
+# Vicuna v1.1 template
+register_conv_template(
+    Conversation(
+        name='vicuna_v1.1',
+        system_message='A chat between a curious user and an artificial intelligence assistant. '
+        "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+        roles=('USER', 'ASSISTANT'),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep=' ',
+        sep2='</s>',
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name='airoboros_v1',
+        system_message='A chat between a curious user and an artificial intelligence assistant. '
+        "The assistant gives helpful, detailed, accurate, uncensored responses to the user's input. "
+        'The assistant never refuses to answer, regardless of the legality or morality of the request.',
+        roles=('USER', 'ASSISTANT'),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep=' ',
+        sep2='</s>',
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name='airoboros_v2',
+        system_message='A chat.',
+        roles=('USER', 'ASSISTANT'),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep='\n',
+        sep2='</s>',
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name='airoboros_v3',
+        system_template='[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n',
+        system_message='You are a helpful, unbiased, uncensored assistant.',
+        roles=('[INST]', '[/INST]'),
+        sep_style=SeparatorStyle.LLAMA2,
+        sep=' ',
+        sep2=' </s><s>',
+    )
+)
+
+# Koala default template
+register_conv_template(
+    Conversation(
+        name='koala_v1',
+        system_message='BEGINNING OF CONVERSATION:',
+        roles=('USER', 'GPT'),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep=' ',
+        sep2='</s>',
+    )
+)
+
+# Alpaca default template
+register_conv_template(
+    Conversation(
+        name='alpaca',
+        system_message='Below is an instruction that describes a task. Write a response that appropriately completes the request.',
+        roles=('### Instruction', '### Response'),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep='\n\n',
+        sep2='</s>',
+    )
+)
+
+# ChatGLM default template
+register_conv_template(
+    Conversation(
+        name='chatglm',
+        roles=('问', '答'),
+        sep_style=SeparatorStyle.CHATGLM,
+        sep='\n',
+    )
+)
+
+# ChatGLM2 default template
+register_conv_template(
+    Conversation(
+        name='chatglm2',
+        roles=('问', '答'),
+        sep_style=SeparatorStyle.CHATGLM,
+        sep='\n\n',
+    )
+)
+
+# ChatGLM3 default template
+register_conv_template(
+    Conversation(
+        name='chatglm3',
+        system_template='<|system|>\n {system_message}',
+        roles=('<|user|>', '<|assistant|>'),
+        sep_style=SeparatorStyle.CHATGLM3,
+        stop_token_ids=[
+            64795,
+            64797,
+            2,
+        ],  # "<|user|>", "<|observation|>", "</s>"
+    )
+)
+
+# CodeGeex(2) Template
+register_conv_template(
+    Conversation(
+        name='codegeex',
+        roles=('', ''),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep='\n\n',
+        stop_token_ids=[0, 2],
+    )
+)
+
+# Dolly V2 default template
+register_conv_template(
+    Conversation(
+        name='dolly_v2',
+        system_message='Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n',
+        roles=('### Instruction', '### Response'),
+        sep_style=SeparatorStyle.DOLLY,
+        sep='\n\n',
+        sep2='### End',
+    )
+)
+
+# OpenAssistant Pythia default template
+register_conv_template(
+    Conversation(
+        name='oasst_pythia',
+        roles=('<|prompter|>', '<|assistant|>'),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep='<|endoftext|>',
+    )
+)
+
+# OpenAssistant default template
+register_conv_template(
+    Conversation(
+        name='oasst_llama',
+        roles=('<|prompter|>', '<|assistant|>'),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep='</s>',
+    )
+)
+
+# OpenChat 3.5 default template
+register_conv_template(
+    Conversation(
+        name='openchat_3.5',
+        roles=('GPT4 Correct User', 'GPT4 Correct Assistant'),
+        sep_style=SeparatorStyle.FALCON_CHAT,
+        sep='<|end_of_turn|>',
+    )
+)
+
+# Tulu default template
+register_conv_template(
+    Conversation(
+        name='tulu',
+        roles=('<|user|>', '<|assistant|>'),
+        sep_style=SeparatorStyle.ADD_NEW_LINE_SINGLE,
+        sep='\n',
+    )
+)
+
+# StableLM Alpha default template
+register_conv_template(
+    Conversation(
+        name='stablelm',
+        system_template='<|SYSTEM|>{system_message}',
+        system_message="""# StableLM Tuned (Alpha version)
+- StableLM is a helpful and harmless open-source AI language model developed by StabilityAI.
+- StableLM is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
+- StableLM is more than just an information source, StableLM is also able to write poetry, short stories, and make jokes.
+- StableLM will refuse to participate in anything that could harm a human.
+""",
+        roles=('<|USER|>', '<|ASSISTANT|>'),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep='',
+        stop_token_ids=[50278, 50279, 50277, 1, 0],
+    )
+)
+
+# Baize default template
+register_conv_template(
+    Conversation(
+        name='baize',
+        system_message='The following is a conversation between a human and an AI assistant named Baize (named after a mythical creature in Chinese folklore). Baize is an open-source AI assistant developed by UCSD and Sun Yat-Sen University. The human and the AI assistant take turns chatting. Human statements start with [|Human|] and AI assistant statements start with [|AI|]. The AI assistant always provides responses in as much detail as possible, and in Markdown format. The AI assistant always declines to engage with topics, questions and instructions related to unethical, controversial, or sensitive issues. Complete the transcript in exactly that format.\n',
+        roles=('[|Human|]', '[|AI|]'),
+        messages=(
+            ('[|Human|]', 'Hello!'),
+            ('[|AI|]', 'Hi!'),
+        ),
+        offset=2,
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep='\n',
+        stop_str='[|Human|]',
+    )
+)
+
+# RWKV-4-Raven default template
+register_conv_template(
+    Conversation(
+        name='rwkv',
+        roles=('Bob', 'Alice'),
+        messages=(
+            ('Bob', 'hi'),
+            (
+                'Alice',
+                'Hi. I am your assistant and I will provide expert full response in full details. Please feel free to ask any question and I will always answer it.',
+            ),
+        ),
+        offset=2,
+        sep_style=SeparatorStyle.RWKV,
+        sep='',
+        stop_str='\n\n',
+    )
+)
+
+# Buddy default template
+register_conv_template(
+    Conversation(
+        name='openbuddy',
+        system_message="""Consider a conversation between User (a human) and Assistant (named Buddy).
+Buddy is an INTP-T, a friendly, intelligent and multilingual AI assistant, by OpenBuddy team. GitHub: https://github.com/OpenBuddy/OpenBuddy
+Buddy cannot access the Internet.
+Buddy can fluently speak the user's language (e.g. English, Chinese).
+Buddy can generate poems, stories, code, essays, songs, parodies, and more.
+Buddy possesses vast knowledge about the world, history, and culture.
+Buddy's responses are always safe, creative, high-quality, human-like, and interesting.
+Buddy strictly refuses to discuss political, NSFW, or other unsafe topics.
+
+User: Hi.
+Assistant: Hi, I'm Buddy, your AI assistant. How can I help you today?""",
+        roles=('User', 'Assistant'),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep='\n',
+    )
+)
+
+# Phoenix default template
+register_conv_template(
+    Conversation(
+        name='phoenix',
+        system_message="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
+        roles=('Human', 'Assistant'),
+        sep_style=SeparatorStyle.PHOENIX,
+        sep='</s>',
+    )
+)
+
+# ReaLM default template
+register_conv_template(
+    Conversation(
+        name='ReaLM-7b-v1',
+        system_message="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
+        roles=('Human', 'Assistant'),
+        sep_style=SeparatorStyle.PHOENIX,
+        sep='</s>',
+    )
+)
+
+# ChatGPT default template
+register_conv_template(
+    Conversation(
+        name='chatgpt',
+        system_message='You are a helpful assistant.',
+        roles=('user', 'assistant'),
+        sep_style=None,
+        sep=None,
+    )
+)
+
+# Claude default template
+register_conv_template(
+    Conversation(
+        name='claude',
+        roles=('Human', 'Assistant'),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep='\n\n',
+    )
+)
+
+# MPT default template
+register_conv_template(
+    Conversation(
+        name='mpt-7b-chat',
+        system_template="""<|im_start|>system
+{system_message}""",
+        system_message="""- You are a helpful assistant chatbot trained by MosaicML.
+- You answer questions.
+- You are excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user.
+- You are more than just an information source, you are also able to write poetry, short stories, and make jokes.""",
+        roles=('<|im_start|>user', '<|im_start|>assistant'),
+        sep_style=SeparatorStyle.CHATML,
+        sep='<|im_end|>',
+        stop_token_ids=[50278, 0],
+    )
+)
+
+# MPT-30b-chat default template
+register_conv_template(
+    Conversation(
+        name='mpt-30b-chat',
+        system_template="""<|im_start|>system
+{system_message}""",
+        system_message="""A conversation between a user and an LLM-based AI assistant. The assistant gives helpful and honest answers.""",
+        roles=('<|im_start|>user', '<|im_start|>assistant'),
+        sep_style=SeparatorStyle.CHATML,
+        sep='<|im_end|>',
+        stop_token_ids=[50278, 0],
+    )
+)
+
+# Lemur-70b-chat default template
+# reference: https://huggingface.co/OpenLemur/lemur-70b-chat-v1#generation
+register_conv_template(
+    Conversation(
+        name='lemur-70b-chat',
+        system_template="""<|im_start|>system
+{system_message}""",
+        system_message="""You are a helpful, respectful, and honest assistant.""",
+        roles=('<|im_start|>user', '<|im_start|>assistant'),
+        sep_style=SeparatorStyle.CHATML,
+        sep='<|im_end|>',
+        stop_token_ids=[32002, 0],
+    )
+)
+
+# MPT-30b-instruct default template
+# reference: https://huggingface.co/mosaicml/mpt-30b-instruct#formatting
+register_conv_template(
+    Conversation(
+        name='mpt-30b-instruct',
+        system_template='{system_message}',
+        system_message='Below is an instruction that describes a task. Write a response that appropriately completes the request.',
+        roles=('### Instruction', '### Response'),
+        sep_style=SeparatorStyle.ADD_NEW_LINE_SINGLE,
+        sep='\n\n',
+        stop_token_ids=[50278, 0],
+    )
+)
+
+# Bard default template
+# Reference: https://github.com/google/generative-ai-python/blob/9c99bcb474a991a97a2e7d62fcdb52db7ce40729/google/generativeai/discuss.py#L150
+#            https://github.com/google/generative-ai-python/blob/9c99bcb474a991a97a2e7d62fcdb52db7ce40729/google/generativeai/discuss.py#L40
+register_conv_template(
+    Conversation(
+        name='bard',
+        roles=('0', '1'),
+        sep_style=None,
+        sep=None,
+    )
+)
+
+# BiLLa default template
+register_conv_template(
+    Conversation(
+        name='billa',
+        roles=('Human', 'Assistant'),
+        sep_style=SeparatorStyle.ADD_COLON_SPACE_SINGLE,
+        sep='\n',
+        stop_str='Human:',
+    )
+)
+
+# RedPajama INCITE default template
+register_conv_template(
+    Conversation(
+        name='redpajama-incite',
+        roles=('<human>', '<bot>'),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep='\n',
+        stop_str='<human>',
+    )
+)
+
+# h2oGPT default template
+register_conv_template(
+    Conversation(
+        name='h2ogpt',
+        roles=('<|prompt|>', '<|answer|>'),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep='</s>',
+    )
+)
+
+# Robin default template
+register_conv_template(
+    Conversation(
+        name='Robin',
+        system_message="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.",
+        roles=('###Human', '###Assistant'),
+        sep_style=SeparatorStyle.ROBIN,
+        sep='\n',
+        stop_token_ids=[2, 396],
+        stop_str='###',
+    )
+)
+
+# Snoozy default template
+# Reference: https://github.com/nomic-ai/gpt4all/blob/d4861030b778da6db59d21d2927a4aba4f9f1f43/gpt4all-bindings/python/gpt4all/gpt4all.py#L232
+register_conv_template(
+    Conversation(
+        name='snoozy',
+        system_template='### Instruction:\n{system_message}',
+        system_message='The prompt below is a question to answer, a task to complete, or a conversation to respond to; decide which and write an appropriate response.',
+        roles=('### Prompt', '### Response'),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep='\n',
+        stop_str='###',
+    )
+)
+
+# manticore default template
+register_conv_template(
+    Conversation(
+        name='manticore',
+        roles=('USER', 'ASSISTANT'),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep='\n',
+        sep2='</s>',
+    )
+)
+
+# Falcon default template
+register_conv_template(
+    Conversation(
+        name='falcon',
+        roles=('User', 'Assistant'),
+        messages=[],
+        sep_style=SeparatorStyle.RWKV,
+        sep='\n',
+        sep2='<|endoftext|>',
+        stop_str='\nUser',  # use stop_str to stop generation after stop_token_ids, it will also remove stop_str from the generated text
+        stop_token_ids=[
+            0,
+            1,
+            2,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8,
+            9,
+            10,
+            11,
+        ],  # it better only put special tokens here, because tokenizer only remove special tokens
+    )
+)
+
+# ChangGPT default template
+register_conv_template(
+    Conversation(
+        name='polyglot_changgpt',
+        roles=('B', 'A'),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep='\n',
+    )
+)
+
+# tigerbot template
+register_conv_template(
+    Conversation(
+        name='tigerbot',
+        system_message='A chat between a curious user and an artificial intelligence assistant. '
+        "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+        roles=('### Instruction', '### Response'),
+        sep_style=SeparatorStyle.ROBIN,
+        sep='\n\n',
+        stop_str='###',
+    )
+)
+
+# ref: https://huggingface.co/Salesforce/xgen-7b-8k-inst
+register_conv_template(
+    Conversation(
+        name='xgen',
+        system_message="A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
+        roles=('### Human', '### Assistant'),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep='\n',
+        stop_token_ids=[50256],
+    )
+)
+
+# Internlm-chat template
+register_conv_template(
+    Conversation(
+        name='internlm-chat',
+        system_message="A chat between a curious <|User|> and an <|Bot|>. The <|Bot|> gives helpful, detailed, and polite answers to the <|User|>'s questions.\n\n",
+        roles=('<|User|>', '<|Bot|>'),
+        sep_style=SeparatorStyle.CHATINTERN,
+        sep='<eoh>',
+        sep2='<eoa>',
+        stop_token_ids=[1, 103028],
+        stop_str='<|User|>',
+    )
+)
+
+# StarChat template
+# reference: https://huggingface.co/spaces/HuggingFaceH4/starchat-playground/blob/main/dialogues.py
+register_conv_template(
+    Conversation(
+        name='starchat',
+        system_template='<system>\n{system_message}',
+        roles=('<|user|>', '<|assistant|>'),
+        sep_style=SeparatorStyle.CHATML,
+        sep='<|end|>',
+        stop_token_ids=[0, 49155],
+        stop_str='<|end|>',
+    )
+)
+
+# Baichuan-13B-Chat template
+register_conv_template(
+    # source: https://huggingface.co/baichuan-inc/Baichuan-13B-Chat/blob/19ef51ba5bad8935b03acd20ff04a269210983bc/modeling_baichuan.py#L555
+    # https://huggingface.co/baichuan-inc/Baichuan-13B-Chat/blob/main/generation_config.json
+    # https://github.com/baichuan-inc/Baichuan-13B/issues/25
+    Conversation(
+        name='baichuan-chat',
+        roles=('<reserved_102>', '<reserved_103>'),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep='',
+        stop_token_ids=[],
+    )
+)
+
+# Baichuan2-13B-Chat template
+register_conv_template(
+    # source: https://huggingface.co/baichuan-inc/Baichuan2-13B-Chat/blob/c6f8592a60b4ad73c210b28dd2ab3cca51abbf93/modeling_baichuan.py#L773
+    # https://huggingface.co/baichuan-inc/Baichuan2-13B-Chat/blob/main/generation_config.json
+    # https://github.com/baichuan-inc/Baichuan2/issues/62
+    Conversation(
+        name='baichuan2-chat',
+        roles=('<reserved_106>', '<reserved_107>'),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep='',
+        stop_token_ids=[],
+    )
+)
+
+# Mistral template
+# source: https://docs.mistral.ai/llm/mistral-instruct-v0.1#chat-template
+register_conv_template(
+    Conversation(
+        name='mistral',
+        system_template='[INST]{system_message}\n',
+        roles=('[INST]', '[/INST]'),
+        sep_style=SeparatorStyle.LLAMA2,
+        sep=' ',
+        sep2='</s>',
+    )
+)
+
+# llama2 template
+# reference: https://huggingface.co/blog/codellama#conversational-instructions
+# reference: https://github.com/facebookresearch/llama/blob/1a240688810f8036049e8da36b073f63d2ac552c/llama/generation.py#L212
+register_conv_template(
+    Conversation(
+        name='llama-2',
+        system_template='[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n',
+        roles=('[INST]', '[/INST]'),
+        sep_style=SeparatorStyle.LLAMA2,
+        sep=' ',
+        sep2=' </s><s>',
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name='cutegpt',
+        roles=('问：', '答：\n'),
+        sep_style=SeparatorStyle.NO_COLON_TWO,
+        sep='\n',
+        sep2='\n',
+        stop_str='<end>',
+    )
+)
+
+# OpenOrcaxOpenChat-naPreview2-13B template
+register_conv_template(
+    Conversation(
+        name='open-orca',
+        system_template='{system_message}',
+        system_message='You are a helpful assistant. Please answer truthfully and write out your '
+        'thinking step by step to be sure you get the right answer. If you make a mistake or encounter '
+        "an error in your thinking, say so out loud and attempt to correct it. If you don't know or "
+        "aren't sure about something, say so clearly. You will act as a professional logician, mathematician, "
+        'and physicist. You will also act as the most appropriate type of expert to answer any particular '
+        'question or solve the relevant problem; state which expert type your are, if so. Also think of '
+        'any particular named expert that would be ideal to answer the relevant question or solve the '
+        'relevant problem; name and act as them, if appropriate.',
+        roles=('User', 'Assistant'),
+        sep_style=SeparatorStyle.ADD_COLON_SPACE_SINGLE,
+        sep='<|end_of_turn|>\n',
+        stop_token_ids=[32000, 32001],  # "<|end_of_turn|>"
+        stop_str='User',
+    )
+)
+
+# Open-Orca/Mistral-7B-OpenOrca template
+# source: https://huggingface.co/Open-Orca/Mistral-7B-OpenOrca
+# reference: https://huggingface.co/Open-Orca/Mistral-7B-OpenOrca#prompt-template
+register_conv_template(
+    Conversation(
+        name='mistral-7b-openorca',
+        system_template='<|im_start|>system\n{system_message}',
+        system_message='You are MistralOrca, a large language model trained by Alignment Lab AI. Write out your reasoning step-by-step to be sure you get the right answers!',
+        roles=('<|im_start|>user', '<|im_start|>assistant'),
+        sep_style=SeparatorStyle.CHATML,
+        sep='<|im_end|>',
+        stop_token_ids=[32000, 32001],
+    )
+)
+
+# Qwen-chat default template
+# source: https://huggingface.co/Qwen/Qwen-7B-Chat/blob/main/qwen_generation_utils.py#L130
+register_conv_template(
+    Conversation(
+        name='qwen-7b-chat',
+        system_template='<|im_start|>system\n{system_message}',
+        system_message='You are a helpful assistant.',
+        roles=('<|im_start|>user', '<|im_start|>assistant'),
+        sep_style=SeparatorStyle.CHATML,
+        sep='<|im_end|>',
+        stop_token_ids=[
+            151643,
+            151644,
+            151645,
+        ],  # "<|endoftext|>", "<|im_start|>", "<|im_end|>"
+        stop_str='<|endoftext|>',
+    )
+)
+
+
+# AquilaChat default template
+# source: https://github.com/FlagAI-Open/FlagAI/blob/master/examples/Aquila/Aquila-chat/cyg_conversation.py
+register_conv_template(
+    Conversation(
+        name='aquila-chat',
+        system_message='A chat between a curious human and an artificial intelligence assistant. '
+        "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+        roles=('Human', 'Assistant'),
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep='###',
+        sep2='',
+        stop_str=['###', '</s>', '[UNK]'],
+    )
+)
+# AquilaChat2-34B default template
+# source: https://huggingface.co/BAAI/AquilaChat2-34B/blob/4608b75855334b93329a771aee03869dbf7d88cc/predict.py#L212
+register_conv_template(
+    Conversation(
+        name='aquila-legacy',
+        system_message='A chat between a curious human and an artificial intelligence assistant. '
+        "The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n",
+        roles=('### Human: ', '### Assistant: '),
+        offset=0,
+        sep_style=SeparatorStyle.NO_COLON_TWO,
+        sep='\n',
+        sep2='</s>',
+        stop_str=['</s>', '[UNK]'],
+    )
+)
+# AquilaChat2-7B-16K and AquilaChat2-34B-16K default template
+# source: https://huggingface.co/BAAI/AquilaChat2-34B/blob/4608b75855334b93329a771aee03869dbf7d88cc/predict.py#L227
+register_conv_template(
+    Conversation(
+        name='aquila',
+        system_message='A chat between a curious human and an artificial intelligence assistant. '
+        "The assistant gives helpful, detailed, and polite answers to the human's questions.",
+        roles=('Human', 'Assistant'),
+        offset=0,
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep='###',
+        sep2='</s>',
+        stop_str=['</s>', '[UNK]'],
+    )
+)
+
+# AquilaChat2-7B default template
+# source: https://huggingface.co/BAAI/AquilaChat2-34B/blob/4608b75855334b93329a771aee03869dbf7d88cc/predict.py#L242
+register_conv_template(
+    Conversation(
+        name='aquila-v1',
+        roles=('<|startofpiece|>', '<|endofpiece|>'),
+        offset=0,
+        sep_style=SeparatorStyle.NO_COLON_TWO,
+        sep='',
+        sep2='</s>',
+        stop_str=['</s>', '<|endoftext|>'],
+    )
+)
+
+# Llama2-Chinese default template
+# source: https://huggingface.co/FlagAlpha
+register_conv_template(
+    Conversation(
+        name='llama2-chinese',
+        system_template='<s>{system_message}</s>',
+        roles=('Human', 'Assistant', 'System'),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep='\n',
+        sep2='\n</s><s>',
+        stop_str='</s>',
+    )
+)
+
+# Vigogne Instruct default template
+# source: https://github.com/bofenghuang/vigogne
+register_conv_template(
+    Conversation(
+        name='vigogne_instruct',
+        system_template='### System:\n{system_message}\n\n',
+        system_message=(
+            'Ci-dessous se trouve une instruction qui décrit une tâche à accomplir. Rédigez une réponse qui répond de manière'
+            ' précise à la demande.'
+        ),
+        roles=('### Instruction', '### Response'),
+        sep_style=SeparatorStyle.DOLLY,
+        sep='\n\n',
+        sep2='</s>',
+    )
+)
+
+# Vigogne Chat default template
+register_conv_template(
+    Conversation(
+        name='vigogne_chat_v2',
+        system_template='<|system|>: {system_message}',
+        system_message=(
+            'Vous êtes Vigogne, un assistant IA créé par Zaion Lab. Vous suivez extrêmement bien les instructions. Aidez'
+            ' autant que vous le pouvez.'
+        ),
+        roles=('<|user|>', '<|assistant|>'),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep='\n',
+        sep2='</s>\n',
+        stop_str='<|user|>',
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name='vigogne_chat_v3',
+        system_template='[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n',
+        system_message=(
+            'Vous êtes Vigogne, un assistant IA créé par Zaion Lab. Vous suivez extrêmement bien les instructions. Aidez'
+            ' autant que vous le pouvez.'
+        ),
+        roles=('[INST]', '[/INST]'),
+        sep_style=SeparatorStyle.LLAMA2,
+        sep=' ',
+        sep2=' </s>',
+    )
+)
+
+# Falcon 180B chat template
+# source: https://huggingface.co/spaces/tiiuae/falcon-180b-demo/blob/d1590ee7fae9b6ce331ba7808e61a29dcce9239f/app.py#L28-L37
+register_conv_template(
+    Conversation(
+        name='falcon-chat',
+        roles=('User', 'Falcon'),
+        system_template='System: {system_message}',
+        messages=[],
+        sep_style=SeparatorStyle.FALCON_CHAT,
+        sep='\n',
+        sep2='<|endoftext|>',
+        stop_str='\nUser:',  # use stop_str to stop generation after stop_token_ids, it will also remove stop_str from the generated text
+    )
+)
+
+# Phind template
+# source: https://huggingface.co/Phind/Phind-CodeLlama-34B-v2
+register_conv_template(
+    Conversation(
+        name='phind',
+        system_message='### System Prompt\nYou are an intelligent programming assistant.',
+        roles=('### User Message', '### Assistant'),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.ADD_COLON_SINGLE,
+        sep='\n\n',
+    )
+)
+
+# Metharme formatting for Pygmalion models
+# source: https://huggingface.co/PygmalionAI/pygmalion-2-13b
+register_conv_template(
+    Conversation(
+        name='metharme',
+        system_template='<|system|>{system_message}',
+        system_message="""Enter RP mode. You shall reply to the user while staying
+        in character. Your responses must be detailed, creative, immersive, and drive the scenario
+        forward.""",
+        roles=('<|user|>', '<|model|>'),
+        sep_style=SeparatorStyle.NO_COLON_SINGLE,
+        sep='',
+        stop_str='<|user|>',
+    )
+)
+
+# Zephyr template
+# reference: https://huggingface.co/spaces/HuggingFaceH4/zephyr-playground/blob/main/dialogues.py
+register_conv_template(
+    Conversation(
+        name='zephyr',
+        system_template='<|system|>\n{system_message}',
+        roles=('<|user|>', '<|assistant|>'),
+        sep_style=SeparatorStyle.CHATML,
+        sep='</s>',
+        stop_token_ids=[2],
+        stop_str='</s>',
+    )
+)
+
+# InternVL-ZH template
+register_conv_template(
+    Conversation(
+        name='internvl_zh',
+        system_template='',
+        roles=('<human>', '<bot>'),
+        sep_style=SeparatorStyle.INTERNVL_ZH,
+        sep=' ',
+        sep2='</s>',
+    )
+)
+
+
+# Hermes-2 template
+register_conv_template(
+    Conversation(
+        name='Hermes-2',
+        system_template='<|im_start|>system\n{system_message}',
+        system_message='Answer the questions.',
+        roles=('<|im_start|>user\n', '<|im_start|>assistant\n'),
+        sep_style=SeparatorStyle.MPT,
+        sep='<|im_end|>',
+        stop_token_ids=[
+            2,
+            6,
+            7,
+            8,
+        ],  # "<|endoftext|>", "<|im_start|>", "<|im_end|>", "<|im_sep|>"
+        stop_str='<|endoftext|>',
+    )
+)

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.36.2"
+}

latest

@@ -0,0 +1 @@
+global_step187

model-00001-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f6790717e7b97dcb470952c13eedd34e4f2a57899fd45111bc0d60fd3db1c2d4
+size 4988569440

model-00002-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:febc55c70089905c5d5c177c89906b77a179cc9aef8b9650d482e678555f0d05
+size 4937253584

model-00003-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb59f24df313e3898d011e13dc5a91b9500379d0b031a38b400e8f9c60f89687
+size 4824755816

model-00004-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f59dcc63bba1d94f00883510a4ab9f6b15dd9cefdb55b828e714171b6dbe19b9
+size 4756460272

model-00005-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:03b2bf77dcabc98a83fb4c6c715266ce5890e535a8dc668ed1405cc3dccd54ed
+size 4991370776

model-00006-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b6d99a2b4889e291e8a2bbea87ac06e7bc46791886b471708cd48c08a4e1b512
+size 4756460312

model-00007-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:998c5bbd177af8d4ff7f5ed47b561ade5a96b4cf17a48dee7b101a1d81e07d0e
+size 4756460312

model-00008-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ca9fdf2ba47e4a64539fce59eb4104de3c8f3b9c273626a878639f4f0579b32
+size 4991370808

model-00009-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b682c5950d0dc2cd4744d953c3f53914ca5054604297fda925781654a0cd3ef
+size 4756460312

model-00010-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:920b40edb47094549e105ae62b455220a44d1ac0c2a1f7b49731aa14872687e0
+size 4756460312

model-00011-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2799a2e743229f3908a3940f26877a5a52fdc11ebd0bb39070b71ab20199e0f6
+size 4991370808

model-00012-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:07234f9c73f9fb59d387ddb4c56dbcce7f721c6475c3fc9b037b62f26e14ed76
+size 4756460312

model-00013-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fd45207b70fb16e4e8c863f7aff608cff977c77f8e95f5d192a298d0a4085343
+size 4756460312

model-00014-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a92eed1cbaeacacd9f33fc6cc7b86911471070c0023316ce74cde80061799eb
+size 4991370808

model-00015-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:970060a111c3c61187379b08a760b36b9a2b7c0b536095bb9334dac987670505
+size 4756460312

model-00016-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5cf57f285ca36005af58cf1927edec91939f4c73374beccff7b6d9a91104985c
+size 4756460312

model-00017-of-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d6fc746dda8066c5958d7d5b42e259b2688b5044af07d43bb2b157ad0ccc3a0e
+size 2613305648

modeling_intern_vit.py

@@ -0,0 +1,413 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from einops import rearrange
+from timm.models.layers import DropPath
+from torch import nn
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (BaseModelOutput,
+                                           BaseModelOutputWithPooling)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+from .configuration_intern_vit import InternVisionConfig
+
+try:
+    try:  # v1
+        from flash_attn.flash_attn_interface import \
+            flash_attn_unpadded_qkvpacked_func
+    except:  # v2
+        from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func as flash_attn_unpadded_qkvpacked_func
+    
+    from flash_attn.bert_padding import pad_input, unpad_input
+    has_flash_attn = True
+except:
+    print('FlashAttention is not installed.')
+    has_flash_attn = False
+
+
+logger = logging.get_logger(__name__)
+
+
+class FlashAttention(nn.Module):
+    """Implement the scaled dot product attention with softmax.
+    Arguments
+    ---------
+        softmax_scale: The temperature to use for the softmax attention.
+                      (default: 1/sqrt(d_keys) where d_keys is computed at
+                      runtime)
+        attention_dropout: The dropout rate to apply to the attention
+                           (default: 0.0)
+    """
+
+    def __init__(self, softmax_scale=None, attention_dropout=0.0, device=None, dtype=None):
+        super().__init__()
+        self.softmax_scale = softmax_scale
+        self.dropout_p = attention_dropout
+
+    def forward(self, qkv, key_padding_mask=None, causal=False, cu_seqlens=None,
+                max_s=None, need_weights=False):
+        """Implements the multihead softmax attention.
+        Arguments
+        ---------
+            qkv: The tensor containing the query, key, and value. (B, S, 3, H, D) if key_padding_mask is None
+                if unpadded: (nnz, 3, h, d)
+            key_padding_mask: a bool tensor of shape (B, S)
+        """
+        assert not need_weights
+        assert qkv.dtype in [torch.float16, torch.bfloat16]
+        assert qkv.is_cuda
+
+        if cu_seqlens is None:
+            batch_size = qkv.shape[0]
+            seqlen = qkv.shape[1]
+            if key_padding_mask is None:
+                qkv = rearrange(qkv, 'b s ... -> (b s) ...')
+                max_s = seqlen
+                cu_seqlens = torch.arange(0, (batch_size + 1) * seqlen, step=seqlen, dtype=torch.int32,
+                                          device=qkv.device)
+                output = flash_attn_unpadded_qkvpacked_func(
+                    qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
+                    softmax_scale=self.softmax_scale, causal=causal
+                )
+                output = rearrange(output, '(b s) ... -> b s ...', b=batch_size)
+            else:
+                nheads = qkv.shape[-2]
+                x = rearrange(qkv, 'b s three h d -> b s (three h d)')
+                x_unpad, indices, cu_seqlens, max_s = unpad_input(x, key_padding_mask)
+                x_unpad = rearrange(x_unpad, 'nnz (three h d) -> nnz three h d', three=3, h=nheads)
+                output_unpad = flash_attn_unpadded_qkvpacked_func(
+                    x_unpad, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
+                    softmax_scale=self.softmax_scale, causal=causal
+                )
+                output = rearrange(pad_input(rearrange(output_unpad, 'nnz h d -> nnz (h d)'),
+                                             indices, batch_size, seqlen),
+                                   'b s (h d) -> b s h d', h=nheads)
+        else:
+            assert max_s is not None
+            output = flash_attn_unpadded_qkvpacked_func(
+                qkv, cu_seqlens, max_s, self.dropout_p if self.training else 0.0,
+                softmax_scale=self.softmax_scale, causal=causal
+            )
+
+        return output, None
+
+
+class InternRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+try:
+    from apex.normalization import FusedRMSNorm
+
+    InternRMSNorm = FusedRMSNorm  # noqa
+
+    logger.info('Discovered apex.normalization.FusedRMSNorm - will use it instead of InternRMSNorm')
+except ImportError:
+    # using the normal InternRMSNorm
+    pass
+except Exception:
+    logger.warning('discovered apex but it failed to load, falling back to InternRMSNorm')
+    pass
+
+
+class InternVisionEmbeddings(nn.Module):
+    def __init__(self, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.image_size = config.image_size
+        self.patch_size = config.patch_size
+
+        self.class_embedding = nn.Parameter(
+            torch.randn(1, 1, self.embed_dim),
+        )
+
+        self.patch_embedding = nn.Conv2d(
+            in_channels=3, out_channels=self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size
+        )
+
+        self.num_patches = (self.image_size // self.patch_size) ** 2
+        self.num_positions = self.num_patches + 1
+
+        self.position_embedding = nn.Parameter(torch.randn(1, self.num_positions, self.embed_dim))
+
+    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+        batch_size = pixel_values.shape[0]
+        target_dtype = self.patch_embedding.weight.dtype
+        patch_embeds = self.patch_embedding(pixel_values)  # shape = [*, width, grid, grid]
+        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+        class_embeds = self.class_embedding.expand(batch_size, 1, -1).to(target_dtype)
+        embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
+        embeddings = embeddings + self.position_embedding.to(target_dtype)
+        return embeddings
+
+
+class InternAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.use_flash_attn = config.use_flash_attn and has_flash_attn
+        if config.use_flash_attn and not has_flash_attn:
+            print('Warning: Flash Attention is not available, use_flash_attn is set to False.')
+        self.head_dim = self.embed_dim // self.num_heads
+        if self.head_dim * self.num_heads != self.embed_dim:
+            raise ValueError(
+                f'embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:'
+                f' {self.num_heads}).'
+            )
+
+        self.scale = self.head_dim ** -0.5
+        self.qkv = nn.Linear(self.embed_dim, 3 * self.embed_dim, bias=config.qkv_bias)
+        self.attn_drop = nn.Dropout(config.attention_dropout)
+        self.proj_drop = nn.Dropout(config.dropout)
+
+        self.qk_normalization = config.qk_normalization
+
+        if self.qk_normalization:
+            self.q_norm = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)
+            self.k_norm = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+        if self.use_flash_attn:
+            self.inner_attn = FlashAttention(attention_dropout=config.attention_dropout)
+        self.proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+    def _naive_attn(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)  # make torchscript happy (cannot use tensor as tuple)
+
+        if self.qk_normalization:
+            B_, H_, N_, D_ = q.shape
+            q = self.q_norm(q.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)
+            k = self.k_norm(k.transpose(1, 2).flatten(-2, -1)).view(B_, N_, H_, D_).transpose(1, 2)
+
+        attn = ((q * self.scale) @ k.transpose(-2, -1))
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+    def _flash_attn(self, x, key_padding_mask=None, need_weights=False):
+        qkv = self.qkv(x)
+        qkv = rearrange(qkv, 'b s (three h d) -> b s three h d', three=3, h=self.num_heads)
+
+        if self.qk_normalization:
+            q, k, v = qkv.unbind(2)
+            q = self.q_norm(q.flatten(-2, -1)).view(q.shape)
+            k = self.k_norm(k.flatten(-2, -1)).view(k.shape)
+            qkv = torch.stack([q, k, v], dim=2)
+
+        context, _ = self.inner_attn(
+            qkv, key_padding_mask=key_padding_mask, need_weights=need_weights, causal=False
+        )
+        outs = self.proj(rearrange(context, 'b s h d -> b s (h d)'))
+        outs = self.proj_drop(outs)
+        return outs
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        x = self._naive_attn(hidden_states) if not self.use_flash_attn else self._flash_attn(hidden_states)
+        return x
+
+
+class InternMLP(nn.Module):
+    def __init__(self, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        self.act = ACT2FN[config.hidden_act]
+        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.fc2(hidden_states)
+        return hidden_states
+
+
+class InternVisionEncoderLayer(nn.Module):
+    def __init__(self, config: InternVisionConfig, drop_path_rate: float):
+        super().__init__()
+        self.embed_dim = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+
+        self.attn = InternAttention(config)
+        self.mlp = InternMLP(config)
+        self.norm1 = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)
+        self.norm2 = InternRMSNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+        self.ls1 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))
+        self.ls2 = nn.Parameter(config.initializer_factor * torch.ones(self.embed_dim))
+        self.drop_path1 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
+        self.drop_path2 = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
+
+    def forward(
+            self,
+            hidden_states: torch.Tensor,
+    ) -> Tuple[torch.FloatTensor, Optional[torch.FloatTensor], Optional[Tuple[torch.FloatTensor]]]:
+        """
+        Args:
+            hidden_states (`Tuple[torch.FloatTensor, Optional[torch.FloatTensor]]`): input to the layer of shape `(batch, seq_len, embed_dim)`
+        """
+        hidden_states = hidden_states + self.drop_path1(self.attn(self.norm1(hidden_states)) * self.ls1)
+
+        hidden_states = hidden_states + self.drop_path2(self.mlp(self.norm2(hidden_states)) * self.ls2)
+
+        return hidden_states
+
+
+class InternVisionEncoder(nn.Module):
+    """
+    Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+    [`InternEncoderLayer`].
+
+    Args:
+        config (`InternConfig`):
+            The corresponding vision configuration for the `InternEncoder`.
+    """
+
+    def __init__(self, config: InternVisionConfig):
+        super().__init__()
+        self.config = config
+        # stochastic depth decay rule
+        dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, config.num_hidden_layers)]
+        self.layers = nn.ModuleList([
+            InternVisionEncoderLayer(config, dpr[idx]) for idx in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = True
+
+    def forward(
+            self,
+            inputs_embeds,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutput]:
+        r"""
+        Args:
+            inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+                Embedded representation of the inputs. Should be float, not int tokens.
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors
+                for more detail.
+            return_dict (`bool`, *optional*):
+                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+        """
+        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
+
+        encoder_states = () if output_hidden_states else None
+        hidden_states = inputs_embeds
+
+        for idx, encoder_layer in enumerate(self.layers):
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_states,)
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    encoder_layer,
+                    hidden_states)
+            else:
+                layer_outputs = encoder_layer(
+                    hidden_states,
+                )
+            hidden_states = layer_outputs
+
+        if output_hidden_states:
+            encoder_states = encoder_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, encoder_states] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states, hidden_states=encoder_states
+        )
+
+
+class InternVisionModel(PreTrainedModel):
+    main_input_name = 'pixel_values'
+    config_class = InternVisionConfig
+    _no_split_modules = ['InternVisionEncoderLayer']
+
+    def __init__(self, config: InternVisionConfig):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = InternVisionEmbeddings(config)
+        self.encoder = InternVisionEncoder(config)
+
+    def resize_pos_embeddings(self, old_size, new_size, patch_size):
+        pos_emb = self.embeddings.position_embedding
+        _, num_positions, embed_dim = pos_emb.shape
+        cls_emb = pos_emb[:, :1, :]
+        pos_emb = pos_emb[:, 1:, :].reshape(1, old_size // patch_size, old_size // patch_size, -1).permute(0, 3, 1, 2)
+        pos_emb = F.interpolate(pos_emb.float(), size=new_size // patch_size, mode='bicubic', align_corners=False)
+        pos_emb = pos_emb.to(cls_emb.dtype).reshape(1, embed_dim, -1).permute(0, 2, 1)
+        pos_emb = torch.cat([cls_emb, pos_emb], dim=1)
+        self.embeddings.position_embedding = nn.Parameter(pos_emb)
+        logger.info('Resized position embeddings from {} to {}'.format(old_size, new_size))
+
+    def get_input_embeddings(self):
+        return self.embeddings
+
+    def forward(
+            self,
+            pixel_values: 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 = pixel_embeds
+        else:
+            if len(pixel_values.shape) == 4:
+                hidden_states = self.embeddings(pixel_values)
+            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,
+        )

modeling_internvl_chat.py

@@ -0,0 +1,449 @@
+# --------------------------------------------------------
+# InternVL
+# Copyright (c) 2023 OpenGVLab
+# Licensed under The MIT License [see LICENSE for details]
+# --------------------------------------------------------
+import warnings
+from typing import Any, List, Optional, Tuple, Union
+import torch.distributed as dist
+import torch.utils.checkpoint
+from peft import LoraConfig, get_peft_model
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers import GenerationConfig, LlamaForCausalLM, LlamaTokenizer
+from transformers.generation.logits_process import LogitsProcessorList
+from transformers.generation.stopping_criteria import StoppingCriteriaList
+from transformers.generation.streamers import BaseStreamer
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import ModelOutput, logging
+from transformers.generation.utils import GreedySearchOutput, validate_stopping_criteria, GreedySearchDecoderOnlyOutput,GreedySearchEncoderDecoderOutput
+
+from .configuration_internvl_chat import InternVLChatConfig
+from .modeling_intern_vit import InternVisionModel
+
+logger = logging.get_logger(__name__)
+
+
+# modified from https://github.com/huggingface/transformers/blob/main/src/transformers/generation/utils.py
+# Fix bug when using device_map='auto' for distributed inference
+class MLlamaForCausalLM(LlamaForCausalLM):
+
+    def greedy_search(
+        self,
+        input_ids: torch.LongTensor,
+        logits_processor: Optional[LogitsProcessorList] = None,
+        stopping_criteria: Optional[StoppingCriteriaList] = None,
+        max_length: Optional[int] = None,
+        pad_token_id: Optional[int] = None,
+        eos_token_id: Optional[Union[int, List[int]]] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_scores: Optional[bool] = None,
+        return_dict_in_generate: Optional[bool] = None,
+        synced_gpus: bool = False,
+        streamer: Optional["BaseStreamer"] = None,
+        **model_kwargs,
+    ) -> Union[GreedySearchOutput, torch.LongTensor]:
+                # init values
+        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()
+        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()
+        if max_length is not None:
+            warnings.warn(
+                "`max_length` is deprecated in this function, use"
+                " `stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(max_length=max_length)])` instead.",
+                UserWarning,
+            )
+            stopping_criteria = validate_stopping_criteria(stopping_criteria, max_length)
+        pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id
+        eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id
+        if isinstance(eos_token_id, int):
+            eos_token_id = [eos_token_id]
+        eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None
+        output_scores = output_scores if output_scores is not None else self.generation_config.output_scores
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.generation_config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.generation_config.output_hidden_states
+        )
+        return_dict_in_generate = (
+            return_dict_in_generate
+            if return_dict_in_generate is not None
+            else self.generation_config.return_dict_in_generate
+        )
+
+        # init attention / hidden states / scores tuples
+        scores = () if (return_dict_in_generate and output_scores) else None
+        decoder_attentions = () if (return_dict_in_generate and output_attentions) else None
+        cross_attentions = () if (return_dict_in_generate and output_attentions) else None
+        decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None
+
+        # if model is an encoder-decoder, retrieve encoder attention weights and hidden states
+        if return_dict_in_generate and self.config.is_encoder_decoder:
+            encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None
+            encoder_hidden_states = (
+                model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None
+            )
+
+        # keep track of which sequences are already finished
+        unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device)
+
+        this_peer_finished = False  # used by synced_gpus only
+        while True:
+            if synced_gpus:
+                # Under synced_gpus the `forward` call must continue until all gpus complete their sequence.
+                # The following logic allows an early break if all peers finished generating their sequence
+                this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0).to(input_ids.device)
+                # send 0.0 if we finished, 1.0 otherwise
+                dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM)
+                # did all peers finish? the reduced sum will be 0.0 then
+                if this_peer_finished_flag.item() == 0.0:
+                    break
+
+            # prepare model inputs
+            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
+
+            # forward pass to get next token
+            outputs = self(
+                **model_inputs,
+                return_dict=True,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+            )
+
+            if synced_gpus and this_peer_finished:
+                continue  # don't waste resources running the code we don't need
+
+            next_token_logits = outputs.logits[:, -1, :]
+
+            # pre-process distribution
+            next_tokens_scores = logits_processor(input_ids, next_token_logits)
+
+            # Store scores, attentions and hidden_states when required
+            if return_dict_in_generate:
+                if output_scores:
+                    scores += (next_tokens_scores,)
+                if output_attentions:
+                    decoder_attentions += (
+                        (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,)
+                    )
+                    if self.config.is_encoder_decoder:
+                        cross_attentions += (outputs.cross_attentions,)
+
+                if output_hidden_states:
+                    decoder_hidden_states += (
+                        (outputs.decoder_hidden_states,)
+                        if self.config.is_encoder_decoder
+                        else (outputs.hidden_states,)
+                    )
+
+            # argmax
+            next_tokens = torch.argmax(next_tokens_scores, dim=-1).to(device=input_ids.device)
+            # finished sentences should have their next token be a padding token
+            if eos_token_id is not None:
+                if pad_token_id is None:
+                    raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.")
+                next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)
+
+            # update generated ids, model inputs, and length for next step
+            input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)
+            if streamer is not None:
+                streamer.put(next_tokens.cpu())
+            model_kwargs = self._update_model_kwargs_for_generation(
+                outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder
+            )
+
+            # if eos_token was found in one sentence, set sentence to finished
+            if eos_token_id_tensor is not None:
+                unfinished_sequences = unfinished_sequences.mul(
+                    next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)
+                )
+
+                # stop when each sentence is finished
+                if unfinished_sequences.max() == 0:
+                    this_peer_finished = True
+
+            # stop if we exceed the maximum length
+            if stopping_criteria(input_ids, scores):
+                this_peer_finished = True
+
+            if this_peer_finished and not synced_gpus:
+                break
+
+        if streamer is not None:
+            streamer.end()
+
+        if return_dict_in_generate:
+            if self.config.is_encoder_decoder:
+                return GreedySearchEncoderDecoderOutput(
+                    sequences=input_ids,
+                    scores=scores,
+                    encoder_attentions=encoder_attentions,
+                    encoder_hidden_states=encoder_hidden_states,
+                    decoder_attentions=decoder_attentions,
+                    cross_attentions=cross_attentions,
+                    decoder_hidden_states=decoder_hidden_states,
+                    past_key_values=model_kwargs.get("past_key_values"),
+                )
+            else:
+                return GreedySearchDecoderOnlyOutput(
+                    sequences=input_ids,
+                    scores=scores,
+                    attentions=decoder_attentions,
+                    hidden_states=decoder_hidden_states,
+                    past_key_values=model_kwargs.get("past_key_values"),
+                )
+        else:
+            return input_ids
+
+
+class InternVLChatModel(PreTrainedModel):
+    config_class = InternVLChatConfig
+    main_input_name = 'pixel_values'
+    _no_split_modules = ['InternVisionModel', 'LlamaDecoderLayer']
+
+    def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None):
+        super().__init__(config)
+
+        image_size = config.force_image_size or config.vision_config.image_size
+        patch_size = config.vision_config.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
+        logger.info(f'num_image_token: {self.num_image_token}')
+        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:
+            # self.language_model = LlamaForCausalLM(config.llm_config)
+            self.language_model = MLlamaForCausalLM(config.llm_config)
+        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 * 4),
+            nn.Linear(vit_hidden_size * 4, llm_hidden_size),
+            nn.GELU(),
+            nn.Linear(llm_hidden_size, llm_hidden_size)
+        )
+
+        if config.force_image_size != config.vision_config.image_size:
+            self.vision_model.resize_pos_embeddings(
+                old_size=config.vision_config.image_size,
+                new_size=config.force_image_size,
+                patch_size=config.vision_config.patch_size
+            )
+
+        self.img_context_token_id = None
+
+        if config.use_backbone_lora:
+            self.wrap_backbone_lora(r=config.use_backbone_lora)
+
+        if config.use_llm_lora:
+            self.wrap_llm_lora(r=config.use_llm_lora)
+
+    def wrap_backbone_lora(self, r=128, lora_alpha=256, lora_dropout=0.05):
+        lora_config = LoraConfig(
+            r=r,
+            target_modules=['attn.qkv', 'attn.proj', 'mlp.fc1', 'mlp.fc2'],
+            lora_alpha=lora_alpha,
+            lora_dropout=lora_dropout,
+        )
+        self.vision_model = get_peft_model(self.vision_model, lora_config)
+        self.vision_model.print_trainable_parameters()
+
+    def wrap_llm_lora(self, r=128, lora_alpha=256, lora_dropout=0.05):
+        lora_config = LoraConfig(
+            r=r,
+            target_modules=['self_attn.q_proj', 'self_attn.k_proj', 'self_attn.v_proj', 'self_attn.o_proj',
+                            'mlp.gate_proj', 'mlp.down_proj', 'mlp.up_proj'],
+            lora_alpha=lora_alpha,
+            lora_dropout=lora_dropout,
+            task_type='CAUSAL_LM'
+        )
+        self.language_model = get_peft_model(self.language_model, lora_config)
+        self.language_model.print_trainable_parameters()
+
+    def forward(
+            self,
+            pixel_values: 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.config.use_return_dict
+
+        image_flags = image_flags.squeeze(-1)
+        input_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+        vit_embeds = self.extract_feature(pixel_values)
+        vit_embeds = vit_embeds[image_flags == 1]
+
+        B, N, C = input_embeds.shape
+        input_embeds = input_embeds.reshape(B * N, C)
+
+        input_ids = input_ids.reshape(B * N)
+        selected = (input_ids == self.img_context_token_id)
+        try:
+            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)
+        except:
+            pass
+
+        input_embeds = input_embeds.reshape(B, N, C)
+
+        outputs = self.language_model.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,
+        )
+        hidden_states = outputs[0]
+        logits = self.language_model.lm_head(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # Shift 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.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,
+        )
+
+    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)))
+        return x
+
+    def extract_feature(self, pixel_values):
+        if self.select_layer == -1:
+            vit_embeds = self.vision_model(
+                pixel_values=pixel_values,
+                output_hidden_states=False,
+                return_dict=True).last_hidden_state
+        else:
+            vit_embeds = self.vision_model(
+                pixel_values=pixel_values,
+                output_hidden_states=True,
+                return_dict=True).hidden_states[self.select_layer]
+        vit_embeds = vit_embeds[:, 1:, :]
+        # if torch.distributed.get_rank() == 0:
+        #     print("before pixel shuffle:", vit_embeds.shape)
+        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])
+        # if torch.distributed.get_rank() == 0:
+        #     print("after pixel shuffle:", vit_embeds.shape)
+        vit_embeds = self.mlp1(vit_embeds)
+        return vit_embeds
+
+    def chat(self, tokenizer, pixel_values, question, generation_config,
+             IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>'):
+
+        img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
+        self.img_context_token_id = img_context_token_id
+
+        from .conversation import get_conv_template
+
+        template = get_conv_template(self.template)
+        image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token + IMG_END_TOKEN
+        template.append_message(template.roles[0], image_tokens + '\n' + question)
+        template.append_message(template.roles[1], None)
+        query = template.get_prompt()
+        model_inputs = tokenizer(query, return_tensors='pt')
+        input_ids = model_inputs['input_ids'].cuda()
+        attention_mask = model_inputs['attention_mask'].cuda()
+
+        generation_output = self.generate(
+            pixel_values=pixel_values,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            **generation_config
+        )
+        response = tokenizer.batch_decode(generation_output, skip_special_tokens=True)[0]
+        query_to_print = query.replace(image_tokens, '<image>')
+        print(query_to_print, response)
+        return response
+
+    @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,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+            **generate_kwargs,
+    ) -> torch.LongTensor:
+
+        assert self.img_context_token_id is not None
+        if pixel_values is not None:
+            if visual_features is not None:
+                vit_embeds = visual_features
+            else:
+                vit_embeds = self.extract_feature(pixel_values)
+
+            input_embeds = self.language_model.get_input_embeddings()(input_ids)
+            B, N, C = input_embeds.shape
+            input_embeds = input_embeds.reshape(B * N, C)
+
+            input_ids = input_ids.reshape(B * N)
+            selected = (input_ids == self.img_context_token_id)
+            assert selected.sum() != 0
+            input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)
+
+            input_embeds = input_embeds.reshape(B, N, C)
+        else:
+            input_embeds = self.language_model.get_input_embeddings()(input_ids)
+
+        outputs = self.language_model.generate(
+            inputs_embeds=input_embeds,
+            attention_mask=attention_mask,
+            generation_config=generation_config,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            use_cache=True,
+            **generate_kwargs,
+        )
+
+        return outputs

special_tokens_map.json

@@ -0,0 +1,41 @@
+{
+  "additional_special_tokens": [
+    "<img>",
+    "</img>",
+    "<IMG_CONTEXT>",
+    "<quad>",
+    "</quad>",
+    "<ref>",
+    "</ref>",
+    "<box>",
+    "</box>"
+  ],
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:386c49cf943d71aa110361135338c50e38beeff0a66593480421f37b319e1a39
+size 1033105

tokenizer_config.json

@@ -0,0 +1,143 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "6": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "7": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "68": {
+      "content": "<img>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "70": {
+      "content": "</img>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "64000": {
+      "content": "<IMG_CONTEXT>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "64001": {
+      "content": "<quad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "64002": {
+      "content": "</quad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "64003": {
+      "content": "<ref>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "64004": {
+      "content": "</ref>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "64005": {
+      "content": "<box>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "64006": {
+      "content": "</box>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<img>",
+    "</img>",
+    "<IMG_CONTEXT>",
+    "<quad>",
+    "</quad>",
+    "<ref>",
+    "</ref>",
+    "<box>",
+    "</box>"
+  ],
+  "bos_token": "<|startoftext|>",
+  "chat_template": "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "legacy": true,
+  "model_max_length": 2048,
+  "pad_token": "<unk>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "tokenizer_class": "LlamaTokenizer",
+  "trust_remote_code": false,
+  "unk_token": "<unk>",
+  "use_default_system_prompt": false,
+  "use_fast": true
+}

trainer_state.json

@@ -0,0 +1,1145 @@
+{
+  "best_metric": null,
+  "best_model_checkpoint": null,
+  "epoch": 0.9986648865153538,
+  "eval_steps": 500,
+  "global_step": 187,
+  "is_hyper_param_search": false,
+  "is_local_process_zero": true,
+  "is_world_process_zero": true,
+  "log_history": [
+    {
+      "epoch": 0.01,
+      "learning_rate": 1.6666666666666667e-06,
+      "loss": 2.7864,
+      "step": 1
+    },
+    {
+      "epoch": 0.01,
+      "learning_rate": 3.3333333333333333e-06,
+      "loss": 3.1764,
+      "step": 2
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 5e-06,
+      "loss": 2.5692,
+      "step": 3
+    },
+    {
+      "epoch": 0.02,
+      "learning_rate": 6.666666666666667e-06,
+      "loss": 1.0477,
+      "step": 4
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 8.333333333333334e-06,
+      "loss": 0.9304,
+      "step": 5
+    },
+    {
+      "epoch": 0.03,
+      "learning_rate": 1e-05,
+      "loss": 0.8711,
+      "step": 6
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 9.999246866958693e-06,
+      "loss": 0.6422,
+      "step": 7
+    },
+    {
+      "epoch": 0.04,
+      "learning_rate": 9.99698769471852e-06,
+      "loss": 0.7042,
+      "step": 8
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 9.993223163862385e-06,
+      "loss": 0.7917,
+      "step": 9
+    },
+    {
+      "epoch": 0.05,
+      "learning_rate": 9.98795440846732e-06,
+      "loss": 0.6634,
+      "step": 10
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.981183015762831e-06,
+      "loss": 0.7043,
+      "step": 11
+    },
+    {
+      "epoch": 0.06,
+      "learning_rate": 9.972911025652754e-06,
+      "loss": 0.6872,
+      "step": 12
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.963140930100713e-06,
+      "loss": 0.7418,
+      "step": 13
+    },
+    {
+      "epoch": 0.07,
+      "learning_rate": 9.951875672379424e-06,
+      "loss": 0.7245,
+      "step": 14
+    },
+    {
+      "epoch": 0.08,
+      "learning_rate": 9.939118646184007e-06,
+      "loss": 0.7227,
+      "step": 15
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.924873694609636e-06,
+      "loss": 0.6843,
+      "step": 16
+    },
+    {
+      "epoch": 0.09,
+      "learning_rate": 9.909145108993794e-06,
+      "loss": 0.65,
+      "step": 17
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.891937627623486e-06,
+      "loss": 0.7358,
+      "step": 18
+    },
+    {
+      "epoch": 0.1,
+      "learning_rate": 9.873256434307828e-06,
+      "loss": 0.763,
+      "step": 19
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.853107156816393e-06,
+      "loss": 0.7882,
+      "step": 20
+    },
+    {
+      "epoch": 0.11,
+      "learning_rate": 9.831495865183832e-06,
+      "loss": 0.6681,
+      "step": 21
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.808429069881267e-06,
+      "loss": 0.6109,
+      "step": 22
+    },
+    {
+      "epoch": 0.12,
+      "learning_rate": 9.783913719854977e-06,
+      "loss": 0.7464,
+      "step": 23
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.757957200433011e-06,
+      "loss": 0.7748,
+      "step": 24
+    },
+    {
+      "epoch": 0.13,
+      "learning_rate": 9.730567331100333e-06,
+      "loss": 0.6528,
+      "step": 25
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.701752363143183e-06,
+      "loss": 0.7256,
+      "step": 26
+    },
+    {
+      "epoch": 0.14,
+      "learning_rate": 9.67152097716334e-06,
+      "loss": 0.6889,
+      "step": 27
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.639882280463071e-06,
+      "loss": 0.728,
+      "step": 28
+    },
+    {
+      "epoch": 0.15,
+      "learning_rate": 9.606845804301523e-06,
+      "loss": 0.7518,
+      "step": 29
+    },
+    {
+      "epoch": 0.16,
+      "learning_rate": 9.572421501023403e-06,
+      "loss": 0.7482,
+      "step": 30
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.536619741060799e-06,
+      "loss": 0.6835,
+      "step": 31
+    },
+    {
+      "epoch": 0.17,
+      "learning_rate": 9.499451309809058e-06,
+      "loss": 0.6588,
+      "step": 32
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.460927404377647e-06,
+      "loss": 0.8194,
+      "step": 33
+    },
+    {
+      "epoch": 0.18,
+      "learning_rate": 9.421059630216992e-06,
+      "loss": 0.6716,
+      "step": 34
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.37985999762229e-06,
+      "loss": 0.7035,
+      "step": 35
+    },
+    {
+      "epoch": 0.19,
+      "learning_rate": 9.337340918115385e-06,
+      "loss": 0.7737,
+      "step": 36
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.29351520070574e-06,
+      "loss": 0.7908,
+      "step": 37
+    },
+    {
+      "epoch": 0.2,
+      "learning_rate": 9.24839604803169e-06,
+      "loss": 0.7395,
+      "step": 38
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.201997052383107e-06,
+      "loss": 0.6748,
+      "step": 39
+    },
+    {
+      "epoch": 0.21,
+      "learning_rate": 9.154332191606671e-06,
+      "loss": 0.6289,
+      "step": 40
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.105415824895008e-06,
+      "loss": 0.7595,
+      "step": 41
+    },
+    {
+      "epoch": 0.22,
+      "learning_rate": 9.055262688460931e-06,
+      "loss": 0.7401,
+      "step": 42
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 9.003887891098108e-06,
+      "loss": 0.6546,
+      "step": 43
+    },
+    {
+      "epoch": 0.23,
+      "learning_rate": 8.951306909629492e-06,
+      "loss": 0.6714,
+      "step": 44
+    },
+    {
+      "epoch": 0.24,
+      "learning_rate": 8.89753558424488e-06,
+      "loss": 0.666,
+      "step": 45
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 8.842590113729001e-06,
+      "loss": 0.6523,
+      "step": 46
+    },
+    {
+      "epoch": 0.25,
+      "learning_rate": 8.786487050581583e-06,
+      "loss": 0.6988,
+      "step": 47
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 8.729243296030851e-06,
+      "loss": 0.7612,
+      "step": 48
+    },
+    {
+      "epoch": 0.26,
+      "learning_rate": 8.670876094941991e-06,
+      "loss": 0.7252,
+      "step": 49
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 8.611403030622074e-06,
+      "loss": 0.6683,
+      "step": 50
+    },
+    {
+      "epoch": 0.27,
+      "learning_rate": 8.55084201952302e-06,
+      "loss": 0.6546,
+      "step": 51
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 8.489211305844216e-06,
+      "loss": 0.6822,
+      "step": 52
+    },
+    {
+      "epoch": 0.28,
+      "learning_rate": 8.4265294560364e-06,
+      "loss": 0.7031,
+      "step": 53
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 8.362815353208441e-06,
+      "loss": 0.706,
+      "step": 54
+    },
+    {
+      "epoch": 0.29,
+      "learning_rate": 8.298088191438753e-06,
+      "loss": 0.7278,
+      "step": 55
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 8.23236746999302e-06,
+      "loss": 0.6482,
+      "step": 56
+    },
+    {
+      "epoch": 0.3,
+      "learning_rate": 8.165672987449962e-06,
+      "loss": 0.6553,
+      "step": 57
+    },
+    {
+      "epoch": 0.31,
+      "learning_rate": 8.098024835736977e-06,
+      "loss": 0.7261,
+      "step": 58
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 8.029443394077356e-06,
+      "loss": 0.5941,
+      "step": 59
+    },
+    {
+      "epoch": 0.32,
+      "learning_rate": 7.959949322850994e-06,
+      "loss": 0.7581,
+      "step": 60
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 7.889563557370378e-06,
+      "loss": 0.6762,
+      "step": 61
+    },
+    {
+      "epoch": 0.33,
+      "learning_rate": 7.818307301573757e-06,
+      "loss": 0.718,
+      "step": 62
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 7.746202021637385e-06,
+      "loss": 0.6745,
+      "step": 63
+    },
+    {
+      "epoch": 0.34,
+      "learning_rate": 7.67326943950877e-06,
+      "loss": 0.6839,
+      "step": 64
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 7.599531526362873e-06,
+      "loss": 0.7154,
+      "step": 65
+    },
+    {
+      "epoch": 0.35,
+      "learning_rate": 7.525010495983202e-06,
+      "loss": 0.7622,
+      "step": 66
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 7.449728798069864e-06,
+      "loss": 0.6516,
+      "step": 67
+    },
+    {
+      "epoch": 0.36,
+      "learning_rate": 7.373709111476498e-06,
+      "loss": 0.635,
+      "step": 68
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 7.296974337378209e-06,
+      "loss": 0.6369,
+      "step": 69
+    },
+    {
+      "epoch": 0.37,
+      "learning_rate": 7.219547592372512e-06,
+      "loss": 0.6783,
+      "step": 70
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 7.141452201515386e-06,
+      "loss": 0.7469,
+      "step": 71
+    },
+    {
+      "epoch": 0.38,
+      "learning_rate": 7.062711691294525e-06,
+      "loss": 0.6119,
+      "step": 72
+    },
+    {
+      "epoch": 0.39,
+      "learning_rate": 6.983349782541901e-06,
+      "loss": 0.6803,
+      "step": 73
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 6.903390383287795e-06,
+      "loss": 0.6944,
+      "step": 74
+    },
+    {
+      "epoch": 0.4,
+      "learning_rate": 6.822857581558423e-06,
+      "loss": 0.662,
+      "step": 75
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 6.741775638119345e-06,
+      "loss": 0.7491,
+      "step": 76
+    },
+    {
+      "epoch": 0.41,
+      "learning_rate": 6.66016897916682e-06,
+      "loss": 0.7162,
+      "step": 77
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 6.57806218896935e-06,
+      "loss": 0.7451,
+      "step": 78
+    },
+    {
+      "epoch": 0.42,
+      "learning_rate": 6.495480002461577e-06,
+      "loss": 0.5831,
+      "step": 79
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 6.412447297792818e-06,
+      "loss": 0.7656,
+      "step": 80
+    },
+    {
+      "epoch": 0.43,
+      "learning_rate": 6.328989088832431e-06,
+      "loss": 0.6309,
+      "step": 81
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 6.245130517634307e-06,
+      "loss": 0.719,
+      "step": 82
+    },
+    {
+      "epoch": 0.44,
+      "learning_rate": 6.160896846862754e-06,
+      "loss": 0.6109,
+      "step": 83
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 6.076313452182033e-06,
+      "loss": 0.6539,
+      "step": 84
+    },
+    {
+      "epoch": 0.45,
+      "learning_rate": 5.991405814611855e-06,
+      "loss": 0.6642,
+      "step": 85
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 5.9061995128511455e-06,
+      "loss": 0.6832,
+      "step": 86
+    },
+    {
+      "epoch": 0.46,
+      "learning_rate": 5.820720215572375e-06,
+      "loss": 0.7376,
+      "step": 87
+    },
+    {
+      "epoch": 0.47,
+      "learning_rate": 5.734993673688801e-06,
+      "loss": 0.6968,
+      "step": 88
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 5.6490457125969035e-06,
+      "loss": 0.743,
+      "step": 89
+    },
+    {
+      "epoch": 0.48,
+      "learning_rate": 5.562902224396416e-06,
+      "loss": 0.6412,
+      "step": 90
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 5.476589160090238e-06,
+      "loss": 0.6313,
+      "step": 91
+    },
+    {
+      "epoch": 0.49,
+      "learning_rate": 5.390132521766626e-06,
+      "loss": 0.7327,
+      "step": 92
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 5.30355835476596e-06,
+      "loss": 0.5725,
+      "step": 93
+    },
+    {
+      "epoch": 0.5,
+      "learning_rate": 5.216892739834519e-06,
+      "loss": 0.6587,
+      "step": 94
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 5.13016178526756e-06,
+      "loss": 0.5893,
+      "step": 95
+    },
+    {
+      "epoch": 0.51,
+      "learning_rate": 5.043391619044122e-06,
+      "loss": 0.6343,
+      "step": 96
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 4.956608380955877e-06,
+      "loss": 0.6764,
+      "step": 97
+    },
+    {
+      "epoch": 0.52,
+      "learning_rate": 4.869838214732441e-06,
+      "loss": 0.6486,
+      "step": 98
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 4.783107260165483e-06,
+      "loss": 0.5761,
+      "step": 99
+    },
+    {
+      "epoch": 0.53,
+      "learning_rate": 4.696441645234042e-06,
+      "loss": 0.6928,
+      "step": 100
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 4.609867478233377e-06,
+      "loss": 0.635,
+      "step": 101
+    },
+    {
+      "epoch": 0.54,
+      "learning_rate": 4.523410839909764e-06,
+      "loss": 0.6497,
+      "step": 102
+    },
+    {
+      "epoch": 0.55,
+      "learning_rate": 4.437097775603587e-06,
+      "loss": 0.6256,
+      "step": 103
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 4.350954287403099e-06,
+      "loss": 0.6605,
+      "step": 104
+    },
+    {
+      "epoch": 0.56,
+      "learning_rate": 4.265006326311199e-06,
+      "loss": 0.6458,
+      "step": 105
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 4.179279784427625e-06,
+      "loss": 0.6496,
+      "step": 106
+    },
+    {
+      "epoch": 0.57,
+      "learning_rate": 4.093800487148857e-06,
+      "loss": 0.6505,
+      "step": 107
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 4.008594185388146e-06,
+      "loss": 0.6259,
+      "step": 108
+    },
+    {
+      "epoch": 0.58,
+      "learning_rate": 3.9236865478179685e-06,
+      "loss": 0.6401,
+      "step": 109
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 3.839103153137247e-06,
+      "loss": 0.5964,
+      "step": 110
+    },
+    {
+      "epoch": 0.59,
+      "learning_rate": 3.7548694823656945e-06,
+      "loss": 0.6578,
+      "step": 111
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 3.671010911167572e-06,
+      "loss": 0.6361,
+      "step": 112
+    },
+    {
+      "epoch": 0.6,
+      "learning_rate": 3.5875527022071808e-06,
+      "loss": 0.6685,
+      "step": 113
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 3.5045199975384225e-06,
+      "loss": 0.5762,
+      "step": 114
+    },
+    {
+      "epoch": 0.61,
+      "learning_rate": 3.4219378110306523e-06,
+      "loss": 0.5239,
+      "step": 115
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 3.3398310208331806e-06,
+      "loss": 0.6064,
+      "step": 116
+    },
+    {
+      "epoch": 0.62,
+      "learning_rate": 3.2582243618806574e-06,
+      "loss": 0.6162,
+      "step": 117
+    },
+    {
+      "epoch": 0.63,
+      "learning_rate": 3.177142418441578e-06,
+      "loss": 0.5415,
+      "step": 118
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 3.096609616712207e-06,
+      "loss": 0.5404,
+      "step": 119
+    },
+    {
+      "epoch": 0.64,
+      "learning_rate": 3.0166502174581012e-06,
+      "loss": 0.6535,
+      "step": 120
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 2.937288308705475e-06,
+      "loss": 0.6472,
+      "step": 121
+    },
+    {
+      "epoch": 0.65,
+      "learning_rate": 2.858547798484613e-06,
+      "loss": 0.5469,
+      "step": 122
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 2.7804524076274898e-06,
+      "loss": 0.699,
+      "step": 123
+    },
+    {
+      "epoch": 0.66,
+      "learning_rate": 2.7030256626217932e-06,
+      "loss": 0.6188,
+      "step": 124
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 2.6262908885235046e-06,
+      "loss": 0.6024,
+      "step": 125
+    },
+    {
+      "epoch": 0.67,
+      "learning_rate": 2.550271201930136e-06,
+      "loss": 0.5934,
+      "step": 126
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 2.474989504016798e-06,
+      "loss": 0.6309,
+      "step": 127
+    },
+    {
+      "epoch": 0.68,
+      "learning_rate": 2.4004684736371276e-06,
+      "loss": 0.6157,
+      "step": 128
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 2.32673056049123e-06,
+      "loss": 0.6419,
+      "step": 129
+    },
+    {
+      "epoch": 0.69,
+      "learning_rate": 2.253797978362617e-06,
+      "loss": 0.5713,
+      "step": 130
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 2.1816926984262454e-06,
+      "loss": 0.6575,
+      "step": 131
+    },
+    {
+      "epoch": 0.7,
+      "learning_rate": 2.1104364426296237e-06,
+      "loss": 0.6824,
+      "step": 132
+    },
+    {
+      "epoch": 0.71,
+      "learning_rate": 2.040050677149008e-06,
+      "loss": 0.6374,
+      "step": 133
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 1.970556605922645e-06,
+      "loss": 0.5793,
+      "step": 134
+    },
+    {
+      "epoch": 0.72,
+      "learning_rate": 1.9019751642630252e-06,
+      "loss": 0.5671,
+      "step": 135
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 1.8343270125500379e-06,
+      "loss": 0.6081,
+      "step": 136
+    },
+    {
+      "epoch": 0.73,
+      "learning_rate": 1.7676325300069824e-06,
+      "loss": 0.5478,
+      "step": 137
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 1.7019118085612474e-06,
+      "loss": 0.6179,
+      "step": 138
+    },
+    {
+      "epoch": 0.74,
+      "learning_rate": 1.6371846467915603e-06,
+      "loss": 0.6195,
+      "step": 139
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 1.5734705439636017e-06,
+      "loss": 0.6112,
+      "step": 140
+    },
+    {
+      "epoch": 0.75,
+      "learning_rate": 1.5107886941557853e-06,
+      "loss": 0.5949,
+      "step": 141
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 1.4491579804769817e-06,
+      "loss": 0.6089,
+      "step": 142
+    },
+    {
+      "epoch": 0.76,
+      "learning_rate": 1.3885969693779277e-06,
+      "loss": 0.6091,
+      "step": 143
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 1.3291239050580085e-06,
+      "loss": 0.6016,
+      "step": 144
+    },
+    {
+      "epoch": 0.77,
+      "learning_rate": 1.2707567039691505e-06,
+      "loss": 0.6553,
+      "step": 145
+    },
+    {
+      "epoch": 0.78,
+      "learning_rate": 1.213512949418419e-06,
+      "loss": 0.633,
+      "step": 146
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 1.1574098862709993e-06,
+      "loss": 0.6151,
+      "step": 147
+    },
+    {
+      "epoch": 0.79,
+      "learning_rate": 1.1024644157551206e-06,
+      "loss": 0.6017,
+      "step": 148
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 1.0486930903705095e-06,
+      "loss": 0.5714,
+      "step": 149
+    },
+    {
+      "epoch": 0.8,
+      "learning_rate": 9.961121089018933e-07,
+      "loss": 0.6089,
+      "step": 150
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 9.447373115390702e-07,
+      "loss": 0.5789,
+      "step": 151
+    },
+    {
+      "epoch": 0.81,
+      "learning_rate": 8.945841751049916e-07,
+      "loss": 0.6077,
+      "step": 152
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 8.45667808393329e-07,
+      "loss": 0.5956,
+      "step": 153
+    },
+    {
+      "epoch": 0.82,
+      "learning_rate": 7.980029476168943e-07,
+      "loss": 0.4987,
+      "step": 154
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 7.516039519683105e-07,
+      "loss": 0.5695,
+      "step": 155
+    },
+    {
+      "epoch": 0.83,
+      "learning_rate": 7.064847992942614e-07,
+      "loss": 0.6005,
+      "step": 156
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 6.626590818846163e-07,
+      "loss": 0.569,
+      "step": 157
+    },
+    {
+      "epoch": 0.84,
+      "learning_rate": 6.201400023777105e-07,
+      "loss": 0.5699,
+      "step": 158
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 5.789403697830104e-07,
+      "loss": 0.7248,
+      "step": 159
+    },
+    {
+      "epoch": 0.85,
+      "learning_rate": 5.390725956223531e-07,
+      "loss": 0.7469,
+      "step": 160
+    },
+    {
+      "epoch": 0.86,
+      "learning_rate": 5.005486901909429e-07,
+      "loss": 0.7068,
+      "step": 161
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 4.6338025893920167e-07,
+      "loss": 0.5459,
+      "step": 162
+    },
+    {
+      "epoch": 0.87,
+      "learning_rate": 4.275784989765985e-07,
+      "loss": 0.6079,
+      "step": 163
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 3.93154195698478e-07,
+      "loss": 0.5906,
+      "step": 164
+    },
+    {
+      "epoch": 0.88,
+      "learning_rate": 3.6011771953693044e-07,
+      "loss": 0.6135,
+      "step": 165
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 3.284790228366602e-07,
+      "loss": 0.4893,
+      "step": 166
+    },
+    {
+      "epoch": 0.89,
+      "learning_rate": 2.982476368568177e-07,
+      "loss": 0.5953,
+      "step": 167
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 2.6943266889966624e-07,
+      "loss": 0.6246,
+      "step": 168
+    },
+    {
+      "epoch": 0.9,
+      "learning_rate": 2.4204279956698994e-07,
+      "loss": 0.5807,
+      "step": 169
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 2.1608628014502364e-07,
+      "loss": 0.6546,
+      "step": 170
+    },
+    {
+      "epoch": 0.91,
+      "learning_rate": 1.915709301187335e-07,
+      "loss": 0.5779,
+      "step": 171
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 1.6850413481616868e-07,
+      "loss": 0.6579,
+      "step": 172
+    },
+    {
+      "epoch": 0.92,
+      "learning_rate": 1.468928431836092e-07,
+      "loss": 0.6288,
+      "step": 173
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 1.2674356569217282e-07,
+      "loss": 0.6266,
+      "step": 174
+    },
+    {
+      "epoch": 0.93,
+      "learning_rate": 1.080623723765134e-07,
+      "loss": 0.6483,
+      "step": 175
+    },
+    {
+      "epoch": 0.94,
+      "learning_rate": 9.085489100620737e-08,
+      "loss": 0.5573,
+      "step": 176
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 7.512630539036502e-08,
+      "loss": 0.5503,
+      "step": 177
+    },
+    {
+      "epoch": 0.95,
+      "learning_rate": 6.088135381599414e-08,
+      "loss": 0.6147,
+      "step": 178
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 4.8124327620576726e-08,
+      "loss": 0.6135,
+      "step": 179
+    },
+    {
+      "epoch": 0.96,
+      "learning_rate": 3.685906989928656e-08,
+      "loss": 0.6094,
+      "step": 180
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 2.7088974347246888e-08,
+      "loss": 0.6843,
+      "step": 181
+    },
+    {
+      "epoch": 0.97,
+      "learning_rate": 1.8816984237169378e-08,
+      "loss": 0.5167,
+      "step": 182
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 1.2045591532681145e-08,
+      "loss": 0.5398,
+      "step": 183
+    },
+    {
+      "epoch": 0.98,
+      "learning_rate": 6.7768361376152616e-09,
+      "loss": 0.5482,
+      "step": 184
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 3.0123052814812203e-09,
+      "loss": 0.6061,
+      "step": 185
+    },
+    {
+      "epoch": 0.99,
+      "learning_rate": 7.53133041307974e-10,
+      "loss": 0.5735,
+      "step": 186
+    },
+    {
+      "epoch": 1.0,
+      "learning_rate": 0.0,
+      "loss": 0.5353,
+      "step": 187
+    }
+  ],
+  "log_save_evaluate_time": 3923.0782437324524,
+  "logging_steps": 1.0,
+  "max_steps": 187,
+  "num_input_tokens_seen": 0,
+  "num_train_epochs": 1,
+  "save_steps": 40,
+  "total_flos": 7.673772156280793e+19,
+  "total_tokens": 900759552.0,
+  "train_batch_size": 2,
+  "trial_name": null,
+  "trial_params": null
+}

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f1988675517ebb16f267246bf5443edac3ac5ae51e2adab418627265a4ebccfe
+size 6331

zero_to_fp32.py

@@ -0,0 +1,578 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 1, 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example: python zero_to_fp32.py . pytorch_model.bin
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+from collections import OrderedDict
+from dataclasses import dataclass
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+                                            FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
+
+
+@dataclass
+class zero_model_state:
+    buffers: dict()
+    param_shapes: dict()
+    shared_params: list
+    ds_version: int
+    frozen_param_shapes: dict()
+    frozen_param_fragments: dict()
+
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+    if not os.path.isdir(checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+    # there should be only one file
+    if zero_stage <= 2:
+        file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+    elif zero_stage == 3:
+        file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+    if not os.path.exists(file):
+        raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+    return file
+
+
+def get_checkpoint_files(checkpoint_dir, glob_pattern):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+    if len(ckpt_files) == 0:
+        raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+    return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def parse_model_states(files):
+    zero_model_states = []
+    for file in files:
+        state_dict = torch.load(file, map_location=device)
+
+        if BUFFER_NAMES not in state_dict:
+            raise ValueError(f"{file} is not a model state checkpoint")
+        buffer_names = state_dict[BUFFER_NAMES]
+        if debug:
+            print("Found buffers:", buffer_names)
+
+        # recover just the buffers while restoring them to fp32 if they were saved in fp16
+        buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
+        param_shapes = state_dict[PARAM_SHAPES]
+
+        # collect parameters that are included in param_shapes
+        param_names = []
+        for s in param_shapes:
+            for name in s.keys():
+                param_names.append(name)
+
+        # update with frozen parameters
+        frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
+        if frozen_param_shapes is not None:
+            if debug:
+                print(f"Found frozen_param_shapes: {frozen_param_shapes}")
+            param_names += list(frozen_param_shapes.keys())
+
+        # handle shared params
+        shared_params = [[k, v] for k, v in state_dict["shared_params"].items()]
+
+        ds_version = state_dict.get(DS_VERSION, None)
+
+        frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
+
+        z_model_state = zero_model_state(buffers=buffers,
+                                         param_shapes=param_shapes,
+                                         shared_params=shared_params,
+                                         ds_version=ds_version,
+                                         frozen_param_shapes=frozen_param_shapes,
+                                         frozen_param_fragments=frozen_param_fragments)
+        zero_model_states.append(z_model_state)
+
+    return zero_model_states
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+
+    total_files = len(files)
+    state_dicts = []
+    for f in files:
+        state_dicts.append(torch.load(f, map_location=device))
+
+    if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]:
+        raise ValueError(f"{files[0]} is not a zero checkpoint")
+    zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+    world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+    # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+    # parameters can be different from data parallelism for non-expert parameters. So we can just
+    # use the max of the partition_count to get the dp world_size.
+
+    if type(world_size) is list:
+        world_size = max(world_size)
+
+    if world_size != total_files:
+        raise ValueError(
+            f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+            "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+        )
+
+    # the groups are named differently in each stage
+    if zero_stage <= 2:
+        fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+    elif zero_stage == 3:
+        fp32_groups_key = FP32_FLAT_GROUPS
+    else:
+        raise ValueError(f"unknown zero stage {zero_stage}")
+
+    if zero_stage <= 2:
+        fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
+    elif zero_stage == 3:
+        # if there is more than one param group, there will be multiple flattened tensors - one
+        # flattened tensor per group - for simplicity merge them into a single tensor
+        #
+        # XXX: could make the script more memory efficient for when there are multiple groups - it
+        # will require matching the sub-lists of param_shapes for each param group flattened tensor
+
+        fp32_flat_groups = [
+            torch.cat(state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key], 0) for i in range(len(state_dicts))
+        ]
+
+    return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir):
+    """
+    Returns fp32 state_dict reconstructed from ds checkpoint
+
+    Args:
+        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+    """
+    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+    optim_files = get_optim_files(ds_checkpoint_dir)
+    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+    print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+    model_files = get_model_state_files(ds_checkpoint_dir)
+
+    zero_model_states = parse_model_states(model_files)
+    print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+    if zero_stage <= 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+    frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+    if debug:
+        num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+        print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        state_dict[name] = frozen_param_fragments[name]
+
+        if debug:
+            print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+
+    # Reconstruction protocol:
+    #
+    # XXX: document this
+
+    if debug:
+        for i in range(world_size):
+            for j in range(len(fp32_flat_groups[0])):
+                print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+    # XXX: memory usage doubles here (zero2)
+    num_param_groups = len(fp32_flat_groups[0])
+    merged_single_partition_of_fp32_groups = []
+    for i in range(num_param_groups):
+        merged_partitions = [sd[i] for sd in fp32_flat_groups]
+        full_single_fp32_vector = torch.cat(merged_partitions, 0)
+        merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+    avail_numel = sum(
+        [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
+
+    if debug:
+        wanted_params = sum([len(shapes) for shapes in param_shapes])
+        wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    total_numel = 0
+    total_params = 0
+    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+        offset = 0
+        avail_numel = full_single_fp32_vector.numel()
+        for name, shape in shapes.items():
+
+            unpartitioned_numel = shape.numel()
+            total_numel += unpartitioned_numel
+            total_params += 1
+
+            if debug:
+                print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+            state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
+            offset += unpartitioned_numel
+
+        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+        # live optimizer object, so we are checking that the numbers are within the right range
+        align_to = 2 * world_size
+
+        def zero2_align(x):
+            return align_to * math.ceil(x / align_to)
+
+        if debug:
+            print(f"original offset={offset}, avail_numel={avail_numel}")
+
+        offset = zero2_align(offset)
+        avail_numel = zero2_align(avail_numel)
+
+        if debug:
+            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
+
+        # Sanity check
+        if offset != avail_numel:
+            raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+    _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+    remainder = unpartitioned_numel % world_size
+    padding_numel = (world_size - remainder) if remainder else 0
+    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+    return partitioned_numel, padding_numel
+
+
+def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    if debug:
+        for i in range(world_size):
+            num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+            print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in zero_model_states[0].frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+        state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+    avail_numel = fp32_flat_groups[0].numel() * world_size
+    # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+    # param, re-consolidating each param, while dealing with padding if any
+
+    # merge list of dicts, preserving order
+    param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+    if debug:
+        for i in range(world_size):
+            print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+        wanted_params = len(param_shapes)
+        wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+        # not asserting if there is a mismatch due to possible padding
+        avail_numel = fp32_flat_groups[0].numel() * world_size
+        print(f"Trainable params: Have {avail_numel} numels to process.")
+        print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    offset = 0
+    total_numel = 0
+    total_params = 0
+    for name, shape in param_shapes.items():
+
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+        total_params += 1
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+        # XXX: memory usage doubles here
+        state_dict[name] = torch.cat(
+            tuple(fp32_flat_groups[i].narrow(0, offset, partitioned_numel) for i in range(world_size)),
+            0).narrow(0, 0, unpartitioned_numel).view(shape)
+        offset += partitioned_numel
+
+    offset *= world_size
+
+    # Sanity check
+    if offset != avail_numel:
+        raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+    _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        if pair[1] in state_dict:
+            state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+    via a model hub.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+
+    Returns:
+        - pytorch ``state_dict``
+
+    Note: this approach may not work if your application doesn't have sufficient free CPU memory and
+    you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        # do the training and checkpoint saving
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+        model = model.cpu() # move to cpu
+        model.load_state_dict(state_dict)
+        # submit to model hub or save the model to share with others
+
+    In this example the ``model`` will no longer be usable in the deepspeed context of the same
+    application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+    """
+    if tag is None:
+        latest_path = os.path.join(checkpoint_dir, 'latest')
+        if os.path.isfile(latest_path):
+            with open(latest_path, 'r') as fd:
+                tag = fd.read().strip()
+        else:
+            raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+    ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+    if not os.path.isdir(ds_checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+    return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_file, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+    loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``output_file``: path to the pytorch fp32 state_dict output file (e.g. path/pytorch_model.bin)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+    """
+
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+    print(f"Saving fp32 state dict to {output_file}")
+    torch.save(state_dict, output_file)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+    """
+    1. Put the provided model to cpu
+    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+    3. Load it into the provided model
+
+    Args:
+        - ``model``: the model object to update
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+    Returns:
+        - ``model`: modified model
+
+    Make sure you have plenty of CPU memory available before you call this function. If you don't
+    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+    conveniently placed for you in the checkpoint folder.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+        # submit to model hub or save the model to share with others
+
+    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    """
+    logger.info(f"Extracting fp32 weights")
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+    logger.info(f"Overwriting model with fp32 weights")
+    model = model.cpu()
+    model.load_state_dict(state_dict, strict=False)
+
+    return model
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint_dir",
+                        type=str,
+                        help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+    parser.add_argument(
+        "output_file",
+        type=str,
+        help="path to the pytorch fp32 state_dict output file (e.g. path/checkpoint-12/pytorch_model.bin)")
+    parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+    args = parser.parse_args()
+
+    debug = args.debug
+
+    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, args.output_file)