[論文レビュー] Whale: A Unified Distributed Training Framework.

Ang Wang, Xianyan Jia+4 more|arXiv (Cornell University)|Nov 18, 2020

Advanced Neural Network Applications参考文献 39被引用数 4

ひとこと要約

Whaleは、メタ駆動型コストモデルを用いた自動並列化と新規IR抽象化を通じて、データ並列、モデル並列、パイプライン並列、演算子シャーディング、ハイブリッド並列を含む多様な並列戦略を統合的にサポートする分散学習フレームワークである。64 GPU上でHorovodのデータ並列化と比較して、BERT-Largeの学習を2.32倍高速化し、大規模な画像分類タスクを14.8倍高速化するが、TensorFlowとの互換性を保ちつつコード変更を最小限に抑える。

ABSTRACT

Data parallelism (DP) has been a common practice to speed up the training workloads for a long time. However, with the increase of data size and model size, DP has become less optimal for most distributed training workloads. Moreover, it does not work on models whose parameter size cannot fit into a single GPU's device memory. To enable and further improve the industrial-level giant model training, we present Whale, a unified distributed training framework. It provides comprehensive parallel strategies including data parallelism, model parallelism, operator sharding, pipeline, hybrid strategy, and automatic parallel strategy. To express complex training strategies effectively and efficiently in one framework, Whale IR is designed as the basic unit to explore and implement different distributed strategies. Moreover, Whale enables automatic parallelism upon using a meta-driven cost model. Whale is compatible with TensorFlow and can easily distribute training tasks by adding a few code lines without changing user model code. To the best of our knowledge, Whale is the first work that can support various hybrid distributed strategies within one framework. In our experiment of Bert Large model, Whale pipeline strategy is 2.32 times faster than Horovod data parallelism (HDP) on 64 GPUs. In a large-scale image classification task (100,000 classes), Whale hybrid strategy, which consists of operator sharding and DP, is 14.8 times faster than HDP on 64 GPUs.

研究の動機と目的

• 大規模なモデルやデータセットへのスケーリングにおけるデータ並列化の限界を解消すること。
• 単一GPUのメモリ容量を超えるパラメータを持つ巨大モデルの効率的学習を可能にすること。
• 単一のシステム内で多様な並列戦略を統合的にサポートする統一フレームワークを提供すること。
• メタ駆動型コストモデルを用いた自動並列戦略選択により、学習パフォーマンスを最適化すること。
• 最小限のコード変更で既存のTensorFlowモデルと完全に後方互換性を持つこと。

提案手法

• 複雑な分散学習戦略を表現・組み合わせる基盤抽象化としてWhale IRを設計すること。
• 包括的な並列化戦略の実装：データ並列、モデル並列、演算子シャーディング、パイプライン並列、およびハイブリッド組み合わせ。
• メタ駆動型コストモデルを用いて最適な戦略を選択する自動並列化エンジンの統合。
• 分散学習のデプロイに必要なコード変更を最小限に抑えることで、TensorFlowとのシームレスな統合を実現すること。
• 高レベルのモデルを効率的な分散実行計画にマップする、統一されたIRベースのコンパイルパイプラインの使用。

実験結果

リサーチクエスチョン

• RQ1 単一のフレームワークが、ハイブリッド組み合わせを含む多様な分散学習戦略を効率的にサポートできるか。
• RQ2 コストモデルに基づく自動戦略選択が、手動または固定戦略選択と比較して、学習効率をどの程度向上できるか。
• RQ3 BERT-Large や10万クラスの画像分類タスクといった大規模モデルに対して、既存のデータ並列ベースラインと比較して顕著な高速化を達成できるか。
• RQ4 モデルコードの変更なしに、TensorFlowなどの既存のディープラーニングフレームワークと互換性を保てるか。
• RQ5 大規模な学習ワークロードにおいて、演算子シャーディングとハイブリッド戦略によって達成可能なパフォーマンス向上はどの程度か。

主な発見

• Whaleのパイプライン戦略は、64 GPU上でBERT-Largeの学習においてHorovodのデータ並列化と比較して2.32倍の高速化を達成した。
• Whaleのハイブリッド戦略（演算子シャーディングとデータ並列の組み合わせ）は、64 GPU上で10万クラスの画像分類タスクにおいてHorovodのデータ並列化と比較して14.8倍の高速化を達成した。
• モデル並列と演算子シャーディングを活用することで、単一GPUのメモリ容量を超えるパラメータサイズのモデルのサポートに成功した。
• メタ駆動型コストモデルによる自動並列化により、手動チューニングなしに最適な戦略選択が可能になった。
• WhaleはTensorFlowと完全に互換性を持ち、既存モデルに対して最小限のコード変更で分散学習を可能にした。