NLPCC：預訓練在小米的推理優化落地

導讀：本文主要分享 小米 AI 實驗室 NLP 團隊 在 NLPCC 輕量級語言模型比賽 上的經驗，以及我們在預訓練模型推理優化上所作的工作和達到的實際落地後的效果。此次分享的目的是幫助大家快速進入比賽，以及瞭解工業界中使用 BERT 之類的大型預訓練模型時進行推理的相關優化手段。

01 背景介紹

首先和大家說一下比賽的背景和預訓練模型中存在的問題。

1. NLP 中的預訓練

隨着 BERT 的推出和發展，預訓練模型成爲目前工業界和比賽使用最廣泛的模型。目前在各大 NLP 任務 SOTA 榜單上排前幾名都是大型的預訓練模型，比如原生 BERT 或者它的一些變體。

預訓練模型的應用分爲兩個階段，先進行預訓練階段 ( pre-training )，然後進行微調階段 ( fine-tuning )。預訓練階段利用大規模的無監督預料，通常大於 100g 的量級，加上特定的無監督任務進行訓練。用來預訓練的任務有，經典的 NSP ( next sentence predict )、MLM ( masked language model )、以及一些變體，如阿里發佈的 Structural language model。另外在預訓練階段時，embedding 層的使用方式也有很多，比如 NEZHA 以及 XLNET 使用的相對位置編碼。Pretrain 階段一般比較消耗計算資源，一般都是在大量的 GPU 集羣上進行運算，並且計算很長的時間才能得到一個比較好的結果。

相對於複雜預訓練任務，下游的微調任務就比較簡單了。在預訓練好的模型的基礎上，只需要添加較少的網絡參數，並在在下游有監督數據上進行少量的訓練，就可以得到很不錯的成果，這也是預訓練技術能夠被大量使用的基礎條件。

2. 效率問題

然而就像天下沒有免費的午餐一樣，預訓練技術在帶來好處的同時也存在一些問題，這其中最爲凸顯的恐怕是效率問題了。

微軟在發佈 TuringNLG 模型的同時發佈了這種圖，從上圖我們可以看到在 BERT 發佈到現在，自然語言模型參數量的發展是指數型增大的，從 BERT 的 1 億參數量、到 GPT2 的 10 億、然後 TuringNLG 的 100 億。模型變大帶來的效率問題可以從訓練和推理兩個方面來體現。

在訓練階段，一個好的模型需要大量的語料支持，通常爲了得到一個不錯的模型，需要使用 100g 以上的訓練數據。同時爲了支持如此規模的數據集，使模型能擬合出不錯的效果，就需要使用大量昂貴的 GPU 資源。而在推理階段，大模型的推理時間比較長，那如何在高併發業務線上使用大模型，並且保證推理服務的性能 ( 比如服務延遲 99 分位數 ) 成爲一個難題。由此可見，在預訓練模型應用時必須考慮性價比，在效果提升和成本預算之間做一個合適的權衡。

3. 解決方向

針對這些效率上的問題，衍生出兩個比較清晰的解決方向和思路。

第一種方法是小模型或蒸餾，小模型是指參數特別少的模型，比如在 NLPCC 小模型比賽中聚焦的小於 12M 參數的模型。蒸餾是將訓練好的 BERT 類大模型的能力遷移到較小的模型中，目前業界聚焦的任務是將 12 層的 BERT-base 模型蒸餾到 6 層的 BERT 上，並保證較少的性能損失。這種方法的主要優勢是資源佔用率比較少，適用於線上的高併發業務；並且通過一些方法能夠實現與大模型類似的效果。

第二種方法是對大模型的優化，具體來說是指優化大模型在單張 GPU 上所能承受的最大流量，即吞吐量。這種方法的優勢是效果較好，比較適合低併發業務。

02 NLPCC 預訓練小模型比賽經驗分享

1. 比賽介紹

比賽要求的小模型必須是小於 12M 參數的預訓練模型，即 1/9 的 BERT-base 大小，最後在四個下游任務中測試模型效果並進行打分。評測涉及任務有：指代消歧 ( CLUEWSC2020 )、分類任務 ( CSL 論文關鍵詞識別 )、命名實體識別 ( CLUENER2020 ) 以及閱讀理解 ( CMRC 2018 )。其中，指代消歧任務用的是小數據集 ( 1.2k 左右 )，其餘任務所用數據集大小爲萬級別。在小數據集上微調模型，尤其是小預訓練模型，是比較影響模型效果的，後面我們會介紹這方面所做工作，希望能夠對讀者有所幫助。

用一句話來總結這個比賽，小模型的效果真的很差嗎？我們直接來看比賽的結果。

其中，第一條我們提交的 BERT-base 的參考。可以看到相較於 BERT 模型，小模型在分類、NER 以及閱讀理解任務上都可以將差距縮小到 1.5% 以內。而在 wsc 任務上由於數據集太小，導致微調時效果的抖動非常大效果會略差一點，而實際中很少有這麼小的數據集。

2. 經驗分享

接下來主要介紹一下我們的解決方案。在直接預訓練小模型和蒸餾的方式選擇上，我們採用了前者。

我們在比賽中所做的工作主要分爲了 3 個方面，首先是數據。我們一共收集了 160g 的原始中文語料，並剔除了一些口語、表情、微博段子、網絡流行用語比較多的語料，最終篩選出 35g+ 相對通用的語料。爲了快速得到實驗結果，我們在輸入長度上採用了 256 的大小，同時也正在嘗試 512 的長度，以便在閱讀理解等需要更長的句子輸入的任務上得到更好一點的效果。

在模型方面，我們選用了 6 層的高瘦模型，並通過以下公式選擇對應的 Hidden size 和 Vocab size 來保證模型參數總量小於 12M。

大概介紹一下訓練細節，我們使用了 8 張 V100 卡；採用了混合精度的方法；優化器爲 lamb 優化器；batch size 大概爲 14400，並每過十幾步累積一次梯度；預訓練任務選用了 wwm-MLM。沒有選用 NSP 任務的原因是，我們在測試 MLM + NSP 任務訓練時出現了收斂較慢的現象，原因仍有待分析。最終訓練結果是 MLM loss 達到了 1.5 ~ 1.8 之間，同時我們發現當 loss 到 1.8 左右時，分類任務和 NER 任務對應的表現已經很好了；繼續訓練當 loss 來到 1.5 附近時，閱讀理解任務的效果提升較大，而分類任務和 NER 任務的表現開始有一些抖動。

遇到的問題：

同時我們分享兩個在比賽過程中遇到的問題。

第一個是模型訓練時的梯度爆炸，具體是每當 loss 收斂到 1.5 附近時就會出現梯度爆炸的問題。再排除了數據中的問題後，最終我們找到的解決方法是將初始的學習率調小，之後再沒有遇到同樣的問題。

第二個問題是，我們再嘗試微軟發佈的 Bert of theseus 蒸餾方式時，發現該方法在本次小模型比賽中取得的效果不好，原因是使用該方法時 Hidden Size 不能改變，因此必須壓縮層數來達到參數量的限制。同時，我們嘗試用小的 Hidden Size 預訓練一個模型，再去蒸餾一個 12M 的模型。可能是受限於 Teacher 模型的效果限制，得到的結果是不好的。

3. 後續嘗試

同時，在比賽快結束的時候我們總結性的提出一些後續可以繼續嘗試的思路。

注：NLPCC 比賽結束後，我們又嘗試了一個數據增強和簡單的蒸餾方法，目前在 WSC 和 CLUENER 兩個任務上已經超過了 bert-base 的效果，並且排在了小模型排行榜第一名，詳見如下：

以上就是這次 NLPCC 小模型比賽的經驗分享，下面將分享一下大模型推理服務的優化以及在小米的落地情況。

03 大模型推理效率及在小米的優化落地

1. 相關知識和現狀

爲了方便大家閱讀，我們先介紹兩個工程上評價服務的指標。

爲什麼要介紹這兩個指標，是因爲線上服務要在短時間內接受大量的請求，爲了滿足使用者的體驗以及維護公司利益，必須要保證較高的可用性。

我們再來看一下在用大模型推理時，用一張 T4 卡能達到的最大 QPS。如果模型用的是 BERT base，句子計算長度設定爲 16，那麼最大的 QPS 大概是 100，這時 P99 已經比較高了，如果 QPS 設定爲 150 服務將變得不可用；同時，如果模型是 BERT-large，在同樣的句子長度下，模型根本無法運行。在這種情況下假設有一個日活 1000w+ 的服務要上線，QPS 對應要求爲 2000，並使用 BERT-base 作計算推理。那麼我們計算了一下，一共需要 20 張 T4 卡來滿足需求，並且 P99 延遲會大於 90ms，這是一個非常大的資源浪費。

2. 優化效果

開門見山的說，在優化後同樣一張 T4 卡上可以達到什麼樣的效果呢？大家可以看下錶中 QPS 和對應的 P99 表現，當 QPS 最大達到 3000 時，P99 的延遲仍然保持在 40ms 以下。

如果大家爲了追求更好的效果使用 BERT-large 模型，使用相同的推理優化方法能夠達到的數據如下表，極限 QPS 大概是 800 左右。

總結一下優化效果，優化前如果 BERT-base 模型要放在線上推理，需要使用 20 張 T4 來滿足需求。優化後可以將數量縮減到 1 張，一共節約了 19 倍的成本。

3. 實現方案

下面來介紹一下我們的優化方案。

簡單來說分爲三個：

第一個是 Tensorflow Serving 提供的 Batching 功能，該方法能夠提升 n 倍的服務性能；
第二個是 FP16 的量化，可以提升 2 倍服務性能；
第三個是 NVIDIA 開源的 Faster Transformer 庫。

Tensorflow Serving Batching 所做的是將多個請求合併在一起進行計算，這樣可以利用 GPU 的並行性能。之所以能這樣做是因爲 GPU 在計算一定量的 batch 所消耗的時間和計算單挑所用時間相差不多。具體的原理是服務端在規定好的一段時間裏持續接收請求並放入隊列等待，如果請求數達到規定的數量或者隊列等待時間超過規定時間，再進行推理運算。另外一個經驗是在使用 TF-Serving Batching 時，服務端做 Padding 會用隨機數填充，所以建議大家在 Client 端做 Padding。

在用 FP16 優化中，我們使用了 TensorCore 來做加速，將模型圖中 FP32 精度的節點全部轉換爲 FP16，從而降低模型推理需要計算的運算量 ( 每秒浮點運算次數 TFLOPs )。在 BERT-base 上實驗，這種方法很穩定，損失基本上控制在萬分位，是一個非常值得使用的方法。

第三個方法是用 Fast Transformer，大家感興趣可以去 github 瞭解源碼。我們在這基礎上所做的工作是，將 Fast Transformer 集成到 Tensorflow Serving 上，並且支持對可變長度輸入的處理等等。大家可以訪問英偉達的 github 主頁瞭解更多。

4. 線上效果

前面在評測這些優化方法時用到的數據都是壓測數據，那麼在真實線上業務的表現怎麼樣呢，我們用小愛同學的閒聊業務來舉例說明。

這個任務是一個問答排序問題，比如用戶問 " 你喜歡維尼小熊嘛 "，這時對應的答案可能有好多種，按照先檢索後排序的方法，需要對輸入和每個回答計算一個相似度。我們使用的方案是 Pointwise 的 BERT 二分類，使用的模型爲 6 層的 BERT，輸入長度爲 32，輸入爲 Query+Reply，輸出爲 0-1 的相關性。下面是服務的優化效果，我們一共分爲了三個版本，在使用了這三種優化方法後成功的將 13 張 GPU 壓縮到了 3 張 GPU，並且 P99 從 200ms+ 降到了 35ms，效果還是比較明顯的。