diff --git a/template/compile/model_export.py b/template/compile/model_export.py index 36527dd3..850d10d3 100755 --- a/template/compile/model_export.py +++ b/template/compile/model_export.py @@ -441,6 +441,12 @@ def load_pretrained(self): self.config = AutoConfig.from_pretrained(self.model_path, trust_remote_code=True) self.model_type = self.config.model_type + if 'chatglm3' in self.model_path: + self.config = AutoConfig.from_pretrained(self.model_path, trust_remote_code=True, model_type='chatglm3') + print('chatglm3 model type:', self.model.config) + self.model_type = self.config.model_type + self.config.num_hidden_layers = self.config.num_layers + if 'qwen2_vl' == self.model_type: from transformers import Qwen2VLForConditionalGeneration self.model = Qwen2VLForConditionalGeneration.from_pretrained(self.model_path) @@ -464,6 +470,8 @@ def load_pretrained(self): elif "Model" in self.config.architectures[0]: self.model = AutoModel.from_pretrained( self.model_path, trust_remote_code=True) + if "chatglm3" in self.model_path: + self.model.config.model_type = "chatglm3" else: raise ValueError(f"Unsupported Architectures:[ {self.config.architectures[0]} ]") diff --git a/template/compile/onnx_rebuilder.py b/template/compile/onnx_rebuilder.py index a0072392..450e4491 100755 --- a/template/compile/onnx_rebuilder.py +++ b/template/compile/onnx_rebuilder.py @@ -46,6 +46,7 @@ def regist_models(self): self.regist_qwen2_5_vl() self.regist_glm() self.regist_glm2() + self.regist_glm3() self.regist_phi() def default_map(self): @@ -204,6 +205,34 @@ def regist_glm2(self): } self.regist('chatglm2', glm2_map) + def regist_glm3(self): + glm3_map = { + 'config': { + 'hidden_size': 'hidden_size', + 'num_attention_heads': 'num_attention_heads', + 'num_key_value_heads': 'multi_query_group_num', + 'num_hidden_layers': 'num_layers', + 'vocab_size': 'vocab_size' + }, + 'model': { + 'lm_': 'transformer.output_layer', + 'embed_': 'transformer.embedding.word_embeddings', + 'blocks_': 'transformer.encoder.layers', + 'final_layernorm_': 'transformer.encoder.final_layernorm', + }, + 'decoder': { + 'self_attn': 'self_attention', + 'mlp': 'mlp', + 'input_layernorm': 'input_layernorm', + 'post_attention_layernorm': 'post_attention_layernorm' + }, + 'attention': { + 'qkv_proj': 'query_key_value', + 'o_proj': 'dense' + } + } + self.regist('chatglm3', glm3_map) + def regist_phi(self): phi_map = { 'config': { @@ -272,6 +301,8 @@ def __init__(self, self.config = config self.visual_length = visual_length self.model_mapper = ModelMapper() + self.num_key_value_heads = None + self.rope_theta = None def _replace_initializer(self, old_init, new_init): """ @@ -685,7 +716,7 @@ def init_rotary_pos_emb(self, seq_length): position_ids = position_ids.float().reshape(-1, 1) idx_theta = position_ids * theta rotary_pos_emb = torch.stack([torch.cos(idx_theta), torch.sin(idx_theta)]) - if self.model_type != 'chatglm2': + if self.model_type != 'chatglm2' and self.model_type != 'chatglm3': rotary_pos_emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) rotary_pos_emb = rotary_pos_emb.unsqueeze(2).unsqueeze(1) return rotary_pos_emb @@ -695,6 +726,8 @@ def apply_rotary_pos(self, x, cos, sin): return self.chatglm_rotary_pos(x, cos, sin) if self.model_type == 'chatglm2': return self.chatglm2_rotary_pos(x, cos, sin) + if self.model_type == 'chatglm3': + return self.chatglm3_rotary_pos(x, cos, sin) if self.model_type == 'phi-msft': return self.phi_rotary_pos(x, cos, sin) return self.llama_rotary_pos(x, cos, sin) @@ -721,6 +754,19 @@ def chatglm2_rotary_pos(self, x, cos, sin): ) return torch.cat((x, x_pass), dim=-1) + def chatglm3_rotary_pos(self, x, cos, sin): + x, x_pass = x[..., :self.rotary_dim], x[..., self.rotary_dim:] + b, s, n, h = x.shape + xshaped = x.view(b, s, n, h//2, 2) + x = torch.concat( + [ + xshaped[..., 0] * cos - xshaped[..., 1] * sin, + xshaped[..., 1] * cos + xshaped[..., 0] * sin, + ], + -1, + ) + return torch.cat((x, x_pass), dim=-1) + def chatglm_rotary_pos(self, x, cos, sin): seq = x.shape[1] x1, x2 = x[..., :self.rotary_dim], x[..., self.rotary_dim:] @@ -928,6 +974,8 @@ def apply_rotary_pos(self, x, cos, sin): return self.chatglm_rotary_pos(x, cos, sin) if self.model_type == 'chatglm2': return self.chatglm2_rotary_pos(x, cos, sin) + if self.model_type == 'chatglm3': + return self.chatglm3_rotary_pos(x, cos, sin) if self.model_type == 'phi-msft': return self.phi_rotary_pos(x, cos, sin) return self.llama_rotary_pos(x, cos, sin) diff --git a/template/demo/pipeline.py b/template/demo/pipeline.py index 670c3804..497ae4d2 100644 --- a/template/demo/pipeline.py +++ b/template/demo/pipeline.py @@ -58,6 +58,15 @@ def map(self, args): history, tokenize=False, add_generation_prompt=True ) self.system_prompt = {"role": "system", "content": "You are a helpful assistant."} + elif self.model_type == "chatglm3": + self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_path, trust_remote_code=True) + self.EOS = [2] + self.append_user = lambda history, input_str: history.append( + "<|user|>\n{}\n<|assistant|>\n".format(input_str) + ) + self.append_assistant = lambda history, answer_str: history.append(answer_str) + self.apply_chat_template = lambda history: "".join(history) + self.system_prompt = "[gMASK]sop<|system|>\nYou are a helpful assistant.\n" elif self.model_type == "qwen": self.tokenizer = AutoTokenizer.from_pretrained(self.tokenizer_path, trust_remote_code=True) self.EOS = [self.tokenizer.im_end_id] @@ -169,7 +178,8 @@ def init_params(self, args): self.model.max_new_tokens = args.max_new_tokens self.model.generation_mode = args.generation_mode self.model.embedding_path = self.embedding_path if os.path.exists(self.embedding_path) else "" - self.model.NUM_LAYERS = self.config["num_hidden_layers"] + if self.model_type == "chatglm3": + self.model.NUM_LAYERS = self.config["num_layers"] self.model.config.model_type = self.model_type self.max_new_tokens = args.max_new_tokens