通過onnx模型可以在支持onnx推理的推理引擎上進(jìn)行推理，從而可以將LLM部署在更加廣泛的平臺上面。此外還可以具有避免pytorch依賴，獲得更好的性能等優(yōu)勢。
這篇博客（大模型LLaMa及周邊項目（二） - 知乎）進(jìn)行了llama導(dǎo)出onnx的開創(chuàng)性的工作，但是依賴于侵入式修改transformers庫，比較不方便。
這里本人實現(xiàn)了避免侵入式修改transformers庫導(dǎo)出LLM為ONNX方法，代碼庫為:
https://github.com/luchangli03/export_llama_to_onnx
導(dǎo)出的LLM進(jìn)行onnxsim優(yōu)化：
一種大于2GB ONNX模型onnxsim優(yōu)化方法_Luchang-Li的博客-CSDN博客
https://github.com/luchangli03/onnxsim_large_model
可以在這個基礎(chǔ)上進(jìn)行簡單修改從而導(dǎo)出其他模型，例如百川，Qwen等模型。當(dāng)前已經(jīng)加入了對llama, 清華ChatGlm2和阿里Qwen的導(dǎo)出支持。百川跟llamma相似性很高，可以直接用llama的導(dǎo)出腳本。
除了導(dǎo)出onnx，目前還增加了一個基于onnx推理chatglm2的demo。
該方案優(yōu)點是無需侵入式修改transformers代碼，缺點是你需要提前了解各個模型的輸入，相關(guān)shape和dtype。這可以在代碼里面增加shape，dtype打印，進(jìn)行一次推理獲取。
阿里的Qwen模型導(dǎo)出onnx存在大量瑣碎的算子，主要是因為其使用了einops庫的rearrange操作。把這些算子替換為reshape，transpose, split等操作后onnx模型會極大簡化。此外，還可以進(jìn)一步優(yōu)化該模型的RotaryEmbedding相關(guān)的代碼以優(yōu)化導(dǎo)出的onnx模型復(fù)雜度。
例如
# context_layer = rearrange(
# context_layer, "b s h d -> b s (h d)"
# ).contiguous()
b, s, h, d = context_layer.shape
context_layer = context_layer.reshape([b, s, -1])# self._rotary_pos_emb_cache = rearrange(emb, "n d -> 1 n 1 d")
emb = torch.unsqueeze(emb, 0) # [1nd]
emb = torch.unsqueeze(emb, 2)def _rotate_half(x):
# from einops import rearrange
# x = rearrange(x, "... (j d) -> ... j d", j=2)
# x1, x2 = x.unbind(dim=-2)
last_dim = x.shape[-1]
x1, x2 = torch.split(x, last_dim//2, dim=-1)
return torch.cat((-x2, x1), dim=-1)
導(dǎo)出的llama decoder會存在if算子，但是經(jīng)過符號shape推導(dǎo)和設(shè)置相應(yīng)的符號shape到onnx模型value_info，然后經(jīng)過onnxsim可以完全去掉。也可以考慮修改llama定義代碼去掉if。優(yōu)化前后：
導(dǎo)出模型調(diào)試
從導(dǎo)出模型中算子的名稱可以大概推斷出該算子在那個模塊定義的。文章來源：http://www.zghlxwxcb.cn/news/detail-638239.html
如/layers.0/attn/rotary_emb/Slice可以知道該算子定義于第0個decoder層的atten.rotary_emb模塊，而slice一般來自于張量切片索引操作。文章來源地址http://www.zghlxwxcb.cn/news/detail-638239.html
到了這里，關(guān)于導(dǎo)出LLaMA ChatGlm2等LLM模型為onnx的文章就介紹完了。如果您還想了解更多內(nèi)容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關(guān)文章，希望大家以后多多支持TOY模板網(wǎng)！
