[論文レビュー] Boosting Distributed Training Performance of the Unpadded BERT Model

Jinle Zeng, Min Li+5 more|arXiv (Cornell University)|Aug 17, 2022

Topic Modeling被引用数 4

ひとこと要約

本論文は、可変長入力を用いてパディングトークン由来の冗長な計算を排除する、高度に最適化されたパディングなしBERTモデルを提案する。この手法は、グループ化されたマルチストリームFMHA、重複するデータ交換を用いたワークロードバランス、およびカーネル/オペレータ最適化を活用する。このアプローチにより、8x A100 GPU上でMLPerf Training v2.0で1位を獲得し、2578サンプル/秒の最先端のスループットを達成。収束にはたった16.598分で完了した。

ABSTRACT

Pre-training models are an important tool in Natural Language Processing (NLP), while the BERT model is a classic pre-training model whose structure has been widely adopted by followers. It was even chosen as the reference model for the MLPerf training benchmark. The distributed training performance optimization of BERT models plays an important role in accelerating the solutions of most NLP tasks. BERT model often uses padding tensors as its inputs, leading to excessive redundant computations. Thus, removing these redundant computations is essential to improve the distributed training performance. This paper designs a new approach to train BERT models with variable-length inputs efficiently. Firstly, we propose a general structure for the variable-length BERT models, and accelerate the encoder layer via our grouped multi-stream FMHA (Fused Multi-Head Attention) method. Secondly, through data exchange, we address the unbalanced workload problem caused by the variable-length inputs, which overlaps highly with the training process. Finally, we optimize the overall performance of the BERT model, such as kernel fusion, and operator optimization. Our experimental results show that our highly optimized BERT model achieves state-of-the-art throughput and ranks first in MLPerf Training v2.0 within the same GPU configuration. The optimizations in this paper can be applied to more BERT-like models in our future works.

研究の動機と目的

• パディングトークン由来の冗長な計算を排除し、分散トレーニングのパフォーマンスを低下させる要因を解消すること。
• 分散トレーニング環境において可変長入力が引き起こすワークロードの不均衡を解消すること。
• 同じハードウェア構成下で、BERT-Largeのトレーニングスループットを最先端水準に達するようにすること。
• アテンション機構、カーネル統合、オペレータレベルのパフォーマンス最適化を含む、トレーニングパイプライン全体を最適化すること。

提案手法

• パディングを伴わない動的シーケンス長をサポートする、可変長BERTモデルの一般化アーキテクチャを設計する。
• エンコーダー層の計算を最適化することで高速化する、グループ化されたマルチストリームFMHA（統合マルチヘッドアテンション）を導入する。
• GPU計算と重なりを取ったデータ交換により、アイドル時間と通信オーバーヘッドを最小限に抑えることで、ワークロードバランスを実現する。
• メモリ帯域幅の削減とGPU利用効率の向上を図るために、カーネル統合とオペレータレベルの最適化を適用する。
• LAMB最適化手法を用いた混合精度トレーニングを実装し、デバイス間での一貫性のあるバッチ処理を保証する。
• ディープラーニングフレームワークにPaddlePaddleを採用し、8x NVIDIA A100 400W GPUを用いてベンチマークを実施する。

実験結果

リサーチクエスチョン

• RQ1 BERTトレーニングにおいて、パディングトークン由来の冗長な計算を効果的に排除する方法は何か？
• RQ2 分散トレーニング環境でシーケンス長が可変である場合、デバイス間のワークロードをバランスさせるためにどのような技術が有効か？
• RQ3 計算と重なりを取ったデータ交換は、分散トレーニングの効率をどのように向上させるか？
• RQ4 カーネルおよびオペレータ最適化のどの組み合わせが、BERT-Largeのスループットを最大限に高めるか？
• RQ5 完全にパディングなしのBERTモデルは、MLPerfのような標準化ベンチマークで、既存の最適化実装を上回る性能を発揮できるか？

主な発見

• 最適化されたパディングなしBERTモデルは、同じGPU構成下で2578サンプル/秒のスループットを達成し、他のすべての実装を上回った。
• MLPerf Training v2.0で1位を獲得し、16.598分で72%のMLM精度に収束した。NVIDIA、HazyResearch、その他の提出物を上回った。
• パディング関連の計算を排除するだけで、ベースラインのパディングあり実装に比べて2.3倍の高速化が達成された。
• カーネル統合最適化は8.9%、オペレータレベル最適化は11.3%のパフォーマンス向上をもたらした。
• 計算と重なりを取ったデータ交換により、2.8%の向上が得られ、通信ボトルネックが軽減された。
• グループ化されたマルチストリームFMHA最適化により、通常のunpad FMHAに比べて3.6%のパフォーマンス向上が達成された。