support export mnn model.

Author: wangzhaode

README.md

@@ -48,7 +48,7 @@ git clone https://modelscope.cn/ZhipuAI/chatglm2-6b.git
 3. 执行LLMExporter导出模型
 ```sh
 cd LLMExporter
-python llm_export.py --path ../chatglm2-6b --export_path ./onnx --export
+python llm_export.py --path ../chatglm2-6b --onnx_path ./onnx --export_split --export_mnn --export_token
 ```
 
 ## 功能
@@ -61,11 +61,13 @@ python llm_export.py --path ../chatglm2-6b --export_path ./onnx --export
 - 支持对模型进行对话测试，使用`--test $query`会返回llm的回复内容
 - 支持在导出onnx模型后使用onnxruntime对结果一致性进行校验，使用`--export_test`
 - 支持将tokenizer导出为文本文件，使用`--export_token`
+- 支持将导出的onnx模型转换为mnn模型，默认转换为非对称4bit量化，使用`--export_mnn`
+- 指定导出路径使用`--onnx_path`和`--mnn_path`
 
 ## 参数
 ```
-usage: llm_export.py [-h] --path PATH [--type {chatglm-6b,chatglm2-6b,codegeex2-6b,Qwen-7B-Chat,Baichuan2-7B-Chat,Llama-2-7b-chat-ms}]
-                     [--export_path EXPORT_PATH] [--export_verbose] [--export_test] [--test TEST] [--export] [--export_split] [--export_token]
+usage: llm_export.py [-h] --path PATH [--type {chatglm-6b,chatglm2-6b,chatglm3-6b,codegeex2-6b,Qwen-7B-Chat,Qwen-1_8B-Chat,Baichuan2-7B-Chat,Llama-2-7b-chat-ms}]
+                     [--onnx_path ONNX_PATH] [--mnn_path MNN_PATH] [--export_mnn] [--export_verbose] [--export_test] [--test TEST] [--export] [--export_split] [--export_token]
                      [--export_embed] [--export_lm] [--export_block EXPORT_BLOCK] [--export_blocks] [--embed_bf16]
 
 LLMExporter
@@ -76,11 +78,13 @@ optional arguments:
                         Can be either:
                         	- A string, the *model id* of a pretrained model like `THUDM/chatglm-6b`. [TODO]
                         	- A path to a *directory* clone from repo like `../chatglm-6b`.
-  --type {chatglm-6b,chatglm2-6b,codegeex2-6b,Qwen-7B-Chat,Baichuan2-7B-Chat,Llama-2-7b-chat-ms}
+  --type {chatglm-6b,chatglm2-6b,chatglm3-6b,codegeex2-6b,Qwen-7B-Chat,Qwen-1_8B-Chat,Baichuan2-7B-Chat,Llama-2-7b-chat-ms}
                         type(`str`, *optional*):
                         	The pretrain llm model type.
-  --export_path EXPORT_PATH
+  --onnx_path ONNX_PATH
                         export onnx model path, defaut is `./onnx`.
+  --mnn_path MNN_PATH   export mnn model path, defaut is `./mnn`.
+  --export_mnn          Whether or not to export mnn model after onnx.
   --export_verbose      Whether or not to export onnx with verbose.
   --export_test         Whether or not to export onnx with test using onnxruntime.
   --test TEST           test model inference with query `TEST`.

llm_export.py

@@ -6,9 +6,33 @@
 import torch
 import numpy as np
 import onnxruntime as ort
+import _tools as MNNTools
 import sentencepiece as spm
 from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer
 
+def onnx2mnn(onnx_path, mnn_dir, quant_bit = 4, asymmetric = True, external_data = False):
+    model_name, model_extension = os.path.splitext(os.path.basename(onnx_path))
+    if model_extension != '.onnx':
+        return
+    mnn_name = model_name + '.mnn'
+    mnn_path = os.path.join(mnn_dir, mnn_name)
+    convert_args = [
+        '',
+        '-f',
+        'ONNX',
+        '--modelFile',
+        str(onnx_path),
+        '--MNNModel',
+        str(mnn_path),
+        '--weightQuantBits',
+        str(quant_bit)
+    ]
+    if asymmetric:
+        convert_args.append("--weightQuantAsymmetric")
+    if external_data:
+        convert_args.append("--saveExternalData")
+    MNNTools.mnnconvert(convert_args)
+
 # some wrapper class for export
 class Embedding(torch.nn.Module):
     def __init__(self, embed, using_bf16: bool = False):
@@ -44,7 +68,13 @@ class LLM(torch.nn.Module):
 
     def __init__(self, args):
         super().__init__()
-        self.export_path = args.export_path
+        self.onnx_path = args.onnx_path
+        self.mnn_path = args.mnn_path
+        if not os.path.exists(self.onnx_path):
+            os.makedirs(self.onnx_path)
+        if not os.path.exists(self.mnn_path):
+            os.makedirs(self.mnn_path)
+        self.export_mnn = args.export_mnn
         self.export_verbose = args.export_verbose
         self.export_test = args.export_test
         self.embed_bf16 = args.embed_bf16
@@ -134,7 +164,7 @@ def assert_equal(self, torch_outs, onnx_outs):
     def export_lm(self):
         model = self.lm
         hidden_states = torch.randn(1, self.hidden_size)
-        onnx_model = f'./{self.export_path}/lm.onnx'
+        onnx_model = f'./{self.onnx_path}/lm.onnx'
         torch.onnx.export(model, (hidden_states),
                         onnx_model,
                         verbose=self.export_verbose,
@@ -151,11 +181,13 @@ def export_lm(self):
             }
             onnx_outs = ort_session.run(None, inputs)
             self.assert_equal(original_outs, onnx_outs)
+        if self.export_mnn:
+            onnx2mnn(onnx_model, self.mnn_path)
 
     def export_embed(self):
         model = self.embed
         input_ids = torch.arange(3, dtype=torch.long)
-        onnx_model = f'./{self.export_path}/embedding.onnx'
+        onnx_model = f'./{self.onnx_path}/embedding.onnx'
         torch.onnx.export(model, (input_ids),
                         onnx_model,
                         verbose=self.export_verbose,
@@ -175,6 +207,8 @@ def export_embed(self):
             }
             onnx_outs = ort_session.run(None, inputs)
             self.assert_equal(original_outs, onnx_outs)
+        if self.export_mnn:
+            onnx2mnn(onnx_model, self.mnn_path)
 
     def export_block(self, block_id: int):
         self.seq_len = 3
@@ -184,7 +218,7 @@ def export_block(self, block_id: int):
         position_ids = self.get_position_ids()
         past_key_values = torch.zeros(self.past_kv_shape[1:])
         model = self.blocks[block_id]
-        onnx_model = f'./{self.export_path}/block_{block_id}.onnx'
+        onnx_model = f'./{self.onnx_path}/block_{block_id}.onnx'
         torch.onnx.export(
             model, (inputs_embeds, attention_mask, position_ids, past_key_values),
             onnx_model,
@@ -207,6 +241,8 @@ def export_block(self, block_id: int):
             }
             onnx_outs = ort_session.run(None, inputs)
             self.assert_equal(original_outs, onnx_outs)
+        if self.export_mnn:
+            onnx2mnn(onnx_model, self.mnn_path)
 
     def export_blocks(self):
         for i in range(self.block_nums):
@@ -220,7 +256,7 @@ def export(self):
         attention_mask =  self.get_attention_mask()
         position_ids = self.get_position_ids()
         past_key_values = torch.zeros(self.past_kv_shape)
-        onnx_model = f'./{self.export_path}/llm.onnx'
+        onnx_model = f'./{self.onnx_path}/llm.onnx'
         torch.onnx.export(
             model, (input_ids, attention_mask, position_ids, past_key_values),
             onnx_model,
@@ -244,9 +280,12 @@ def export(self):
             }
             onnx_outs = ort_session.run(None, inputs)
             self.assert_equal(original_outs, onnx_outs)
+        if self.export_mnn:
+            # single model is > 2G, using external_data
+            onnx2mnn(onnx_model, self.mnn_path, 4, True, True)
 
     def export_tokenizer(self):
-        file_path = os.path.join(self.export_path, "tokenizer.txt")
+        file_path = os.path.join(self.onnx_path, "tokenizer.txt")
         if self.sp_model is not None:
             # senetencepiece
             NORMAL = 1; UNKNOWN = 2; CONTROL = 3
@@ -644,7 +683,9 @@ def get_position_ids(self) -> torch.Tensor:
                         help='type(`str`, *optional*):'
                         '\n\tThe pretrain llm model type.'
                         )
-    parser.add_argument('--export_path', type=str, default='./onnx', help='export onnx model path, defaut is `./onnx`.')
+    parser.add_argument('--onnx_path', type=str, default='./onnx', help='export onnx model path, defaut is `./onnx`.')
+    parser.add_argument('--mnn_path', type=str, default='./mnn', help='export mnn model path, defaut is `./mnn`.')
+    parser.add_argument('--export_mnn', action='store_true', default=False, help='Whether or not to export mnn model after onnx.')
     parser.add_argument('--export_verbose', action='store_true', default=False, help='Whether or not to export onnx with verbose.')
     parser.add_argument('--export_test', action='store_true', help='Whether or not to export onnx with test using onnxruntime.')
     parser.add_argument('--test', type=str, help='test model inference with query `TEST`.')

requirements.txt

@@ -1,3 +1,4 @@
+MNN==2.8.0
 numpy==1.25.2
 onnxruntime==1.15.1
 torch==2.0.1