NVIDIA BERT 推理解決方案 Faster Transformer 開源了

NVIDIA BERT推理解決方案Faster Transformer開源了

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 的調用即可。具體代碼如下所示。

# import op
transformer_op_module = tf.load_op_library(os.path.join('../../build/lib/libtf_transformer.so'))
...
def fast_transformer_model_trans(...)
    ...
      # original code
      ...
      layer_output = layer_norm(layer_output + attention_output)

      # calling faster transformer op
      trainable_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=tf.get_variable_scope().name)
      layer_output = transformer_op_module.bert_transformer(
        layer_input,
        layer_input,
        trainable_vars[0], trainable_vars[2], trainable_vars[4], trainable_vars[1], trainable_vars[3], trainable_vars[5], 
        attention_mask,
        trainable_vars[6], trainable_vars[7], trainable_vars[8], trainable_vars[9], trainable_vars[10], trainable_vars[11],
        trainable_vars[12], trainable_vars[13], trainable_vars[14], trainable_vars[15],
        batch_size=batch_size, from_seq_len=seq_length, to_seq_len=seq_length, head_num=num_attention_heads, size_per_head=attention_head_size)
      
      # original code
      ...
      all_layer_outputs.append(layer_output)
    ...

使用 C++ API 或者 TensorRT 調用 Faster Transformer

考慮到封裝成 TensorFlow OP 會引入一些額外的開銷，我們更建議用戶直接使用 C++ API 或者 TensorRT Plugin 的方式去集成。目前這兩種方式不支持直接解析訓練好的模型。Transformer 層所需要的 weights 參數，需要用戶手動從訓練好的模型中導出。調用方式相對簡單，將導出的 weights 賦值給參數結構體，創建相應的對象，調用 initialize 或者 build_engine 函數初始化對象。運行時，每次調用 forward 或者do_inference 即可。具體代碼如下所示。

/* C++ interface */
typedef BertEncoderTransformerTraits<OperationType::HALF,  cuda::OpenMultiHeadAttention> EncoderTraits_;
fastertransformer::Allocator<AllocatorType::CUDA> allocator(0);
EncoderInitParam<__half> encoder_param; //init param here

encoder_param.from_tensor = d_from_tensor;
...

BertEncoderTransformer<EncoderTraits_> *encoder_transformer_ = new 
 BertEncoderTransformer<EncoderTraits_>(allocator, batch_size, from_seq_len, to_seq_len, head_num, size_per_head);
encoder_transformer_->initialize(encoder_param);
encoder_transformer_->forward();

/* TensorRT Plugin */
std::vector<std::vector<T *> > params;
/* push all weight pointers into the vector params*/
TRT_Transformer<T>* trt_transformer = new TRT_Transformer<T>(batch_size, seq_len, head_num, hidden_dim, layers);
trt_transformer->build_engine(params);
trt_transformer->do_inference(batch_size, h_from_tensor, h_attr_mask, h_transformer_out, stream);

優化原理

在深入了解 Faster Transformer 的優化原理之前，我們先來看下 TensorFlow 的實現情況。圖1是 TensorFlow 在默認計算模式（不使用 XLA 優化）下的時間線片段。

圖1. TensorFlow計算GELU的時間線

其中，黃色矩形框中對應的是激活函數 GELU。可以看到，在 TensorFlow 中，這個函數是通過 8 個類似Pow，Add，和 Tanh 等基本 OP 來實現的。Layer Normalization 操作也是類似的情況（圖2）。

圖2. TensorFlow計算Layer Normalization的時間線

在 TensorFlow 中，每一個基本 OP 都會對應一次 GPU kernel 的調用，和多次顯存讀寫，這些都會增加大量額外的開銷。TensorFlow XLA 可以在一定程度上緩解這個問題，它會對一些基本的 OP 進行合併，以減少GPU kernel 的調度和顯存讀寫。但在大多數情況下，XLA 依然無法達到最優的性能，特別是對於 BERT 這種計算密集的情況，任何性能的提升都將節省巨量的計算資源。

如我們前面提到的，OP 融合可以降低 GPU 調度和顯存讀寫，進而提升性能。出於性能最大化的考慮，在Faster Transformer 內部，我們將除矩陣乘法以外的所有 kernel 都進行了儘可能的融合，單層Transformer 的計算流程如下圖所示：

圖3. BERT中Transformer Layer 的計算流程圖

如圖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/

（本文作者：NVIDIA GPU計算專家團隊，賈曉瑩）

關於NVIDIA

NVIDIA（納斯達克股票代碼：NVDA）在1999年發明的GPU激發了PC遊戲市場的增長，重新定義了現代計算機顯卡，並且對并行計算進行了革新。最近，通過將GPU作為可以感知和理解世界的計算機、機器人乃至自動駕駛汽車的大腦，GPU深度學習再度點燃了全新的計算時代——現代人工智慧。更多信息，請訪問http://nvidianews.nvidia.com/。