NVIDIA BERT推理解決方案Faster Transformer開源啦
作者：英偉達GPU計算專家團隊，賈曉瑩
Faster Transformer是一個基於CUDA和cuBLAS的Transformer Encoder前向計算實現，其優越的性能將助力於多種BERT的應用場景。
2017年12月Google在論文「Attention is All You Need」[1] 中首次提出了Transformer，將其作為一種通用高效的特徵抽取器。至今，Transformer已經被多種NLP模型採用，比如BERT[2]以及上月發布重刷其記錄的XLNet[3]，這些模型在多項NLP任務中都有突出表現。在NLP之外， TTS，ASR等領域也在逐步採用Transformer。可以預見，Transformer這個簡潔有效的網路結構會像CNN和RNN一樣被廣泛採用。雖然Transformer在多種場景下都有優秀的表現，但是在推理部署階段，其計算性能卻受到了巨大的挑戰：以BERT為原型的多層Transformer模型，其性能常常難以滿足在線業務對於低延遲（保證服務質量）和高吞吐（考慮成本）的要求。以BERT-BASE為例，超過90%的計算時間消耗在12層Transformer的前向計算上。因此，一個高效的Transformer 前向計算方案，既可以為在線業務帶來降本增效的作用，也有利於以Transformer結構為核心的各類網路在更多實際工業場景中落地。本文將介紹NVIDIA GPU計算專家團隊針對Transformer推理提出的性能優化方案：Faster Transformer。
Faster Transformer是一個BERT Transformer 單層前向計算的高效實現，其代碼簡潔明了，後續可以通過簡單修改支持多種Transformer結構。目前優化集中在編碼器（encoder）的前向計算（解碼器decoder開發在後續特性規劃中）。底層由CUDA和cuBLAS實現，支持FP16和FP32兩種計算模式，其中FP16可以充分利用Volta和Turing架構GPU上的Tensor Core計算單元。
Faster Transformer共接收4個輸入參數。首先是attention head的數量以及每個head的維度。這兩個參數是決定Transformer網路結構的關鍵參數。這兩個參數的動態傳入，可以保證Faster Transformer既支持標準的BERT-BASE（12 head x 64維），也支持裁剪過的模型（例如，4 head x 32維），或者其他各式專門定製化的模型。其餘兩個參數是Batch Size 和句子最大長度。出於性能考慮，目前句子最大長度固定為最常用的32，64 和128三種，未來會支持任意長度。Faster Transformer對外提供C++ API，TensorFlow OP 介面，以及TensorRT [4]插件，並提供了相應的示例，用以支持用戶將其集成到不同的線上應用代碼中。
Faster Transformer目前已經開源，可以訪問https://github.com/NVIDIA/DeepLearningExamples/tree/master/FasterTransformer
獲取項目全部源代碼，最新的性能數據以及支持的特性。歡迎大家前往使用，加星和反饋。
性能數據
Faster Transformer在不同的應用場景下都有著突出的表現。我們在這裡測試了不同生產環境下Faster Transformer前向計算的執行時間以及與TensorFlow XLA的性能比較。測試環境如表1所示：
表1. 性能數據測試環境（本地伺服器）
首先針對線上QPS較低的業務（例如問答），我們將batch size設置為1，測試了BERT標準模型在不同的句子長度下，12層Transformer在P4和T4上的性能。由於這種場景下TensorFlow的性能非常依賴於CPU，因此這裡不予列出。
表2. 小batch size情況下Faster Transformer的性能
接著我們來觀察Faster Transformer在搜索或者廣告推薦等大batch size場景下的加速效果。表3和表4分別測試了固定句子長度為32，標準模型（12 head x 64維）和裁剪模型（4 head x 32維）在不同batch size下，12層Transformer 在V100上使用了FP16計算精度的性能。
表3. 標準模型不同Batch Size下TensorFlow XLA和Faster Transformer在V100上的性能對比
表4. 裁剪模型不同Batch Size下TensorFlow XLA和Faster Transformer在V100上的性能對比
可以看出，在標準模型和裁剪模型上，Faster Transformer都有很好的加速效果。
使用方法
Faster Transformer提供了TensorFlow OP ，C++ API和TensorRT Plugin 三種介面。
在TensorFlow中使用Faster Transformer
在TensorFlow中使用Faster Transformer最為簡單。只需要先import .so文件，然後在代碼段中添加對Faster Transformer OP的調用即可。具體代碼如下所示。
使用C++ API或者TensorRT 調用Faster Transformer
考慮到封裝成TensorFlow OP會引入一些額外的開銷，我們更建議用戶直接使用C++ API或者TensorRT Plugin的方式去集成。目前這兩種方式不支持直接解析訓練好的模型。Transformer層所需要的weights參數，需要用戶手動從訓練好的模型中導出。調用方式相對簡單，將導出的weights賦值給參數結構體，創建相應的對象，調用initialize或者build_engine函數初始化對象。運行時，每次調用forward或者do_inference即可。具體代碼如下所示。
優化原理
在深入了解Faster Transformer的優化原理之前，我們先來看下TensorFlow的實現情況。圖1是TensorFlow在默認計算模式（不使用XLA優化）下的時間線片段。
其中，黃色矩形框中對應的是激活函數GELU。可以看到，在TensorFlow中，這個函數是通過8個類似Pow，Add，和Tanh等基本OP來實現的。Layer Normalization 操作也是類似的情況（圖2）。
在TensorFlow中，每一個基本OP都會對應一次GPU kernel的調用，和多次顯存讀寫，這些都會增加大量額外的開銷。TensorFlow XLA可以在一定程度上緩解這個問題，它會對一些基本的OP進行合併，以減少GPU kernel的調度和顯存讀寫。但在大多數情況下，XLA依然無法達到最優的性能，特別是對於BERT這種計算密集的情況，任何性能的提升都將節省巨量的計算資源。
如我們前面提到的，OP融合可以降低GPU調度和顯存讀寫，進而提升性能。出於性能最大化的考慮，在Faster Transformer內部，我們將除矩陣乘法以外的所有kernel都進行了儘可能的融合，單層Transformer的計算流程如下圖所示：
如圖3所示，Faster Transformer只用了14個kernel就完成了原來將近60個kernel的計算邏輯。這其中，8個kernel是通過調用cuBLAS介面計算矩陣乘法（綠色框），其餘6個是自定義kernel （藍色框）。
針對batch size比較小的場景（例如問答，TTS等），簡單的融合後，基本上就可以達到很好的性能。這類場景下，TensorFlow原生實現的最大瓶頸就在於頻繁的kernel launch，融合後大大降低了launch的開銷，因此可以比較輕易地獲得很好的加速效果。
針對大batch的場景，我們需要對矩陣乘法和所有的自定義kernel做精細的調優，才能達到很好的加速效果。我們從矩陣乘法演算法選擇，非矩陣乘法操作的參數配置，SoftMax多版本實現，以及數據結構類型等幾個方面對大batch的情況進行了專門的調優。
首先針對矩陣乘法，在調用cuBLAS的介面時，可以指定性能最優的演算法。特別是針對Volta和Turing架構的GPU，使用Tensor Core進行半精度計算時，當精度滿足需求的情況下，累加器也可以選擇半精度，從而進一步提升性能。
除矩陣乘法以外的6個kernel，大部分都是對矩陣乘的結果進行一些element-wise的操作。輸入矩陣的大小，跟4個參數有關，batch size，句子長度，attention的head數量以及每個head的維度。針對不同的應用場景，參數大小可能極為不同。比如在線問答類的場景，batch size可能為會很小，通常為1。而廣告推薦或者搜索類的場景，batch size通常跟候選集大小有關，一般會是幾百的規模。這樣，輸入矩陣的行數變化範圍可能是幾十到上千。因此，我們需要針對不同的情況，動態的調整kernel launch時的配置參數（grid和block的大小），甚至要針對同一個功能實現多個不同版本的kernel函數，例如，SoftMax的計算就有兩個不同的實現版本。
針對半精度FP16，我們對各個kernel也進行了相應優化。首先，在kernel的實現中，將輸入的half指針轉成half2類型，並使用了half2相關的數學函數。這樣不僅僅可以達到2倍於half的訪存帶寬和計算吞吐，還可以極大地減少指令的發射數量。其次，在SoftMax以及Layer Normalization的操作中，為防止求和溢出，將數據以half2的形式讀入後，會轉成float2類型，來做求和計算。
除上述優化之外，Faster Transformer還優化了前向計算中耗時較高的GELU激活函數，Layer Normalization以及SoftMax等操作。比如利用warp shuffle實現高效的矩陣按行求和操作， 將1/sqrtf計算替換為rsqrtf函數，以及power (x, 3.0) 替換為x * x * x等。總之，我們針對Transformer進行了各種優化以保證它的高效執行。
結論
Faster Transformer是一個開源的高效Transformer實現，相比TensorFlow XLA 可以帶來1.5-2x的提速。Faster Transformer對外提供C++ API， TensorFlow OP，以及TensorRT Plugin三種介面。對每種介面的調用方式，我們提供了完整的示例，方便用戶集成。
Faster Transformer目前已經開源，可以訪問https://github.com/NVIDIA/DeepLearningExamples/tree/master/FasterTransformer
獲取項目全部源代碼，最新的性能數據以及支持的特性。歡迎大家前往使用，加星和反饋。
[1] Vaswani, Ashish, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 「Attention Is All You Need.」 ArXiv:1706.03762 [Cs], June 12, 2017. http://arxiv.org/abs/1706.03762.
[2] Devlin, Jacob, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. 「BERT: Pre-Training of Deep Bidirectional Transformers for Language Understanding.」 ArXiv:1810.04805 [Cs], October 10, 2018. http://arxiv.org/abs/1810.04805.
[3] Yang, Zhilin, Zihang Dai, Yiming Yang, Jaime Carbonell, Ruslan Salakhutdinov, and Quoc V. Le. 「XLNet: Generalized Autoregressive Pretraining for Language Understanding.」 ArXiv:1906.08237 [Cs], June 19, 2019. http://arxiv.org/abs/1906.08237.
[4] TensorRT: https://developer.nvidia.com/tensorrt
[5] Turing T4 GPU, more information: https://www.nvidia.com/en-us/data-center/tesla-t4/
[6] Volta V100 GPU, more information: https://www.nvidia.com/en-us/data-center/tesla-v100/
[7] Pascal P4 GPU, more information: https://www.nvidia.com/en-us/deep-learning-ai/solutions/inference-platform/hpc/
推薦閱讀：
