讓 Transformer 的推理速度提高 4.5 倍，這個小 trick 還能給你省十幾萬
點擊上方“程序員大白”，選擇“星標”公眾號
重磅干貨，第一時間送達
豐色 發(fā)自 凹非寺
量子位 報道 | 公眾號 QbitAI
最近，NLP明星公司Hugging Face發(fā)布了一個叫做Infinity的產(chǎn)品，可以以1ms延時完成Transformer的推理，性能相當高了。
但是，厲害歸厲害，還是有點貴——1年至少要十幾萬塊?（2萬美元）。
那有沒有什么平替的方法呢？
有的！還是開源的、“不費吹灰之力”就可以達到Infinity一些公共基準的那種。
并且現(xiàn)在，通過在該方法上施加一個小trick，將Transformer的推理速度提高4.5倍！
那么，一個“平替”到底為什么能達到“付費”的效果呢？
一個trick讓Transformer推理速度提高4.5倍
先來認識一下這個方法：Transformer-deploy。
它可以用一行命令優(yōu)化和部署Hugging Face上的Transformer模型，并支持大多數(shù)基于Transformer編碼器的模型，比如Bert、Roberta、miniLM、Camembert、Albert、XLM-R、Distilbert等。
Transformer-deploy推理服務器用的是Nvidia Triton。
推理引擎為Microsoft ONNX Runtime（用于CPU和GPU推理）和Nvidia TensorRT（僅限 GPU）。
如果想在GPU上獲得一流的性能，Nvidia Triton+Nvidia TensorRT這樣的組合無疑是最佳選擇。
雖然TensorRT用起來有點難，但它確實能比用Pytorch快5～10倍。
在實際性能測試中，Transformer-deploy在batch size為1、token分別為16和128的輸入序列中的推理速度，都比付費的Hugging Face Infinity要快：
Transformer-deploy在token為16時要1.52ms，Infinity則需要1.7ms；token為128時需要1.99ms，Infinity則需要2.5ms。
那前面說的能讓Transformer的推理性能進一步提高的小trick是什么呢？
GPU量化（quantization）。
作者表示：
據(jù)我所知，目前任何OOS云服務都還沒用到過這個方法。
不過執(zhí)行GPU量化需要修改模型源代碼（需在矩陣乘法等代價高昂的操作上添加一些叫做QDQ的特定節(jié)點），既容易出錯，又很無聊，并且還需自己維護修改后的代碼。
因此作者已經(jīng)為多個基于Transformer的模型手動完成了這項工作。
后來，他們又發(fā)現(xiàn)似乎只需修補模型模塊的抽象語法樹?（AST）也可以自動完成。
在用戶端，在GPU上執(zhí)行模型的基本量化類似這樣：
最終，該方法在Roberta-base模型和MNLI數(shù)據(jù)集（分類任務）上實現(xiàn)了4.53倍的推理速度。
當然這也犧牲了0.4個點的精度；如果一點不犧牲的話，也可以加速3.2倍左右。
作者表示，與Transformer-deploy原來的版本相比，這已經(jīng)是一個很大的改進了，畢竟原版本的加速成本需要超過1個點的精確度。
最終他們用Albert、Bert（包括miniLM）、Distilbert、Roberta（包括 Camembert、XLM-R、DistilRoberta等）、Electra測試了該trick。
結果是對于任何可以導出為ONNX格式的Transformer模型，都可以“開箱即用”。
參考鏈接：
https://www.reddit.com/r/MachineLearning/comments/rr17f9/p_45_times_faster_hugging_face_transformer/
推薦閱讀
關于程序員大白
程序員大白是一群哈工大，東北大學，西湖大學和上海交通大學的碩士博士運營維護的號，大家樂于分享高質量文章，喜歡總結知識，歡迎關注[程序員大白]，大家一起學習進步！