這應該是業界第一款完整支援 Transformer、GPT 等多種模型高速推理的開源引擎。
2017 年 Google 提出了 Transformer [1] 模型，之後在它基礎上誕生了許多優秀的預訓練語言模型和機器翻譯模型，如 BERT [2] 、GPT 系列[13]等，不斷重新整理著眾多自然語言處理任務的能力水平。與此同時，這些模型的引數量也在呈現近乎指數增長（如下圖所示）。例如最近引發熱烈討論的 GPT-3 [3]，擁有 1750 億引數，再次重新整理了引數量的記錄。
如此巨大的引數量，也為模型推理部署帶來了挑戰。以機器翻譯為例，目前 WMT[4]比賽中 SOTA 模型已經達到了 50 層以上。主流深度學習框架下，翻譯一句話需要好幾秒。這帶來了兩個問題：一是翻譯時間太長，影響產品使用者體驗；二是單卡 QPS （每秒查詢率）太低，導致服務成本過高。
因此，今天給大家安利一款速度非常快，同時支援非常多特性的高效能序列推理引擎——LightSeq。它對以 Transformer 為基礎的序列特徵提取器（Encoder）和自迴歸的序列解碼器（Decoder）做了深度優化，早在 2019 年 12 月就已經開源，應用在了包括火山翻譯等眾多業務和場景。據瞭解，這應該是業界第一款完整支援 Transformer、GPT 等多種模型高速推理的開源引擎。
LightSeq 可以應用於機器翻譯、自動問答、智慧寫作、對話回覆生成等眾多文字生成場景，大大提高線上模型推理速度，改善使用者的使用體驗，降低企業的運營服務成本。
相比於目前其他開源序列推理引擎，LightSeq具有如下幾點優勢：
- 高效能
LightSeq推理速度非常快。例如在翻譯任務上，LightSeq相比於Tensorflow實現最多可以達到14倍的加速。同時領先目前其他開源序列推理引擎，例如最多可比Faster Transformer快1.4倍。
- 支援模型功能多
LightSeq支援BERT、GPT、Transformer、VAE 等眾多模型，同時支援beam search、diverse beam search[5]、sampling等多種解碼方式。下表詳細列舉了Faster Transformer[7]、Turbo Transformers[6]和LightSeq三種推理引擎在文字生成場景的功能差異：
- 簡單易用，無縫銜接Tensorflow、PyTorch等深度學習框架
LightSeq通過定義模型協議，支援各種深度學習框架訓練好的模型靈活匯入。同時包含了開箱即用的端到端模型服務，即在不需要寫一行程式碼的情況下部署高速模型推理，同時也靈活支援多層次複用。
使用方法
利用 LightSeq 部署線上服務比較簡便。LightSeq 支援了 Triton Inference Server[8]，這是 Nvidia 開源的一款 GPU 推理 server ，包含眾多實用的服務中介軟體。LightSeq 支援了該 server 的自定義推理引擎 API 。因此只要將訓練好的模型匯出到 LightSeq 定義的模型協議[9]中，就可以在不寫程式碼的情況下，一鍵啟動端到端的高效模型服務。更改模型配置（例如層數和 embedding 大小）都可以方便支援。具體過程如下：
首先準備好模型倉庫，下面是目錄結構示例，其中 transformer.pb 是按模型協議匯出的模型權重，libtransformer.so 是 LightSeq 的編譯產物。
- model_zoo/
- model_repo/
- config.pbtxt
- transformer.pb
- 1/
- libtransformer.so然後就可以啟動Triton Inference Server[8]，搭建起模型服務。
1. trtserver --model-store=${model_zoo}效能測試
在 NVIDIA Tesla P4 和 NVIDIA Tesla T4 顯示卡上，筆者測試了 LightSeq 的效能，選擇了深度學習框架 Tensorflow v1.13 和解碼場景支援較為豐富的 Faster Transformer v2.1 實現作為對比。Turbo Transformers 解碼方法比較單一（只支援 Beam Search ，不支援文字生成中常用的取樣解碼），尚未滿足實際應用需求，因此未作對比。
機器翻譯效能
在機器翻譯場景下，筆者測試了 Transformer base 模型（6層 encoder、6層 decoder 、隱層維度 512 ）採用 beam search 解碼的效能，實驗結果如下：
可以發現，在小 batch 場景下，Faster Transformer 和 LightSeq 對比 Tensorflow 都達到了 10 倍左右的加速。而隨著 batch 的增大，由於矩陣乘法運算佔比越來越高，兩者對 Tensorflow 的加速比都呈衰減趨勢。LightSeq 衰減相對平緩，特別是在大 batch 場景下更加具有優勢，最多能比 Faster Transformer 快 1.4 倍。這也對未來的一些推理優化工作提供了指導：小 batch 場景下，只要做好非計算密集型運算元融合，就可以取得很高的加速收益；而大 batch 場景下則需要繼續優化計算密集型運算元，例如矩陣乘法等。
最後在 WMT14 標準的法英翻譯任務上，筆者測試了 Transformer big 模型的效能。LightSeq 在 Tesla P4 顯示卡上平均每句翻譯延遲為 167ms ，Tesla T4 上減小到了 82ms。而作為對比， TensorFlow 延遲均為 1071ms，LightSeq 分別達到了 6.41 和 13.06 倍加速。另外，筆者嘗試了其他多種模型配置，得到了比較一致的加速效率。例如更深層的模型結構上（encoder加深至 16 層），LightSeq 得到的加速比，分別是 6.97 和 13.85 倍。
文字生成效能
上述機器翻譯通常採用 Beam Search 方法來解碼， 而在文字生成場景，經常需要使用取樣（ Sampling ）來提升生成結果的多樣性。下圖展示了 Transformer base 模型採用 top-k/top-p sampling 的效能測試對比：
可以發現，在需要使用取樣解碼的任務中，LightSeq 在大部分配置下領先於 Faster Transformer，最多也能達到 1.4 倍的額外加速。此外，相比於 TensorFlow 實現，LightSeq 對 GPT 和 VAE 等生成模型也達到了 5 倍以上的加速效果。
服務壓力測試
在雲服務上，筆者測試了在實際應用中 GPT 場景下，模型服務從 Tensorflow 切換到LightSeq 的延遲變化情況（服務顯示卡使用 NVIDIA Tesla P4）。可以觀察到，pct99 延遲降低了 3 到 5 倍，峰值從 360 毫秒左右下降到 80 毫秒左右，詳細結果如下圖所示：
更多的對比實驗結果可以在 LightSeq 效能評測報告 [10] 中檢視到。
技術原理
以 Transformer 為例，一個機器翻譯/文字生成模型推理過程包括兩部分：序列編碼模組特徵計算和自迴歸的解碼演算法。其中特徵計算部分以自注意力機制及特徵變換為核心（矩陣乘法，計算密集型），並伴隨大量 Elementwise（如 Reshape）和 Reduce（如Layer Normalization）等 IO 密集型運算；解碼演算法部分包含了詞表 Softmax、beam 篩選、快取重新整理等過程，運算瑣碎，並引入了更復雜的動態 shape。這為模型推理帶來了眾多挑戰：
- IO 密集型計算的細粒度核函式呼叫帶來大量冗餘視訊記憶體讀寫，成為特徵計算效能瓶頸。
- 複雜動態 shape 為計算圖優化帶來挑戰，導致模型推理期間大量視訊記憶體動態申請，耗時較高。
- 解碼生成每一步字元過程邏輯複雜，難以並行化計算從而發揮硬體優勢。
LightSeq 取得這麼好的推理加速效果，對這些挑戰做了哪些針對性的優化呢？筆者分析發現，核心技術包括這幾項：融合了多個運算操作來減少 IO 開銷、複用視訊記憶體來避免動態申請、解碼演算法進行層級式改寫來提升推理速度。下面詳細介紹下各部分的優化挑戰和 LightSeq 的解決方法。
運算元多運算融合
近年來，由於其高效的特徵提取能力，Transformer encoder/decoder 結構被廣泛應用於各種 NLP 任務中，例如海量無標註文字的預訓練。而多數深度學習框架（例如 Tensorflow、Pytorch 等）通常都是呼叫基礎運算庫中的核函式（kernel function）來實現 encoder/decoder 計算過程。這些核函式往往粒度較細，通常一個元件需要呼叫多個核函式來實現。
以層歸一化（Layer Normalization）為例，Tensorflow 是這樣實現的：
mean = tf.reduce_mean(x, axis=[-1], keepdims=True)
variance = tf.reduce_mean(tf.square(x - mean), axis=[-1], keepdims=True)
result = (x - mean) * tf.rsqrt(variance + epsilon) * scale + bias可以發現，即使基於編譯優化技術（自動融合廣播（Broadcast）操作和按元素（Elementwise）運算），也依然需要進行三次核函式呼叫（兩次 reduce_mean，一次計算最終結果）和兩次中間結果的視訊記憶體讀寫（mean 和 variance）。而基於 CUDA，我們可以定製化一個層歸一化專用的核函式，將兩次中間結果的寫入暫存器。從而實現一次核函式呼叫，同時沒有中間結果視訊記憶體讀寫，因此大大節省了計算開銷。有興趣的同學可以在文末參考連結中進一步檢視具體實現[11]。
基於這個思路，LightSeq 利用 CUDA 矩陣運算庫 cuBLAS[12]提供的矩陣乘法和自定義核函式實現了 Transformer，具體結構如下圖所示：
藍色部分是自定義核函式，黃色部分是矩陣乘法。可以發現，矩陣乘法之間的運算全部都用一個定製化核函式實現了，因此大大減少了核函式呼叫和視訊記憶體讀寫，最終提升了運算速度。
動態視訊記憶體複用
為了避免計算過程中的視訊記憶體申請釋放並節省視訊記憶體佔用，LightSeq 首先對模型中所有動態的 shape 都定義了最大值（例如最大序列長度），將所有動態shape轉換為靜態。接著在服務啟動的時候，為計算過程中的每個中間計算結果按最大值分配視訊記憶體，並對沒有依賴的中間結果共用視訊記憶體。這樣對每個請求，模型推理時不再申請視訊記憶體，做到了：不同請求的相同 Tensor 複用視訊記憶體；同請求的不同 Tensor 按 shape 及依賴關係複用視訊記憶體。
通過該視訊記憶體複用策略，在一張 T4 顯示卡上，LightSeq 可以同時部署多達 8 個 Transformer big 模型（batch_size=8，最大序列長度=8，beam_size=4，vocab_size=3萬）。從而在低頻或錯峰等場景下，大大提升顯示卡利用率。
層級式解碼計算
在自迴歸序列生成場景中，最複雜且耗時的部分就是解碼。LightSeq 目前已經支援了 beam search、diversity beam search、top-k/top-p sampling 等多種解碼方法，並且可以配合 Transformer、GPT使用，達到數倍加速。這裡我們以應用最多的 beam search 為例，介紹一下 LightSeq 對解碼過程的優化。
首先來看下在深度學習框架中傳統是如何進行一步解碼計算的：
1.計算以每個token為結尾的序列的log probability
log_token_prob = tf.nn.log_softmax(logit) # [batch_size, beam_size, vocab_size]
log_seq_prob += log_token_prob # [batch_size, beam_size, vocab_size]
log_seq_prob = tf.reshape(log_seq_prob, [-1, beam_size * vocab_size])2. 為每個序列（batch element）找出排名topk的token
topk_log_probs, topk_indices = tf.nn.top_k(log_seq_prob, k=K)3. 根據beam id，重新整理decoder中的self attention模組中的key和value的快取
refresh_cache(cache, topk_indices)可以發現，為了挑選概率 top-k 的 token ，必須在 [batch_size, beam_size, vocab_size]大小的 logit 矩陣上進行 softmax 計算及視訊記憶體讀寫，然後進行 batch_size 次排序。通常 vocab_size 都是在幾萬規模，因此計算量非常龐大，而且這僅僅只是一步解碼的計算消耗。因此實踐中也可以發現，解碼模組在自迴歸序列生成任務中，累計延遲佔比很高（超過 30%）。
LightSeq 的創新點在於結合 GPU 計算特性，借鑑搜尋推薦中常用的粗選-精排的兩段式策略，將解碼計算改寫成層級式，設計了一個 logit 粗選核函式，成功避免了 softmax 的計算及對十幾萬元素的排序。該粗選核函式遍歷 logit 矩陣兩次：
• 第一次遍歷，對每個 beam，將其 logit 值隨機分成k組，每組求最大值，然後對這k個最大值求一個最小值，作為一個近似的top-k值（一定小於等於真實top-k值），記為R-top-k。在遍歷過程中，同時可以計算該beam中logit的log_sum_exp值。
• 第二次遍歷，對每個 beam，找出所有大於等於 R-top-k 的 logit 值，將（logit - log_sum_exp + batch_id offset, beam_id vocab_size + vocab_id）寫入候選佇列，其中 offset 是 logit 的下界。
在第一次遍歷中，logit 值通常服從正態分佈，因此算出的R-top-k值非常接近真實top-k值。同時因為這一步只涉及到暫存器的讀寫，且演算法複雜度低，因此可以快速執行完成（十幾個指令週期）。實際觀察發現，在top-4設定下，根據R-top-k只會從幾萬token中粗選出十幾個候選，因此非常高效。第二次遍歷中，根據R-top-k粗選出候選，同時對 logit 值按 batch_id 做了值偏移，多執行緒併發寫入視訊記憶體中的候選佇列。
粗選完成後，在候選佇列中進行一次排序，就能得到整個batch中每個序列的準確top-k值，然後更新快取，一步解碼過程就快速執行完成了。
下面是k=2，詞表大小=8的情況下一個具體的示例（列代表第幾個字元輸出，行代表每個位置的候選）。可以看出，原來需要對 16 個元素進行排序，而採用層級解碼之後，最後只需要對 5 個元素排序即可，大大降低了排序的複雜度。
視覺化分析計算延遲
為了驗證上面幾種優化技術的實際效果，筆者用 GPU profile 工具，對 LightSeq 的一次推理過程進行了延遲分析。下圖展示了 32 位浮點數和 16 位浮點數精度下，各計算模組的延遲佔比：
可以發現，在兩種計算精度下：
- 經過優化後，cuBLAS 中的矩陣乘法計算延遲分別佔比 82% 和 88% ，成為推理加速新的主要瓶頸。而作為對比，我們測試了 Tensorflow 模型，矩陣乘法計算延遲只佔了 25% 。這說明 LightSeq 的 beam search 優化已經將延遲降到了非常低的水平。
- 快取重新整理分別佔比 10% 和 6% ，比重也較高，但很難繼續優化。今後可以嘗試減少快取量（如降低 decoder 層數，降低快取精度等）來繼續降低延遲。
- 其他運算總計佔比 8% 和 6% ，包括了 Layer Normalization、beam search 和中間結果的視訊記憶體讀寫等。
視覺化結果說明了 LightSeq 已經做到了極致優化，大大提升了推理速度。
傳送門：
GitHub專案地址：
https://github.com/bytedance/...
[1] Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems. 2017.
[2] Devlin, Jacob, et al. "Bert: Pre-training of deep bidirectional transformers for language understanding." arXiv preprint arXiv:1810.04805 (2018).
[3] Brown, Tom B., et al. "Language models are few-shot learners." arXiv preprint arXiv:2005.14165 (2020).
[4] WMT2020, http://www.statmt.org/wmt20/
[5] Li, Jiwei, Will Monroe, and Dan Jurafsky. "A simple, fast diverse decoding algorithm for neural generation." arXiv preprint arXiv:1611.08562 (2016).
[6] TurboTransformers, https://github.com/Tencent/Tu...
[7] FasterTransformer, https://github.com/NVIDIA/Dee...
[8] NVIDIA Triton Inference Server, https://github.com/triton-inf...
[9] LightSeq proto, https://github.com/bytedance/...
[10] LightSeq效能評測報告, https://github.com/bytedance/...
[11] LightSeq Layer Normalization, https://github.com/bytedance/...
[12] cuBLAS, https://docs.nvidia.com/cuda/...
[13] GPT2,"Language Models are Unsupervised Multitask Learners"