支援Transformer全流程訓練加速，加速3倍！位元組跳動LightSeq上新

Transformer 已經成為眾多 NLP 任務以及部分 CV 任務的主流模型，但由於硬體資源匱乏，很多高校實驗室或者公司都無法訓練很大的模型，而降低批處理大小等措施又會導致訓練時間成倍增加。針對這一痛點，位元組跳動推出了 LightSeq 訓練加速引擎，對 Transformer 訓練的整個計算過程進行了最佳化，最多可以實現 3 倍以上的加速。

如今，NLP 和 CV 領域的大部分任務都將 Transformer 作為基礎模型。而早在 2019 年 12 月，位元組跳動就開源過一款 Transformer類模型推理加速引擎——LightSeq [1]。作為業界第一款支援多種模型和解碼方法的推理加速引擎，LightSeq 的推理速度快於其它同類軟體，更是遠遠超過了 TensorFlow 和 PyTorch。

今天，LightSeq 的最新版本釋出了！

這次釋出引入了引擎方面的重大更新—— 支援了 Transformer 全流程訓練加速，在不同的批處理大小下相比主流訓練庫最高可加速 3 倍多！至此從訓練到推理部署的整個過程都已被 LightSeq 打通。

那麼它到底採用了哪些技術呢？筆者將根據 LightSeq 公佈的資料為你逐個揭曉。

LightSeq 能做什麼？

圖 1：Transformer 模型結構圖（以機器翻譯為例）

自 2017 年被谷歌提出之後，Transformer 模型 [2] 成為了眾多 NLP 任務以及部分 CV 任務的主流模型，尤其是機器翻譯、文字生成、文字摘要、時間序列預測等任務。圖 1 是機器翻譯任務使用 Transformer 進行訓練的一個例子。

但由於硬體資源匱乏，很多高校實驗室或者公司都無法訓練很大的模型，而降低批處理大小等措施又會導致訓練時間成倍增加。因此，如何利用有限的資源儘快訓練出模型成為了亟待解決的問題。

針對這一痛點，位元組跳動推出了 LightSeq 訓練加速引擎，對 Transformer 訓練的整個計算過程進行了最佳化。從詞嵌入層、編碼層、解碼層，到最後的損失函式層，從前向傳播、反向傳播、梯度同步，到最後的引數更新，LightSeq 都進行了細緻的效能分析和最佳化。

以機器翻譯任務為例，在該任務中，我們僅需要修改幾行程式碼開啟 LightSeq 訓練加速，就可以實現最多 3 倍以上的加速。

總的來說，LightSeq 具有如下幾個 優點：

1. 支援Transformer 的完整加速。

LightSeq 是業界第一款完整支援整個 Transformer 模型加速的訓練引擎，包括了詞嵌入層、編碼層、解碼層、損失函式層等高效自定義層。相比之下，另一款知名的深度學習最佳化引擎 DeepSpeed [3] 僅支援編碼層的加速，因此只能用在 BERT 等模型上，侷限性較大。

2. 訓練速度快。

LightSeq 訓練速度非常快。例如在 WMT14 [4] 英德機器翻譯任務上，利用英偉達最新的訓練顯示卡 A100，相比於主流序列生成庫，LightSeq 最快僅需要三分之一的訓練時間。

3. 功能全面，簡單易用。

LightSeq 提供了高效的 TensorFlow 和 PyTorch 自定義層供使用者靈活使用，可以自由插入到 Hugging Face 等主流訓練庫中。除此之外，它還和當前流行的訓練庫如 Fairseq [5]、NeurST [6] 等做了深度整合。使用者在安裝 LightSeq 後，只需要修改幾個命令列引數，就能在這些訓練庫上使用 LightSeq。

4. 提供豐富的二次開發工具。

LightSeq 提供了完整的 CUDA kernel 和 Transformer 自定義層的單元測試功能，可以測試自定義運算元的正確性，同時分析出執行時間和加速比，幫助開發者更快地驗證功能的正確性和有效性。

下表 1 列舉了不同訓練加速引擎之間支援的功能對比：

表 1：不同訓練加速引擎支援功能對比

如何快速上手 LightSeq？

LightSeq 提供了多種便捷的執行方式，點選文末連結可以快速體驗。這裡簡單講解一下快速接入 LightSeq 進行訓練的兩種方式。

使用自定義層

LightSeq 提供了許多自定義層供使用者靈活使用。

例如，如果你想將 Transformer 模型中的編碼層替換為 LightSeq 的編碼層，只需要提供一個編碼層引數，用來初始化 LightSeq 編碼層，然後就可以用它來替換原始的編碼層，加速模型訓練。詳細程式碼如下（這裡省略了部分配置引數）：

一鍵啟動

LightSeq 還和幾個當前流行的訓練庫（例如 Fairseq 和 NeurST）做了深度整合，僅需修改很少的程式碼就能開啟 LightSeq 加速。

Fairseq

LightSeq 為 Fairseq 提供了一套完整便捷的 Transformer 訓練樣例。

首先，你需要安裝 Fairseq 以及必要的第三方庫，然後用如下命令安裝 LightSeq 加速庫：

接著就可以透過 LightSeq 提供的啟動器，靈活地指定使用 LightSeq 最佳化版本的 Transformer 模型、引數最佳化器和損失函式。下面是啟動引數樣例（這裡省略了和 LightSeq 無關的引數）：

NeurST

NeurST 是一款同時支援 TensorFlow 和 PyTorch 的開源序列生成庫，可以用來做文字生成、機器翻譯和語音翻譯。LightSeq 已經與其進行了深度融合，無需使用者修改程式碼和啟動引數即可直接進行加速訓練。

首先需要安裝 NeurST，官方提供了詳細的安裝教程 [6]。

然後安裝 TensorFlow 版本的 LightSeq，命令如下：

這樣，NeurST 就會自動識別 LightSeq 已經安裝成功，呼叫 lightseq 庫進行模型構建來加速訓練，無需修改啟動引數。執行命令詳見 NeurST 提供的機器翻譯樣例 [7]。

效能測試

在 WMT14 標準的英德翻譯任務上，LightSeq 做了評測實驗。以當前流行的 Fairseq 訓練庫（基於 PyTorch）和被廣泛使用 Apex 工具庫 [8] 為基準，測試了 LightSeq 的訓練效能。實驗在 NVIDIA Tesla V100 和 NVIDIA Ampere A100 上進行，採用單機八卡資料並行訓練和 16 位浮點數混合精度。

在不同模型大小和批處理大小下，LightSeq 對單步訓練速度的提升結果如圖 2 所示：

圖 2：A100 和 V100 顯示卡下不同層數模型加速比

這裡所有模型的詞表大小為 40k，編碼層和解碼層隱層維度是 1024，注意力頭數為 16。e 和 d 分別表示編碼器和解碼器的層數。加速比的計算採用了每秒訓練有效單詞數（real word per second）這一常見指標。

可以發現：

1. 使用了 LightSeq 後，單步訓練速度有 45%-250% 的提升。作為對比，Apex 僅有 5%-30% 的提升。LightSeq 能取得這種大幅度的效能提升，源自其對模型進行了全流程的細緻最佳化。

2. Apex 視訊記憶體使用量對比原生 Fairseq 略有提升，例如 V100 上，Apex 導致 6e6d 的模型在 15000 批處理大小上出現了視訊記憶體溢位，而 LightSeq 仍然可以正常訓練。這證明了 LightSeq 在保證高效計算的同時，也做到了高效使用視訊記憶體，這個特徵在視訊記憶體有限或者訓練大模型時非常關鍵。

3. 隨著批處理大小的增加，LightSeq 加速比逐漸降低。其原因是，經過 LightSeq 最佳化後，單步訓練中矩陣乘法佔比提高，顯示卡的計算吞吐成為訓練速度的瓶頸。這表明 LightSeq 已經對計算資源做到了充分利用。也解釋了為什麼計算吞吐更高的 A100，平均會取得比 V100 高 15% 左右的加速比。

最後在 WMT14 英德翻譯資料集上測試了圖 1 中 Transformer 模型訓練至收斂的時間，結果如圖 3 所示：

圖 3：A100 和 V100 顯示卡下不同模型訓練總時間

由於 LightSeq 的計算最佳化是無損的，不會影響模型訓練至收斂的訓練步數，所以收斂時間的提升和單步訓練時間的提升趨於一致。觀察可以發現，LightSeq 最多可將模型訓練時間由 8.5 小時降低到 3.8 小時。

視覺化分析

為了更清楚地看出 LightSeq 最佳化前後模型的運算情況，我們用 Nsight Systems [9] 視覺化模型訓練過程中單步的運算元呼叫情況來展示。

首先是 Fairseq+Apex 的視覺化，結果如圖 4 所示。總耗時在 288ms 左右，三個紅色框分別表示前向傳播、反向傳播、梯度同步與引數更新。可以看出前向傳播的運算元排列比較稀疏，存在很大的最佳化空間。

圖 4：Fairseq+Apex 單步訓練過程視覺化

然後是 Fairseq+LightSeq 的視覺化，結果如圖 5 所示，總耗時降到了 185ms 左右。而且 LightSeq 的運算元排列更加緊密，大大增加了顯示卡的利用率。

圖 5：Fairseq+LightSeq 單步訓練過程視覺化

CUDA kernel效能

此外還測試了 LightSeq 單卡情況下所有 CUDA kernel 的效能，對比了 PyTorch、TensorFlow（XLA 編譯最佳化）、DeepSpeed 和 LightSeq 四種實現方式。由於 kernel 太多，這裡只列舉了部分實驗結果。

首先對比了最常用的 dropout，圖 6 是 V100 顯示卡上 16 位和 32 位浮點數 dropout 不同實現的加速對比結果：

圖 6：dropout 加速對比

從圖 6 中可以看出，LightSeq 的實現要遠遠快於 PyTorch 和 DeepSpeed。DeepSpeed 在元素數量過百萬之後逐漸不如 PyTorch，而 LightSeq 始終比 PyTorch 快。隨著元素數量的增加，LightSeq 和 DeepSpeed 的速度都會有明顯下降。TensorFlow 在開啟了 XLA 之後速度依然落後於 PyTorch，且隨著元素數量的增加差距逐漸縮小。

然後對比了注意力機制中的 softmax 函式，測試了實際訓練場景中批處理大小為 8192 情況下的加速比。圖 7 是 V100 顯示卡上 16 位和 32 位浮點數 softmax 不同實現的加速對比結果，因為 DeepSpeed 只支援句子長度為 8 的整數倍，所以這裡只測試了長度為 32 的整數倍的句子計算速度：

圖 7：softmax 加速對比

可以看出，LightSeq 幾乎在所有情況下都遠遠快於 DeepSpeed。且隨著序列長度增加，LightSeq 加速比逐漸增大。而 DeepSpeed 在 16 位浮點數情況下加速比會逐漸減小，在 32 位浮點數情況下甚至會出現比 PyTorch 還要慢的情況。TensorFlow 即便使用 XLA 融合運算元，運算速度也遠遠落後於其它三種實現。

在其它多數 kernel 測試中，LightSeq 都要快於其它三種實現。

關鍵技術

圖 8：模型訓練過程

如圖 8 所示，以 2 卡為例，一個完整的 Transformer 模型訓練過程主要包括四個步驟：前向傳播、後向傳播、梯度同步和引數更新。其中前向傳播和後向傳播佔模型訓練總時間的 70% 多，包含了眾多計算密集型和 I/O 密集型操作，因此是最佳化的重點。而梯度同步雖然耗時相對較少，但是也可以透過和反向傳播並行化等方法隱藏掉大部分延時。最後最佳化器更新引數方面也大有文章可做，可以從計算和 I/O 兩個方面降低延時。

下面詳細介紹最佳化這四個步驟的幾種技術： 運算元多運算融合、輸入輸出層融合、動態視訊記憶體複用和引數連續化，最後介紹一下 LightSeq 提供的 單元測試功能。

運算元多運算融合

在大多數深度學習框架（例如 TensorFlow 和 PyTorch）中，一個簡單的運算通常都需要很多細粒度的核函式來實現。例如在 TensorFlow 中，一次層歸一化（Layer Normalization）操作需要呼叫三次核函式以及兩次中間視訊記憶體讀寫，非常耗時。而基於 CUDA，LightSeq 定製化了一個層歸一化專用的核函式，將兩次中間結果的寫入暫存器。從而實現一次核函式呼叫，同時沒有中間結果視訊記憶體讀寫，因此大大節省了計算開銷。

基於這個思路，LightSeq 利用 CUDA 矩陣運算庫 cuBLAS [10] 提供的矩陣乘法和自定義核函式實現了 Transformer 的編碼器和解碼器。以編碼層為例，具體結構如圖 9 所示：

圖 9：編碼層計算過程

藍色部分是自定義核函式，黃色部分是矩陣乘法。可以發現，矩陣乘法之間的運算全部都用一個定製化核函式實現了，因此大大減少了核函式呼叫和視訊記憶體讀寫，最終提升了運算速度。

LightSeq 還最佳化了核函式的實現，採用 float4 資料型別來進行資料讀寫，大大增加了資料的吞吐量，減小了核函式執行的延時。LightSeq 支援任意長度的序列輸入，根據不同的序列長度選擇最合適的核函式來呼叫執行。相比之下，DeepSpeed 多數核函式只支援句子長度為 8 的整數倍，因此速度較慢，侷限性較大。

輸入輸出層融合

此外 LightSeq 還實現了詞嵌入層和損失函式層的運算元融合。對於詞嵌入層，LightSeq 將詞表查詢與放縮、位置向量融合以及 dropout 操作都寫成了一個核函式。對於損失函式層，將交叉熵損失融合成一個核函式。透過輸入輸出層的融合，進一步減小了模型訓練的時間，增加了顯示卡利用率。

以詞嵌入層為例，圖 10 展示了輸入句子中單詞 id 之後，詞嵌入層進行的計算過程：

圖 10：詞嵌入層計算過程

可以看出，在融合之前，一個詞嵌入層需要經過詞向量查詢與放縮、位置向量查詢、兩者相加、dropout 五種運算，因此需要頻繁呼叫核函式，非常耗時。而將這五個操作融合成一個核函式可以大大加快獲取最終詞表示的速度。

動態視訊記憶體複用

為了避免計算過程中的視訊記憶體申請釋放並節省視訊記憶體佔用，LightSeq 首先對模型中所有動態的矩陣大小都定義了最大值（例如最大序列長度）。接著在模型初始化的時候，為計算過程中的每個中間計算結果按最大值分配視訊記憶體，並對沒有依賴的中間結果共用視訊記憶體。

引數連續化

LightSeq 將 Transformer 每一層中所有的引數都繫結在一起，做連續化處理。初始化模型某一層的時候只需要定義一個引數，引數量為原始模型該層總的引數量。在後續計算時，只需要從這個引數的對應位置處取出原始引數值即可。

以編碼層為例，自注意力層和前饋層共有 16 個引數，假設總引數量為 S。於是可以定義一個大小為 S 的引數，按照原始各個引數的引數量來劃分出 16 個塊，連續儲存 16 個引數，在計算時只需要取出對應的引數塊即可。

引數連續化能顯著降低引數更新期間複製、同步、計算的次數。實驗分析發現，進行了這項最佳化後，最佳化器效能提升了 40%-50%。

開發工具

LightSeq 提供了豐富的單元測試功能，不僅可以測試所有的自定義核函式和自定義層的正確性，還可以對比測試不同實現之間的速度。使用者可以自由指定測試的組數、每組重複執行次數、容差和資料型別。

具體步驟上，首先使用者需要構造隨機資料，然後根據資料型別選擇不同的 CUDA 核函式，最後分別實現自定義和基準的計算函式即可。這樣就保證了使用者可以基於自身場景進行更進一步的二次定製開發。

總結

LightSeq新版訓練加速引擎全方面提升了Transformer模型的訓練速度，打通了訓練和推理部署全流程，使用靈活方便。大大縮減了科研工作者們訓練模型的成本。可以期待未來像機器翻譯、文字生成、摘要、對話生成、情感分析等大量NLP應用場景可以使用LightSeq來訓練和推理。

LightSeq 地址：

參考文獻：

[1] LightSeq:

[2] Transformer: "Attention is all you need", NIPS (2017)

[3] DeepSpeed:

[4] WMT14:

[5] Fairseq:

[6] NeurST:

[7] NeurST 機器翻譯樣例:
/examples/translation

[8] Apex:

[9] Nsight Systems: https://developer.nvidia.com/nsight-systems

[10] cuBLAS: https://docs.nvidia.com/cuda/cublas/index.html