feat(template): add support for chatglm3-6b

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]} ]")
 

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)

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