support lora weight export.

Author: wangzhaode

README.md

@@ -8,7 +8,7 @@ llm-export是一个llm模型导出工具，能够将llm模型导出为onnx和mnn
 - 🚀 优化原始代码，支持动态形状
 - 🚀 优化原始代码，减少常量部分
 - 🚀 使用[OnnxSlim](https://github.com/WeLoveAI/OnnxSlim)优化onnx模型，性能提升约5%; by [@inisis](https://github.com/inisis)
-
+- 🚀 支持将lora权重导出为onnx和mnn
 
 ## 模型支持与下载
 - [![Download][download-chatglm-6b-onnx]][release-chatglm-6b-onnx]

README_en.md

@@ -7,7 +7,7 @@ llm-export is a tool for exporting llm models, capable of converting llm models
 - 🚀 Optimized the original code to support dynamic shapes
 - 🚀 Optimized the original code to reduce the constant portion
 - 🚀 Using [OnnxSlim](https://github.com/WeLoveAI/OnnxSlim) slim onnx model，speed up 5%; by [@inisis](https://github.com/inisis)
-
+- 🚀 Support export lora weight to onnx or MNN model
 
 ## Model Support and Downloads
 

llm_export.py

@@ -990,6 +990,28 @@ def get_position_ids(self) -> torch.Tensor:
     def get_attention_mask(self) -> torch.Tensor:
         return torch.ones([1, 1, 1, self.seq_len], dtype=torch.long)
 
+class LoraModule(torch.nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.onnx_path = args.onnx_path
+        self.mnn_path = args.mnn_path
+        self.export_mnn = args.export_mnn
+        import peft
+        lora_weight = peft.load_peft_weights(args.path)
+        for k, v in lora_weight.items():
+            k = k.replace('.', '/')
+            self.register_buffer(k, v.cpu())
+
+    def forward(self, dummpy):
+        return self._buffers
+
+    def export(self):
+        onnx_model = f'./{self.onnx_path}/lora.onnx'
+        torch.onnx.export(self.eval(), torch.tensor([]), onnx_model)
+        if self.export_mnn:
+            onnx2mnn(onnx_model, self.mnn_path)
+
+
 if __name__ == '__main__':
     llm_models = {
         'chatglm-6b': Chatglm_6b,
@@ -1006,7 +1028,8 @@ def get_attention_mask(self) -> torch.Tensor:
         'Yi-6B-Chat': Llama2_7b_Chat,
         'deepseek-llm-7b-chat': Llama2_7b_Chat,
         'phi-2': phi_2,
-        'bge-large-zh': bge
+        'bge-large-zh': bge,
+        'lora': LoraModule
     }
     parser = argparse.ArgumentParser(description='llm_exporter', formatter_class=argparse.RawTextHelpFormatter)
     parser.add_argument('--path', type=str, default='THUDM/chatglm-6b', required=True,