Perseus-BERT——業內性能極致優化的BERT訓練方案

一，背景——橫空出世的BERT全面超越人類

2018年在天然語言處理（NLP）領域最具爆炸性的一朵「蘑菇雲」莫過於Google Research提出的BERT（Bidirectional Encoder Representations from Transformers）模型。做爲一種新型的語言表示模型，BERT以「摧枯拉朽」之勢橫掃包括語言問答、理解、預測等各項NLP錦標的桂冠，見圖1和圖2。

【圖1】SQuAD是基於Wikipedia文章的標準問答數據庫的NLP錦標。目前SQuAD2.0排名前十名均爲基於BERT的模型（圖中列出前五名），前20名有16席均是出自BERT

【圖2】GLUE是一項通用語言理解評估的benchmark，包含11項NLP任務。BERT自誕生日起長期壓倒性霸佔榜首（目前BERT排名第二，第一爲Microsoft提交的BIGBIRD模型，因爲沒有URL連接無從知曉模型細節，網傳BIGBIRD的名稱上有借鑑BERT BIG模型之嫌）

業內將BERT在天然語言處理的地位比做ResNet之於計算機視覺領域的里程碑地位。在BERT橫空出世以後，全部的天然語言處理任務均可以基於BERT模型爲基礎展開。

一言以蔽之，現現在，做爲NLP的研究者，若是不瞭解BERT，那就是落後的科技工做者；做爲以天然語言處理爲重要依託的科技公司，若是不落地BERT，那就是落後生產力的表明。

二，痛點——算力成爲BERT落地的攔路虎

BERT強大的緣由在哪裏？讓咱們拂去雲靄，窺探下硝煙下的奧祕。

BERT模型分爲預訓練模型（Pretrain）和精調模型（Finetune）。Pretrain模型爲通用的語言模型。Finetune只須要在Pretrain的基礎上增長一層適配層就能夠服務於從問答到語言推理等各種任務，無需爲具體任務修改總體模型架構，如圖3所示。這種設計方便BERT預處理模型適配於各種具體NLP模型（相似於CV領域基於ImageNet訓練的各類Backbone模型）。

【圖3】左圖基於BERT pretrain的模型用於語句問答任務（SQuAD）的finetune模型，右圖爲用於句對分類（Sentence Pair Classification Tasks）的finetune模型。他們均是在BERT Pretrain模型的基礎上增長了一層具體任務的適配層

所以，BERT的強大主要歸功於精確度和魯棒性俱佳的Pretrain語言模型。大部分的計算量也出自Pretrain模型。其主要運用瞭如下兩項技術，都是極其耗費計算資源的模塊。

1. 雙向Transformer架構

圖4可見，與其餘pre-training的模型架構不一樣，BERT從左到右和從右到左地同時對語料進行transformer處理。這種雙向技術能充分提取語料的時域相關性，但同時也大大增長了計算資源的負擔。【關於Transformer是Google 17年在NLP上的大做，其用全Attention機制取代NLP經常使用的RNN及其變體LSTM等的經常使用架構，大大改善了NLP的預測準確度。本文不展開，該興趣的同窗能夠自行搜索一下】。

【圖4】Pretrain架構對比。其中OpenAI GPT採用從左到右的Transformer架構，ELMo採用部分從左到右和部分從右到左的LSTM的級聯方式。BERT採用同時從左到右和從右到左的雙向Transformer架構。

1. 詞/句雙任務隨機預測

BERT預訓練模型在迭代計算中會同時進行單詞預測和語句預測兩項非監督預測任務。

其一，單詞預測任務對語料進行隨機MASK操做（Masked LM）。在全部語料中隨機選取15%的單詞做爲Mask數據。被選中Mask的語料單詞在迭代計算過程當中80%時間會被掩碼覆蓋用於預測、10%時間保持不變、10%時間隨機替換爲其餘單詞，如圖5所示。

其二，語句預測任務（Next Sentence Prediction）。對選中的先後句A和B，在整個迭代預測過程當中，50%的時間B做爲A的真實後續語句（Label=IsNext），另外50%的時間則從語料庫裏隨機選取其餘語句做爲A的後續語句（Label=NotNext），如圖5所示

【圖5】詞/句雙任務隨機預測輸入語料實例。藍框和紅框爲同一個語料輸入在不一樣時刻的隨機狀態。對單詞預測任務，藍框中的「went」爲真實數據，到了紅框則被[MASK]，紅框中的「the」 則相反；對於語句預測任務，藍框中的句組爲真實的先後句，而紅框中的句組則爲隨機的組合。

這種隨機選取的單詞/語句預測方式在功能上實現了非監督數據的輸入的功能，有效防止模型的過擬合。可是按比例隨機選取須要大大增長對語料庫的迭代次數才能消化全部的語料數據，這給計算資源帶來了極大的壓力。

綜上，BERT預處理模型功能須要創建在極強的計算力基礎之上。BERT論文顯示，訓練BERT BASE 預訓練模型（L=12, H=768, A=12, Total Parameters=110M, 1000,000次迭代）須要1臺Cloud TPU工做16天；而做爲目前深度學習主流的Nvidia GPU加速卡面對如此海量的計算量更是力不從心。即便是目前主流最強勁的Nvidia V100加速卡，訓練一個BERT-Base Pretrain模型須要一兩個月的時間。而訓練Large模型，須要花至少四五個月的時間。

花幾個月訓練一個模型，對於絕大部分在GPU上訓練BERT的用戶來講真是傷不起。

三，救星——擎天雲加速框架爲BERT披荊斬棘

阿里雲彈性人工智能團隊依託阿里雲強大的基礎設施資源打磨業內極具競爭力的人工智能創新方案。基於BERT的訓練痛點，團隊打造了擎天優化版的Perseus-BERT, 極大地提高了BERT pretrain模型的訓練速度。在雲上一臺V100 8卡實例上，只需4天不到便可訓練一份BERT模型。

Perseus-BERT是如何打造雲上最佳的BERT訓練實踐？如下乾貨爲您揭祕Perseus-BERT的獨門絕技。

1. Perseus 統一分佈式通訊框架 —— 賦予BERT分佈式訓練的輕功

Perseus（擎天）統一分佈式通訊框架是團隊針對人工智能雲端訓練的痛點，針對阿里雲基礎設施極致優化的分佈式訓練框架。其可輕便地嵌入主流人工智能框架的單機訓練代碼，在保證訓練精度的同時高效地提高訓練的多機擴展性。擎天分佈式框架的乾貨介紹詳見團隊另外一篇文章《Perseus(擎天)：統一深度學習分佈式通訊框架》。

針對tensorflow代碼的BERT，Perseus提供horovod的python api方便嵌入BERT預訓練代碼。基本流程以下：

• 讓每塊GPU對應一個Perseus rank進程；
• 對global step和warmup step作基於rank數的校準；
• 對訓練數據根據rank-id作劃分；
• 給Optimizer增長DistributeOptimizer的wrapper。

值得注意的是，BERT源碼用的自定義的Optimizer，在計算梯度時採用瞭如下api

grads = tf.gradients(loss, tvars)

Perseus的DistributeOptimizer繼承標準的Optimizer實現，並在compute_gradients api 上實現分佈式的梯度更新計算。所以對grads獲取作了以下微調

grads_and_vars  = optimizer.compute_gradients(loss, tvars)

grads = list()

for grad, var in grads_and_vars:

  grads.append(grad)

2. 混合精度訓練和XLA編譯優化——提高BERT單機性能的內功

混合精度

在深度學習中，混合精度訓練指的是float32和float16混合的訓練方式，通常的混合精度模式如圖6所示

【圖6】混合精度訓練示例。在Forward+Backward計算過程當中用float16作計算，在梯度更新時轉換爲float32作梯度更新。

混合梯度對Bert訓練帶來以下好處，

• 增大訓練時的batch size和sequence_size以保證模型訓練的精度。

  目前阿里雲上提供的主流的Nvidia顯卡的顯存最大爲16GB，對一個BERT-Base模型在float32模式只能最高設置爲sequence_size=256，batch_size=26。BERT的隨機預測模型設計對sequence_size和batch_size的大小有必定要求。爲保證匹配BERT的原生訓練精度，須要保證sequece_size=512的狀況下batch_size不小於16。Float16的混合精度能夠保證如上需求。

• 混合精度能充分利用硬件的加速資源。

  NVidia從Volta架構開始增長了Tensor Core資源，這是專門作4x4矩陣乘法的fp16/fp32混合精度的ASIC加速器，一塊V100能提供125T的Tensor Core計算能力，只有在混合精度下計算才能利用上這一塊強大的算力。

受限於float16的表示精度，混合精度訓練的代碼須要額外的編寫，NVidia提供了在Tensorflow下作混合精度訓練的教程 。其主要思路是經過tf.variable_scope的custom_getter 參數保證存儲的參數爲float32並用float16作計算。

在BERT預訓練模型中，爲了保證訓練的精度，Perseus-BERT沒有簡單的利用custom_getter參數，而是顯式指定訓地參數中哪些能夠利用float16不會影響精度，哪些必須用float32已保證精度。咱們的經驗以下：

• Embedding部分要保證float32精度；
• Attetion部分能夠利用float16加速；
• Gradients相關的更新和驗證須要保證float32精度；
• 非線性激活等模塊須要保證float32精度。

XLA編譯器優化

XLA是Tensorflow新近提出的模型編譯器，其能夠將Graph編譯成IR表示，Fuse冗餘Ops，並對Ops作了性能優化、適配硬件資源。然而官方的Tensorflow release並不支持xla的分佈式訓練，爲了保證分佈式訓練能夠正常進行和精度，咱們本身編譯了帶有額外patch的tensorflow來支持分佈式訓練，Perseus-BERT 經過啓用XLA編譯優化加速訓練過程並增長了Batch size大小。

3. 數據集預處理的加速

Perseus BERT 同時對文本預處理作的word embedding和語句劃分作了並行化的優化。這裏就不展開說明。

四，性能——計算時間單位從月下降到天

圖7展現了Perseus BERT在P100實例上的性能，與開源主流的horovod相比，Peseus-BERT雙機16卡的分佈式性能是前者的5倍之多。

目前某大客戶已在阿里雲P100集羣上大規模上線了Perseus BERT，用10臺4卡P100只須要2.5天便可訓練完成業務模型，若是用開源的horovod（Tensorflow分佈式性能優化版）大概須要1個月的時間。

【圖7】Bert在阿里雲上P100實例的對比（實驗環境Bert on P100； Batch size: 22 ；Max seq length: 256 ；Data type:float32； Tensorflow 1.12； Perseus: 0.9.1；Horovod: 0.15.2）

爲了和Google TPU作對比，咱們量化了TPU的性能，性能依據如圖8。一個Cloud TPU可計算的BERT-Base性能 256 （1000000/4/4/24/60/60） = 185 exmaples/s。 而一臺阿里雲上的V100 單機八卡實例在相同的sequence_size=512下, 經過Perseus-BERT優化的Base模型訓練能夠作到 680 examples/s，接近一臺Cloud TPU的4倍性能。對一臺Cloud TPU花費16天才能訓練完的BERT模型，一臺阿里雲的V100 8卡實例只須要4天不到即可訓練完畢。*

【圖8】BERT Pretain在Google Cloud TPU上的性能依據*

五，總結——基於阿里雲基礎設施的AI極致性能優化

彈性人工智能團隊一直致力基於阿里雲基礎設施的AI極致性能優化的創新方案。Perseus-BERT就是一個很是典型的案例，咱們在框架層面上基於阿里雲的基礎設施作深度優化，充分釋放阿里雲上基礎資源的計算能力，讓阿里雲的客戶充分享受雲上的AI計算優點，讓天下沒有難算的AI。

 

原文連接 更多技術乾貨 請關注阿里云云棲社區微信號 ：yunqiinsight