import json import os from typing import Any, Dict, List, Tuple, Union import numpy as np import torch from einops import rearrange from transformers import AutoTokenizer from maga_transformer.config.gpt_init_model_parameters import GptInitModelParameters, TemplateType from maga_transformer.distribute.worker_info import ParallelInfo, g_parallel_info from maga_transformer.model_factory_register import register_model from maga_transformer.models.eva2clip_vit import EVA2CLIPImageEmbedding from maga_transformer.models.cogvlm2_weight import CogVLM2WeightInfo, CogVLM2VitWeights from maga_transformer.models.llama import Llama from maga_transformer.models.multimodal.multimodal_mixin import MultiModalMixin from maga_transformer.utils.util import to_torch_dtype LANGUAGE_TOKEN_TYPE = 0 VISION_TOKEN_TYPE = 1 class CogVLM2(Llama, MultiModalMixin): def __init__(self, config: GptInitModelParameters): quant_algo = config.quant_algo if quant_algo.isGptq() or quant_algo.isAwq() or quant_algo.isSmoothQuant() or quant_algo.isOmniQuant(): raise Exception("CogVLM2 only support FP32, BF16, FP16, INT8, not support other quant algorithm") super().__init__(config) def _init_multimodal(self, config: GptInitModelParameters): self.mm_part = EVA2CLIPImageEmbedding(config) config.mm_related_params.vit_weights = CogVLM2VitWeights( {"vit": self.mm_part.vit} ) def load(self, device: str): if os.environ.get("VIT_TRT", "0") == "1": weights_info = self.get_weight_cls()(self.config, g_parallel_info.tp_size, g_parallel_info.tp_rank) self.init_mm_trt( weights_info, self.config.ckpt_path, self.config.mm_related_params, device, to_torch_dtype(self.config.data_type) ) super().load(device=device) @staticmethod def _create_config(ckpt_path): config = GptInitModelParameters( head_num=0, size_per_head=0, layer_num=0, max_seq_len=0, vocab_size=0, ckpt_path=ckpt_path, activation_type="SiGLU", norm_type="rmsnorm", rotary_embedding_dim=128, rotary_embedding_style=1, rotary_embedding_base=500000, has_post_decoder_layernorm=True ) # hugggingface config_path = os.path.join(ckpt_path, "config.json") if os.path.exists(config_path): with open(config_path) as reader: content = reader.read() config_json = json.loads(content) CogVLM2.from_huggingface(config, config_json) else: raise Exception("cogvlm2 parameter from unkown source") config.tp_split_emb_and_lm_head = False # cogvlm2 embedding can't tp return config @staticmethod def get_weight_cls(): return CogVLM2WeightInfo @staticmethod def from_huggingface(config: GptInitModelParameters, config_json: Dict[str, Any]): config.use_expert_attention = True config.add_special_tokens = False config.head_num = config_json["num_attention_heads"] config.head_num_kv = 8 config.hidden_size = config_json["hidden_size"] config.size_per_head = ( config_json["hidden_size"] // config_json["num_attention_heads"] ) config.layer_num = config_json["num_hidden_layers"] config.max_seq_len = config_json.get("max_position_embeddings", 8192) config.vocab_size = config_json["vocab_size"] config.layernorm_eps = config_json.get("rms_norm_eps", 1e-05) config.inter_size = config_json["intermediate_size"] config.rotary_embedding_dim = config.size_per_head config.tie_word_embeddings = config_json.get("tie_word_embeddings", False) config.build_position_ids = True try: template_type_str = config_json.get("template_version", "chat") config.template_type = TemplateType[template_type_str] except KeyError: raise Exception(f"unknown template_type: {template_type_str}") config.reserve_runtime_mem_mb = 2048 vit_config = config_json["vision_config"] config.mm_related_params.config.update(vit_config) # use vision hidden size for linear_proj and conv layer in eva2clip config.mm_related_params.config['use_vision_hidden_size'] = True config.special_tokens.bos_token_id = config_json["bos_token_id"] config.special_tokens.pad_token_id = config_json["pad_token_id"] if isinstance(config_json['eos_token_id'], list): config.special_tokens.eos_token_id = config_json['eos_token_id'][0] config.special_tokens.stop_words_id_list = [[x] for x in config_json['eos_token_id']] else: config.special_tokens.eos_token_id = config_json['eos_token_id'] @classmethod def get_tokenizer(cls, config: GptInitModelParameters): tokenizer = AutoTokenizer.from_pretrained( config.tokenizer_path, trust_remote_code=True ) tokenizer.pad_token_id = 128002 # llama3 adapt for cogvlm return tokenizer def extend_generate_combo_token_types(self, combo_tokens: List[int]) -> List[int]: return [LANGUAGE_TOKEN_TYPE] * len(combo_tokens) def extend_context_combo_token_types(self, token_types: List[int]) -> List[int]: return token_types def extend_generate_position_ids( self, generate_batch_size: int, num_beams: int, vision_token_length: List[int], seq_lengths_list: List[int] ) -> List[int]: position_ids = [] for i in range(generate_batch_size): start_idx = i * num_beams if vision_token_length[start_idx] == 0: # only text position_ids.append(seq_lengths_list[start_idx] - 1) else: # ccontain image position_ids.append(seq_lengths_list[start_idx] - vision_token_length[start_idx] + 2) return position_ids def extend_context_position_ids( self, context_begin_position: int, context_end_position: int, token_type_ids: torch.Tensor, token_ids: torch.Tensor ) -> List[int]: # construct position ids for rotary embedding, assuming the token_type_ids is [T, V, V, V, V, V, T, T, T] # the expected position ids is [0, 1, 2, 2, 2, 3, 4, 5, 6] # see https://huggingface.co/THUDM/cogvlm2-llama3-chat-19B/blob/main/modeling_cogvlm.py#318 tmp = token_type_ids.numpy().copy() is_boi_eoi = np.zeros_like(tmp, dtype=bool) is_boi_eoi[1:] |= (tmp[1:] == VISION_TOKEN_TYPE) & (tmp[:-1] == LANGUAGE_TOKEN_TYPE) is_boi_eoi[0] |= tmp[0] == VISION_TOKEN_TYPE is_boi_eoi[:-1] |= (tmp[:-1] == VISION_TOKEN_TYPE) & (tmp[1:] == LANGUAGE_TOKEN_TYPE) is_boi_eoi[-1] |= tmp[-1] == VISION_TOKEN_TYPE tmp[is_boi_eoi] = LANGUAGE_TOKEN_TYPE y = np.zeros_like(tmp, dtype=np.int32) y[1:] = (tmp[1:] == LANGUAGE_TOKEN_TYPE) | ( (tmp[1:] == VISION_TOKEN_TYPE) & (tmp[:-1] == LANGUAGE_TOKEN_TYPE) ) y = y.cumsum() return y.tolist() def expand_token_id( self, token_ids: List[int], images: List[torch.Tensor] ) -> Tuple[List[int], List[torch.Tensor], List[int]]: if len(images) > 1: raise Exception("CogVLM2 support processes one image at a time") input_ids = [self.tokenizer.bos_token_id] token_type_ids = [LANGUAGE_TOKEN_TYPE] if len(images) > 0: patch_size: int = self.config.mm_related_params.config["patch_size"] image_size: int = self.config.mm_related_params.config["image_size"] vision_token_num = (image_size // patch_size // 2) * ( image_size // patch_size // 2 ) + 2 input_ids += [self.tokenizer.pad_token_id] * vision_token_num token_type_ids += [VISION_TOKEN_TYPE] * vision_token_num input_ids += token_ids token_type_ids += [LANGUAGE_TOKEN_TYPE] * len(token_ids) return input_ids, images, token_type_ids def multimodal_embedding( self, input_ids: torch.Tensor, image_features: List[torch.Tensor], token_type_ids: torch.Tensor, ): token_type_ids = token_type_ids.reshape(1, -1) input_embeds = self.word_embedding(input_ids) if len(image_features) > 0: image_features = torch.stack(image_features, dim=0) image_features = rearrange(image_features, "b n d -> (b n) d") image_features = image_features.to( dtype=input_embeds.dtype, device=input_embeds.device ) input_embeds = input_embeds.index_put( [token_type_ids == VISION_TOKEN_TYPE], image_features ) return input_embeds register_model( "cogvlm2", CogVLM2, ["CogVLMForCausalLM"], [ "THUDM/cogvlm2-llama3-chat-19B", "THUDM/cogvlm2-llama3-chinese-chat-19B" ], )