讓Transformer的推理速度提高4.5倍,這個trick還能給你省十幾萬
本文經AI新媒體量子位(公眾號ID:QbitAI)授權轉載,轉載請聯系出處。
最近,NLP明星公司Hugging Face發布了一個叫做Infinity的產品,可以以1ms延時完成Transformer的推理,性能相當高了。

但是,厲害歸厲害,還是有點貴——1年至少要十幾萬塊 (2萬美元)。
那有沒有什么平替的方法呢?
有的!還是開源的、“不費吹灰之力”就可以達到Infinity一些公共基準的那種。

并且現在,通過在該方法上施加一個小trick,將Transformer的推理速度提高4.5倍!
△ 帖子發布不到一天就收獲了250+熱度
那么,一個“平替”到底為什么能達到“付費”的效果呢?
一個trick讓Transformer推理速度提高4.5倍
先來認識一下這個方法:Transformer-deploy。

它可以用一行命令優化和部署Hugging Face上的Transformer模型,并支持大多數基于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)。
作者表示:
據我所知,目前任何OOS云服務都還沒用到過這個方法。
不過執行GPU量化需要修改模型源代碼(需在矩陣乘法等代價高昂的操作上添加一些叫做QDQ的特定節點),既容易出錯,又很無聊,并且還需自己維護修改后的代碼。
因此作者已經為多個基于Transformer的模型手動完成了這項工作。
后來,他們又發現似乎只需修補模型模塊的抽象語法樹 (AST)也可以自動完成。
在用戶端,在GPU上執行模型的基本量化類似這樣:

最終,該方法在Roberta-base模型和MNLI數據集(分類任務)上實現了4.53倍的推理速度。

當然這也犧牲了0.4個點的精度;如果一點不犧牲的話,也可以加速3.2倍左右。
作者表示,與Transformer-deploy原來的版本相比,這已經是一個很大的改進了,畢竟原版本的加速成本需要超過1個點的精確度。
最終他們用Albert、Bert(包括miniLM)、Distilbert、Roberta(包括 Camembert、XLM-R、DistilRoberta等)、Electra測試了該trick。
結果是對于任何可以導出為ONNX格式的Transformer模型,都可以“開箱即用”。