[論文レビュー] Multi-node Bert-pretraining: Cost-efficient Approach

Jiahuang Lin, Xin Li+1 more|arXiv (Cornell University)|Aug 1, 2020

Topic Modeling参考文献 26被引用数 7

ひとこと要約

本稿では、10 Gbpsネットワークを介した通信ボトル neck を軽減し、アルゴリズム的およびソフトウェア的最適化を用いて単一デバイスのスループットを向上させるとともに、複数ノード間でのトレーニングワークロードを効率的に分散させることで、256台のNVIDIA T4 GPUを搭載した32ノードの学術的クラスタ上で、12日間でBERT-largeの事前学習を実現するコスト効率の高いマルチノードフレームワークを提示している。これにより、弱スケーリング効率70%を達成し、工業的スケールの設定と同等のモデルパフォーマンスを、コストの僅か数分の1で実現した。

ABSTRACT

Recently, large scale Transformer-based language models such as BERT, GPT-2, and XLNet have brought about exciting leaps in state-of-the-art results for many Natural Language Processing (NLP) tasks. One of the common trends in these recent models is a significant increase in model complexity, which introduces both more weights and computation. Moreover, with the advent of large-scale unsupervised datasets, training time is further extended due to the increased amount of data samples within a single training epoch. As a result, to train these models within a reasonable time, machine learning (ML) programmers often require advanced hardware setups such as the premium GPU-enabled NVIDIA DGX workstations or specialized accelerators such as Google's TPU Pods. Our work addresses this limitation and demonstrates that the BERT pre-trained model can be trained within 2 weeks on an academic-size cluster of widely available GPUs through careful algorithmic and software optimizations. In this paper, we present these optimizations on how to improve single device training throughput, distribute the training workload over multiple nodes and GPUs, and overcome the communication bottleneck introduced by the large data exchanges over the network. We show that we are able to perform pre-training on BERT within a reasonable time budget (12 days) in an academic setting, but with a much less expensive and less aggressive hardware resource requirement than in previously demonstrated industrial settings based on NVIDIA DGX machines or Google's TPU Pods.

研究の動機と目的

• 限られたハードウェア予算を持つ学術的環境で、効率的かつコスト効率の良いBERT-large事前学習を可能にすること。
• BERT-largeのような大規模なTransformerモデルを学習するための高い計算的・財務的障壁を克服すること。
• コモディティハードウェア上で単一GPUのスループットとマルチノード分散トレーニングのパフォーマンスを最適化すること。
• 勾配蓄積とシステムレベルのチューニングを通じて、分散トレーニングにおける通信オーバーヘッドを低減すること。
• 業界基準の規模のモデル品質が、工業的ベンチマークと比較してはるかに低いハードウェアコストで達成可能であることを示すこと。

提案手法

• ノードあたり8台のNVIDIA T4 GPUを搭載した32ノードのクラスタを用い、合計256台のGPUを10 Gbpsネットワークで接続した。
• ノード間通信の頻度とネットワークトラフィックを低減するため、4ステップにわたる勾配蓄積を適用した。
• 混合精度トレーニングと効率的なメモリ管理による単一GPUトレーニングスループットの最適化を実施した。
• 2段階の事前学習を実装した：第1段階はシーケンス長128（36エポック）、第2段階はシーケンス長512（6エポック）、それぞれ異なるバッチサイズと周波数を用いた。
• パラメータの更新に同期的確率的勾配降下法を用い、ノード間でモデル並列とデータ並列を実装した。
• 収束問題の診断とトレーニング安定性の向上のため、損失曲線解析とハイパーパramータチューニングを実施した。

実験結果

リサーチクエスチョン

• RQ1 学術的グレードのハードウェアと10 Gbpsネットワークのみを用いて、BERT-largeを12日以内に事前学習できるか？
• RQ2 コモディティマルチGPUクラスタ上で高いスケーリング効率を達成するために必要なシステムレベルの最適化は何か？
• RQ3 勾配蓄積は、分散BERT事前学習における通信オーバーヘッドとトレーニングスループットにどのように影響するか？
• RQ4 なぜBERT事前学習の第2段階で収束の不安定さが生じるのか、そしてハードウェアコストを増加させることなくこれを解消できるか？
• RQ5 高価なDGXやTPUシステムと比較して、低コストのT4 GPUでトレーニングした場合、モデルパフォーマンスはどの程度維持できるか？

主な発見

• 著者らは、32ノードの学術的クラスタ上に256台のT4 GPUを用いて、12日間でBERT-largeの事前学習を完了し、70%の弱スケーリング効率を達成した。
• 総ハードウェアコストは約62万4千ドルであり、同等のDGXやTPU Podセットアップの1500万～1280万ドルと比較して24倍低い。
• 微調整後、SQuAD v1.1で81%～83%のF1スコアを達成し、Googleの90.9%およびNVIDIAの90%～91%と比較して9%～10%以内の差にとどまった。
• 第2段階の収束問題は、システム最適化ではなくハイパーパramータの誤設定に起因しており、より良いチューニングにより第2段階は6エポックではなく4エポックで完了可能であった。
• 最適化済みと非最適化の両方の実行における損失曲線の比較から、パフォーマンス最適化が収束に悪影響を及げず、むしろ第2段階でより安定した挙動を示したことがわかった。
• Google Cloud上で12日間256台のT4 GPUをレンタルするコストは2万5739ドルであり、ハードウェアの資本コストの24倍も低いが、ハードウェアの耐用年数が3年間あるため、90回分の実験が可能である。