[論文レビュー] Data Movement Is All You Need: A Case Study on Optimizing Transformers

Andrei Ivanov, Nikoli Dryden+3 more|arXiv (Cornell University)|Jun 30, 2020

Advanced Neural Network Applications参考文献 96被引用数 24

ひとこと要約

この論文は、Transformerの学習におけるデータ移動を主なボトルネックとして特定し、計算の進歩にもかかわらず、アムダールの法則に起因して、学習がメモリ制約となったことを示している。著者らは、体系的なデータフロー解析とレイアウト変換を通じて、データレイアウトと移動をグローバルに最適化することで、データ移動量を最大22.91％削減し、最先端のフレームワークと比較して、BERTエンコーダ層で1.30倍の高速化、全BERTモデルで1.19倍の高速化を達成した。

ABSTRACT

Transformers are one of the most important machine learning workloads today. Training one is a very compute-intensive task, often taking days or weeks, and significant attention has been given to optimizing transformers. Despite this, existing implementations do not efficiently utilize GPUs. We find that data movement is the key bottleneck when training. Due to Amdahl's Law and massive improvements in compute performance, training has now become memory-bound. Further, existing frameworks use suboptimal data layouts. Using these insights, we present a recipe for globally optimizing data movement in transformers. We reduce data movement by up to 22.91% and overall achieve a 1.30x performance improvement over state-of-the-art frameworks when training a BERT encoder layer and 1.19x for the entire BERT. Our approach is applicable more broadly to optimizing deep neural networks, and offers insight into how to tackle emerging performance bottlenecks.

研究の動機と目的

• Transformerモデルの学習における主なパフォーマンスボトルネックを特定すること。
• 高度に最適化されたカーネルを備えても、既存のディープラーニングフレームワークがGPU計算を十分に活用できない理由を分析すること。
• データ移動と最適でないデータレイアウトが、学習パフォーマンスを低下させる役割を調査すること。
• 計算グラフ全体にわたるデータ移動を体系的かつ一貫して削減するグローバル最適化戦略を開発すること。
• データ移動最適化が、アルゴリズム的改善やハードウェアレベルの改善単独よりも大きなパフォーマンス向上をもたらすかどうかを検証すること。

提案手法

• 著者らは、Transformerの計算を状態保持型データフローマルチグラフ（SDFG）としてモデル化することで、演算間でのデータ移動を体系的に分析できるようにした。
• 詳細なデータフロー解析を実施し、特にマルチヘッドアテンションにおいて、実行時間の37％がデータ移動に費やされていることが判明した。
• 本手法は、演算間でデータレイアウトを再配置・再構築することで、重複するデータ移動を最小限に抑えるグローバルレイアウト変換戦略を導入した。
• パラメトリックで平坦なデータフロー表現を用いることで、複数の粒度レベルにわたる粗いおよび細かいレイアウト最適化および統合最適化を可能にした。
• 著者らは、独自のSDFGベースのコンパイラスタックを用いて最適化の有効性を評価し、PyTorch、XLAを搭載したTensorFlow、cuDNNと比較した。
• メモリアクセスのオーバーヘッドを低減するレイアウト変換を適用し、行列同士の積（MMM）のパフォーマンスを最大52％向上させた。

実験結果

リサーチクエスチョン

• RQ1 現代のTransformerモデルの学習における主なパフォーマンスボトルネックは何か？
• RQ2 高度に最適化されたカーネルを備えても、既存のディープラーニングフレームワークがGPU計算を効率的に活用できないのはなぜか？
• RQ3 データレイアウトは、Transformerにおけるデータ移動と全体の学習パフォーマンスにどのように影響するか？
• RQ4 局所的またはカーネルレベルの最適化に比べて、グローバルかつエンドツーエンドのデータ移動最適化は、より優れた性能向上を達成できるか？
• RQ5 データ移動の削減は、BERTおよびその部品における学習スループットの向上にどの程度寄与するか？

主な発見

• BERTの学習イテレーションにおいて、データ移動が実行時間の37％を占めているにもかかわらず、テンソルの畳み込み演算が99％以上の算術演算を占めている。
• グローバルなレイアウト最適化と統合最適化により、データ移動量を最大22.91％削減し、パフォーマンスを顕著に向上させた。
• 提案手法は、最先端のフレームワークと比較して、1つのBERTエンコーダ層で1.30倍の高速化、全BERTモデルで1.19倍の高速化を達成した。
• パフォーマンス向上の主な要因は最適化されたデータレイアウトであり、行列同士の積（MMM）のパフォーマンスを最大52％向上させた。
• アムダールの法則と、計算帯域とメモリ帯域の間の不均衡な成長に起因して、学習がメモリ制約となったことが確認された。
• 本手法はTransformerに限らず、他のディープニューラルネットワークにおけるデータ移動最適化のフレームワークとしても汎用的である。