有錢任性：英偉達訓練80億參數量GPT-2，1475塊V100 53分鐘訓練BERT

機器之心報道

機器之心編輯部

你見過 50 分鐘訓練 BERT-Large、80 億參數量訓練 GPT-2 嗎？快看看經過 CUDA 優化的 Transformer 為什麼這麼強。

英偉達今日宣布，該公司打破了 NLP 領域的三項記錄：

1）將 BERT 的訓練時間縮短到了 53 分鐘；

2）將 BERT 的推理時間縮短到了 2.2 毫米（10 毫秒已經是業界公認的高水平）；

3）將 GPT-2 的參數量推向 80 億（以前 OpenAI GPT-2 最大為 15 億參數量）。

這些突破可以為現實世界中所有使用 NLP 對話 AI 和 GPU 硬體的用戶帶來很多便利，如降低語音助手的反應延時，使其與人類的交流更加自然。

訓練最快的語言模型

英偉達在 BERT 訓練、推理時間上的突破離不開其 SuperPOD 系統。它由 92 個英偉達 DGX-2H 系統組成，運行在 1472 塊 V100 GPU 上。該系統在 T4 GPU 上管理推斷行為，它的性能甚至比高度優化的 CPU 還要好幾個數量級。

英偉達使用 PyTorch 運行整個 BERT-Large 模型，並採用了自動混合精度方法以加速吞吐量。對於一般的研究者，只要有一個 DGX-2 伺服器（16 塊 V100），我們就能在 3 天內完成 BERT-Large 模型的訓練。如下展示了在不同 GPU 數量下的訓練時長：

目前，英偉達已經開源了 BERT 的訓練代碼以及 TensorRT 優化的 BERT 樣本，地址和下面的預訓練 GPT-2 是一樣的。

訓練最大的語言模型

如果訓練變得更快，那麼這意味著什麼？當然是我們能訓練更大的模型啦。英偉達成功地構建並訓練了最大的語言模型 GPT-2 8B，這一模型包含 83 億參數量，是 BERT-Large 模型的 24 倍、GPT-2 的 5.6 倍。想想我們一個 BERT-Large 都訓練不了，英偉達還訓練「24 個」BERT-Large，這也是很優秀了。

英偉達將這一模型稱為「Megatron」（威震天），還開源了用來訓練這一模型的 pytorch 代碼。

地址：https://github.com/NVIDIA/Megatron-LM

所以這樣超大模型的結構應該是什麼樣的？當然基礎結構還是 Transformer，但是超參的配置確實非常驚人，你見過 72 層、每層隱藏單元都是 3072 的 Treansformer 么。。

這麼巨大模型，當然效果也還是挺好的，83 億參數量的 GPT-2 在驗證困惑度上，下降地非常快，差不多 5 個 Epoch 就能到達非常理想的效果。如下所示為模型大小與驗證困惑度之間的關係。

為了訓練如此龐大的模型，英偉達利用了模型並行，用創建大模型的技術將模型分成幾個部分。利用這一技術創建的模型非常大，單個 GPU 內存不足以應付，所以需要模型並行來分解壓力。

雖然模型並行本來就有一些開銷，且如果模型能裝進單塊 GPU 最好不用模型並行，但是對於 80 億參數量的 GPT-2，模型並行就是必不可少的。如下展示了隨 GPU 數量的增加，其所提供有效計算力的增長差不多接近線性。

為什麼這個 Transformer 能這麼快

前面不論是訓練 BERT-Large 還是巨型 GPT-2，它們的基礎都是 Transformer，如果 Transformer 訓練得不夠快，那麼堆再多的算力也不能完成這兩項挑戰。Faster Transformer 就是支持極速訓練極大模型的基礎，與兩個模型開源的同時，英偉達也開源了 Faster Transformer，讀者可以獲取項目全部源代碼，最新的性能數據以及支持的特性。

開源地址：https://github.com/NVIDIA/DeepLearningExamples/tree/master/FasterTransformer

如下英偉達 GPU 計算專家團隊賈曉瑩將向各位讀者介紹 Faster Transformer 這一制勝武器。

什麼是 Faster Transformer

目前 Transformer 在多種場景下都有非常優秀的表現，但是在推理部署階段，其計算性能卻受到了巨大的挑戰：以 BERT 為原型的多層 Transformer 模型，其性能常常難以滿足在線業務對於低延遲（保證服務質量）和高吞吐（考慮成本）的要求。以 BERT-BASE 為例，超過 90% 的計算時間消耗在 12 層 Transformer 的前向計算上。

因此，一個高效的 Transformer 前向計算方案，既可以為在線業務帶來降本增效的作用，也有利於以 Transformer 結構為核心的各類網路在更多實際工業場景中落地，這也就是 Faster Transformer 誕生的前提。

Faster Transformer 是一個基於 CUDA 和 cuBLAS 的 Transformer Encoder 前向計算實現，其代碼簡潔明了，後續可以通過簡單修改支持多種 Transformer 結構。目前優化集中在編碼器的前向計算（解碼器開發在後續特性規劃中）。底層由 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 插件，並提供了相應的示例，用以支持用戶將其集成到不同的線上應用代碼中。

Faster Transformer 性能怎麼樣

Faster Transformer 在不同的應用場景下都有著突出的表現。英偉達在這裡測試了不同生產環境下 Faster Transformer 前向計算的執行時間以及與 TensorFlow XLA 的性能比較。

英偉達考察了 Faster Transformer 在搜索或者廣告推薦等大 batch size 場景下的加速效果。下面兩張表分別測試了固定句子長度為 32，標準模型（12 head x 64 維）和裁剪模型（4 head x 32 維）在不同 batch size 下，12 層 Transformer 在 V100 上使用了 FP16 計算精度的性能。

表：標準模型不同 Batch Size 下 TensorFlow XLA 和 Faster Transformer 在 V100 上的性能對比

表：裁剪模型不同 Batch Size 下 TensorFlow XLA 和 Faster Transformer 在 V100 上的性能對比

可以看出，在標準模型和裁剪模型上，Faster Transformer 都有很好的加速效果。

Faster Transformer 優化原理

Faster Transformer 提供了 TensorFlow OP，C API 和 TensorRT Plugin 三種介面。在 TensorFlow 中使用 Faster Transformer 最為簡單。只需要先 import .so 文件，然後在代碼段中添加對 Faster Transformer OP 的調用即可。

雖然使用起來非常便捷，但 Faster Transformer 的優化原理是什麼樣的？

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

圖 1：TensorFlow 計算 GELU 的時間線

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

圖：TensorFlow 計算 Layer Normalization 的時間線

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

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

圖：BERT 中 Transformer Layer 的計算流程圖

如上圖所示，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 進行了各種優化以保證它的高效執行。

WAIC 2019 開發者日將於 8 月 31 日在上海世博中心舉辦，包含 1 個主單元、4 個分單元、黑客馬拉松比賽和開發者診所互動區。

屆時，全球頂尖 AI 專家、技術大牛、知名企業代表以及數千名開發者將齊聚上海，圍繞人工智慧前沿理論技術和開發實踐進行分享與解讀。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！