import logging import os import functools import json from typing import List, Any, Dict from maga_transformer.utils.model_weight import (W, WeightStyle, CkptWeightInfo, identity, sp_0, sp_head_lora, sp_id, sp_neg1, zeros, transpose, transpose_pad, merge_qkv_b, merge_qkv_hf, merge_qkv_lora_A, merge_qkv_lora_B) from maga_transformer.config.gpt_init_model_parameters import GptInitModelParameters from maga_transformer.models.qwen import QWen from transformers import AutoTokenizer from maga_transformer.model_factory_register import register_model from maga_transformer.model_loader.weight_module import WeightModule, AtomicWeight from maga_transformer.model_loader.ffn_weight import FfnAtomicWeight, FfnWeight, FfnConfig from maga_transformer.model_loader.attn_weight import AttnAtomicWeight, AttnConfig from maga_transformer.model_loader.ffn_weight import FfnWeight, FfnAtomicWeight, FfnConfig from maga_transformer.model_loader.attn_weight import AttnAtomicWeight, AttnConfig from maga_transformer.model_loader.model_weight_info import ModelWeightInfo, ModelDeployWeightInfo def scale_reshape(ts): return ts[0].reshape(-1) class QWenV2Weight(ModelDeployWeightInfo): def __init__(self, *args: Any, **kwargs: Any): self.prefix: str = kwargs.pop('prefix', "") super().__init__(*args, **kwargs) @property def support_lora(self): return True def _process_meta(self, meta_dicts: Any, weight_keys: List[str]): # compat for qwen_v2_video if self._contains(weight_keys, 'language_model.'): self.prefix = 'language_model.' if self.prefix + 'transformer.layers.0.attention.dense.weight' in meta_dicts[0]: self.weight_style = WeightStyle.TRT_ENGINE logging.info(f"weight_style: {self.weight_style}") def _get_weight_info(self): return self._get_hf_weight_info() def _get_hf_ffn_layer_weight_info(self, layer_id: int) -> List[WeightModule]: ffn_config = FfnConfig( is_gated_activation=self._is_gated_activation, inter_padding_size=self._inter_padding_size, is_moe=False ) inter_padding_size = self._layer_inter_padding_size[layer_id] if self._layer_inter_padding_size else self._inter_padding_size return [FfnWeight(sub_weights=[ FfnAtomicWeight(W.ffn_w1, [CkptWeightInfo(self.prefix + 'model.layers.{i}.mlp.gate_proj.weight', identity)], functools.partial(transpose_pad, inter_padding_size=inter_padding_size, dim=0), config=ffn_config, lora_a_process_func=transpose, lora_b_process_func=functools.partial(transpose_pad, inter_padding_size=inter_padding_size, dim=0), lora_a_split_func=sp_id, lora_b_split_func=sp_neg1), FfnAtomicWeight(W.ffn_w3, [CkptWeightInfo(self.prefix + 'model.layers.{i}.mlp.up_proj.weight', identity)], functools.partial(transpose_pad, inter_padding_size=inter_padding_size, dim=0), config=ffn_config, lora_a_process_func=transpose, lora_b_process_func=functools.partial(transpose_pad, inter_padding_size=inter_padding_size, dim=0), lora_a_split_func=sp_id, lora_b_split_func=sp_neg1), FfnAtomicWeight(W.ffn_w2, [CkptWeightInfo(self.prefix + 'model.layers.{i}.mlp.down_proj.weight', identity)], functools.partial(transpose_pad, inter_padding_size=inter_padding_size, dim=1), config=ffn_config, lora_a_process_func=functools.partial(transpose_pad, inter_padding_size=inter_padding_size, dim=1), lora_b_process_func=transpose, lora_a_split_func=sp_0, lora_b_split_func=sp_id) ], config=ffn_config) ] def _get_hf_layer_weight_info(self, layer_id: int): attn_config = AttnConfig( hidden_size=self._hidden_size, size_per_head=self._size_per_head, head_num=self._head_num, head_num_kv=self._head_num_kv) layer_weights = [ AtomicWeight(W.pre_ln_gamma, [CkptWeightInfo(self.prefix + 'model.layers.{i}.input_layernorm.weight', identity)], identity), AttnAtomicWeight(W.attn_qkv_b, [ CkptWeightInfo(self.prefix + 'model.layers.{i}.self_attn.q_proj.bias', identity), CkptWeightInfo(self.prefix + 'model.layers.{i}.self_attn.k_proj.bias', identity), CkptWeightInfo(self.prefix + 'model.layers.{i}.self_attn.v_proj.bias', identity) ], functools.partial(merge_qkv_b), config=attn_config), AttnAtomicWeight(W.attn_qkv_w, [ CkptWeightInfo(self.prefix + 'model.layers.{i}.self_attn.q_proj.weight', identity), CkptWeightInfo(self.prefix + 'model.layers.{i}.self_attn.k_proj.weight', identity), CkptWeightInfo(self.prefix + 'model.layers.{i}.self_attn.v_proj.weight', identity) ], functools.partial(merge_qkv_hf), config=attn_config, lora_a_process_func=functools.partial(merge_qkv_lora_A, allow_empty=False, hidden_size=self._hidden_size, head_num=self._head_num, head_num_kv=self._head_num_kv, size_per_head=self._size_per_head), lora_b_process_func=functools.partial(merge_qkv_lora_B, allow_empty=False, hidden_size=self._hidden_size, head_num=self._head_num, head_num_kv=self._head_num_kv, size_per_head=self._size_per_head), lora_a_split_func=sp_id, lora_b_split_func=sp_head_lora ), AttnAtomicWeight(W.attn_o_w, [CkptWeightInfo(self.prefix + 'model.layers.{i}.self_attn.o_proj.weight', identity)], transpose, config=attn_config, lora_a_process_func=transpose, lora_b_process_func=transpose, lora_a_split_func=sp_0, lora_b_split_func=sp_id), AtomicWeight(W.post_ln_gamma, [CkptWeightInfo(self.prefix + 'model.layers.{i}.post_attention_layernorm.weight', identity)], identity, config=attn_config), ] if self._use_qk_norm: layer_weights.extend([ AttnAtomicWeight(W.q_ln_gamma, [CkptWeightInfo(self.prefix + 'model.layers.{i}.self_attn.q_norm.weight')], config=attn_config), AttnAtomicWeight(W.k_ln_gamma, [CkptWeightInfo(self.prefix + 'model.layers.{i}.self_attn.k_norm.weight')], config=attn_config)] ) layer_weights.extend(self._get_hf_ffn_layer_weight_info(layer_id)) return layer_weights def _get_hf_weight_info(self): if self.weight_style == WeightStyle.TRT_ENGINE: weights = [ AtomicWeight(W.embedding, [CkptWeightInfo('transformer.vocab_embedding.weight', identity)], identity), AtomicWeight(W.lm_head, [CkptWeightInfo('lm_head.weight', identity)], identity), AtomicWeight(W.final_ln_gamma, [CkptWeightInfo('transformer.ln_f.weight', identity)], identity), AtomicWeight(W.final_ln_beta, [], functools.partial(zeros, shape=[self._hidden_size])), ] else: weights = [ AtomicWeight(W.embedding, [CkptWeightInfo(self.prefix + 'model.embed_tokens.weight', identity)], identity), AtomicWeight(W.lm_head, [CkptWeightInfo(self.prefix + 'lm_head.weight', identity)], identity), AtomicWeight(W.final_ln_gamma, [CkptWeightInfo(self.prefix + 'model.norm.weight', identity)], identity), AtomicWeight(W.final_ln_beta, [], functools.partial(zeros, shape=[self._hidden_size])), ] layer_weights: List[List[WeightModule]] = [] for layer in range(self._num_layers): w = self._get_hf_layer_weight_info(layer) layer_weights.append(w) return ModelWeightInfo(layer_weights=layer_weights, weights=weights) class QWenV2(QWen): @classmethod def _create_config(cls, ckpt_path: str): config = GptInitModelParameters( head_num=0, head_num_kv=0, size_per_head=0, layer_num=0, inter_size=0, # 13696 vocab_size=152064, max_seq_len=8192) config.rotary_embedding_dim = 128 config.rotary_embedding_style = 1 config.activation_type = 'SiGLU' config.has_pre_decoder_layernorm = False config.has_post_decoder_layernorm = True config.norm_type = 'rmsnorm' config.special_tokens.bos_token_id = -1 config.special_tokens.eos_token_id = 151643 # <|im_start|> and <|im_end|> config.special_tokens.stop_words_id_list = [[151645], [151644]] config.special_tokens.system.token_ids = [151644, 8948, 198] # '<|im_start|>system\n' config.special_tokens.system.eos_token_ids = [151645, 198] # '<|im_end|>\n' config.special_tokens.user.token_ids = [151644, 872, 198] # '<|im_start|>user\n' config.special_tokens.user.eos_token_ids = [151645, 198] # '<|im_end|>\n' config.special_tokens.assistant.token_ids = [151644, 77091, 198] # '<|im_start|>assistant\n' config.special_tokens.assistant.eos_token_ids = [151645, 198] # '<|im_end|>\n' cls._from_hf(config, ckpt_path) assert config.head_num > 0 and config.head_num_kv > 0 and config.size_per_head > 0 and config.layer_num > 0 and config.inter_size > 0, "error config" return config @classmethod def _from_hf(cls, config: GptInitModelParameters, ckpt_path: str): config_path = os.path.join(ckpt_path, "config.json") if not os.path.exists(config_path): return with open(config_path) as reader: content = reader.read() config_json = json.loads(content) QWenV2._from_config_json(config, config_json) return config @staticmethod def _from_config_json(config: GptInitModelParameters, config_json: Dict[str, Any]): # config.activation_type = config_json["hidden_act"] config.inter_size = config_json["intermediate_size"] config.head_num = config_json["num_attention_heads"] config.head_num_kv = config_json.get("num_key_value_heads", config.head_num) config.size_per_head = int(config_json.get("head_dim")) if "head_dim" in config_json else config_json["hidden_size"] // config.head_num if config_json.get("hidden_size") is not None: config.hidden_size = config_json["hidden_size"] config.layer_num = config_json["num_hidden_layers"] config.rotary_embedding_base = config_json.get("rope_theta", config.rotary_embedding_base) config.vocab_size = config_json["vocab_size"] config.rotary_embedding_dim = config.size_per_head config.layernorm_eps = config_json.get("rms_norm_eps", 1e-06) config.tie_word_embeddings = config_json.get('tie_word_embeddings', False) @staticmethod def get_weight_cls(): return QWenV2Weight @classmethod def get_tokenizer(cls, config: GptInitModelParameters): tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_path, verbose=False, trust_remote_code=True) tokenizer.im_start_id = tokenizer.encode('<|im_start|>')[0] tokenizer.im_end_id = tokenizer.encode('<|im_end|>')[0] return tokenizer class QWenV2Embedding(QWenV2): @classmethod def _create_config(cls, ckpt_path: str): config = QWenV2._create_config(ckpt_path) config.is_causal = False return config register_model('qwen_2', QWenV2, ["Qwen2ForCausalLM"]) register_model('qwen_agent', QWenV2) register_model('qwen_2_embedding', QWenV2Embedding) register_model("qwen_tool", QWenV2)