1024塊TPU在燃燒！BERT訓練從3天縮短到76分鐘

「2019 Python開發者日」全日程揭曉，請掃碼諮詢

作者 | 琥珀

出品 | AI科技大本營（ID:rgznai100）

「Jeff Dean稱讚，TensorFlow官方推特支持，BERT目前工業界最耗時的應用，計算量遠高於ImageNet。我們將BERT的訓練時間從三天縮短到了一小時多。」UC Berkeley大學在讀博士尤洋如是說道。

近日，來自Google、UC Berkeley、UCLA研究團隊再度合作，成功燃燒1024塊TPU，將BERT預訓練模型的訓練時長從3天縮減到了76分鐘。batch size技術是加速神經網路訓練的關鍵，在「Reducing BERT Pre-Training Time from 3 Days to 76 Minutes」這篇論文中，作者提出了LAMB優化器，它支持自適應元素更新和分層校正。

論文傳送門：https://arxiv.org/pdf/1904.00962.pdf

論文摘要：batch size增加到很大時的模型訓練是加速大型分散式系統中深度神經網路訓練的關鍵。但是，這種模型訓練很難，因為它會導致一種泛化差距。直接優化通常會導致測試集上的準確性下降。

BERT是一種先進的深度學習模型，它建立在語義理解的深度雙向轉換器上。當我們增加batch size的大小（如超過8192）時，此前的模型訓練技巧在BERT上表現得並不好。BERT預訓練也需要很長時間才能完成，如在16個TPUv3上大約需要三天。

為了解決這個問題，我們提出了LAMB優化器，可將batch size擴展到65536，且不會降低準確率。LAMB是一個通用優化器，batch size大小均使用，且除了學習率之外不需要別的參數調整。

基線BERT-Large模型需要100萬次迭代才能完成預訓練，而batch size大小為65536/32768的LAMB僅需要8599次迭代。我們還將batch size進行內存限制，接近TPUv3 pod，結果可在76分鐘內完成BERT訓練。

據悉，該論文的一作是來自UC Berkeley計算機科學部的在讀博士尤洋，同時也是Google Brain的實習生。據公開信息顯示，尤洋的導師是美國科學院與工程院院士，ACM/IEEE fellow，伯克利計算機系主任，以及首批中關村海外顧問James Demmel教授。他當前的研究重點是大規模深度學習訓練演算法的分散式優化。2017年9月，尤洋等人的新演算法以24分鐘完成ImageNet訓練，刷新世界紀錄 。

在此之前，他曾在英特爾實驗室、微軟研究院、英偉達、IBM沃森研究中心等機構實習。尤洋本科就讀於中國農業大學計算機系，碩士保送清華大學計算機系，是一名杠杠的理工學霸！

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！