微信也在用的Transformer加速推理工具,現(xiàn)在騰訊開源了
本文經(jīng)AI新媒體量子位(公眾號ID:QbitAI)授權(quán)轉(zhuǎn)載,轉(zhuǎn)載請聯(lián)系出處。
近年來,基于Transformer的模型,可以說是在NLP界殺出了一片天地。
雖然在提高模型模型精度上,Transformer發(fā)揮了不容小覷的作用,但與此同時(shí),卻引入了更大的計(jì)算量。
那么,這個(gè)計(jì)算量有多大呢?
來看下數(shù)據(jù)。

因此,實(shí)現(xiàn)一個(gè)能充分發(fā)揮CPU/GPU硬件計(jì)算能力的Transformer推理方法,就成了急需解決的問題。
近日,騰訊便開源了一個(gè)叫TurboTransformers的工具,對Transformer推理過程起到了加速作用,讓你的推理引擎變得更加強(qiáng)大。

這個(gè)工具已經(jīng)在微信、騰訊云、QQ看點(diǎn)等產(chǎn)品中廣泛應(yīng)用,在線上預(yù)測場景中可以說是“身經(jīng)百戰(zhàn)”。
Turbo具有如下三大特性:
- 優(yōu)異的CPU/GPU性能表現(xiàn)。
- 為NLP推理任務(wù)特點(diǎn)量身定制。
- 簡單的使用方式。
值得一提的是,TurboTransformers,是騰訊通過Github對外開源的第100個(gè)項(xiàng)目。
那么,具有如此“紀(jì)念意義”的開源工具,到底有多厲害?
接下來,我們將一一講解。
多項(xiàng)性能測試“摘桂冠”
Turbo在CPU/GPU性能上的表現(xiàn)可以說是非常優(yōu)異。
在多種CPU和GPU硬件上獲得了超過pytorch/tensorflow和目前主流優(yōu)化引擎的性能表現(xiàn)。

CPU上的測試結(jié)果
首先,是在CPU 硬件平臺上,測試了 TurboTransformers 的性能表現(xiàn)。
選擇 pytorch、pytorch-jit 和 onnxruntime-mkldnn 和 TensorRT 實(shí)現(xiàn)作為對比。
性能測試結(jié)果為迭代 150 次的均值。為了避免多次測試時(shí),上次迭代的數(shù)據(jù)在 cache 中緩存的現(xiàn)象,每次測試采用隨機(jī)數(shù)據(jù),并在計(jì)算后刷新的 cache 數(shù)據(jù)。
下圖是Intel Xeon 6133 CPU的性能測試結(jié)果。


GPU上的測試結(jié)果
其次,是在GPU硬件平臺上,測試了 TurboTransformers 的性能表現(xiàn)。
選擇對比的對象分別是:pytorch、NVIDIA Faster Transformers、onnxruntime-gpuTensorRT。
性能測試結(jié)果為迭代 150 次的均值。
下圖是在NVIDIA RTX 2060 GPU的性能測試結(jié)果。


接下來,是在NVIDIA P40 GPU的性能測試結(jié)果。


最后,是在NVIDIA V100 GPU的性能測試結(jié)果。


Turbo技術(shù)原理
能夠取得如此好的推理性能,這背后的計(jì)算原理又是什么呢?
TurboTransformers的軟件架構(gòu)如下圖,它讓微信內(nèi)部眾多NLP線上應(yīng)用能夠充分榨取底層硬件的計(jì)算能力,讓算法更好地服務(wù)的用戶。
具體來說TurboTransformers可以在算子優(yōu)化、框架優(yōu)化和接口部署方式簡化三個(gè)方面做了工作。

算子層優(yōu)化
Transformer都包含了什么計(jì)算呢?
如下圖所示,圖(a)展示了論文Transformer結(jié)構(gòu)示意圖,這里稱灰色方框內(nèi)的結(jié)構(gòu)為一個(gè)Transformer Cell,BERT encoder堆疊了Nx個(gè)這樣的Transformer Cell。
圖(b)將一個(gè)Cell的細(xì)節(jié)加以展開,每一個(gè)矩形都是一個(gè)獨(dú)立的計(jì)算核心。

Transformer Cell計(jì)算包含了8個(gè)GEMM(通用矩陣乘法,General Matrix Multiplication)運(yùn)算。通過調(diào)優(yōu)Intel MKL和cuBLAS的GEMM調(diào)用方式來獲得最佳GEMM性能。
并且在硬件允許條件下,在GPU上使用tensor core方式進(jìn)行GEMM運(yùn)算。
類似NVIDIA FasterTransformers方案,將所有GEMM運(yùn)算之間的計(jì)算融合成一個(gè)調(diào)用核心。融合會(huì)帶來兩個(gè)好處,一是減少了內(nèi)存訪問開銷,二是減少多線程啟動(dòng)開銷。
對于這些核心,在CPU上采用openmp進(jìn)行并行,在GPU上使用CUDA進(jìn)行優(yōu)化實(shí)現(xiàn)。
對于比較復(fù)雜的LayerNorm和Softmax算子,它們包含了不適合GPU上并行的規(guī)約操作,TurboTransformers為它們設(shè)計(jì)了創(chuàng)新并行算法,極大降低了這些算子的延遲。
理論上Transformers推理延遲應(yīng)該近似于矩陣乘法延遲。
框架層優(yōu)化
TurboTransformers采用了一個(gè)有效的內(nèi)存管理方式。
由于NLP的采用變長輸入特性,每次運(yùn)算中間結(jié)果的大小其實(shí)并不相同。為了避免每次都分配釋放內(nèi)存,研究人員通過Caching方式管理顯存。
為了能夠無縫支持pytorch/tensorflow訓(xùn)練好的序列化模型,提供了一些腳本可以將二者的預(yù)訓(xùn)練模型轉(zhuǎn)化為npz格式,供TurboTransformers讀入。
特別的,考慮到pytorch huggingface/transformers是目前最流行的transformers訓(xùn)練方法,支持直接讀入huggingface/transformers預(yù)訓(xùn)練模型。
應(yīng)用部署
Turbo提供了C++和Python調(diào)用接口,可以嵌入到C++多線程后臺服務(wù)流程中,也可以加入到pytorch服務(wù)流程中。
研究人員建議TurboTransformers通過docker部署,一方面保證了編譯的可移植性,另一方面也可以無縫應(yīng)用于K8S等線上部署平臺。
傳送門
GitHub項(xiàng)目地址:https://github.com/Tencent/TurboTransformers/blob/master/README_cn.md