速度超快！字節跳動開源序列推理引擎LightSeq

作者：佚名 2020-10-24 07:30:05

LightSeq 可以應用于機器翻譯、自動問答、智能寫作、對話回復生成等眾多文本生成場景，大大提高線上模型推理速度，改善用戶的使用體驗，降低企業的運營服務成本。

這應該是業界第一款完整支持 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具有如下幾點優勢：

1. 高性能

LightSeq推理速度非常快。例如在翻譯任務上，LightSeq相比于Tensorflow實現最多可以達到14倍的加速。同時領先目前其他開源序列推理引擎，例如最多可比Faster Transformer快1.4倍。

2. 支持模型功能多

LightSeq支持BERT、GPT、Transformer、VAE 等眾多模型，同時支持beam search、diverse beam search[5]、sampling等多種解碼方式。下表詳細列舉了Faster Transformer[7]、Turbo Transformers[6]和LightSeq三種推理引擎在文本生成場景的功能差異：

3. 簡單易用，無縫銜接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]，搭建起模型服務。

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。這為模型推理帶來了眾多挑戰：

1. IO 密集型計算的細粒度核函數調用帶來大量冗余顯存讀寫，成為特征計算性能瓶頸。

2. 復雜動態 shape 為計算圖優化帶來挑戰，導致模型推理期間大量顯存動態申請，耗時較高。

3. 解碼生成每一步字符過程邏輯復雜，難以并行化計算從而發揮硬件優勢。

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 個元素排序即可，大大降低了排序的復雜度。