[論文レビュー] BAGUA: Scaling up Distributed Learning with System Relaxations

Shaoduo Gan, Jiawei Jiang+12 more|arXiv (Cornell University)|Jan 1, 2021

Stochastic Gradient Optimization Techniques参考文献 79被引用数 3

ひとこと要約

BAGUAは、量子化、分散化、非同期学習などの高度なシステム緩和技術を効率的に実装できるモジュラーでMPIスタイルの通信ライブラリである。BAGUAは、オーバーパイプラグ、統合、階層的通信を備えた柔軟な最適化フレームワークにより、最先端のアルゴリズムをネイティブでサポートすることで、さまざまなワークロードにおいて、PyTorch-DDP、Horovod、BytePSを最大2倍速く、エンドツーエンドの学習時間を短縮する。

ABSTRACT

Recent years have witnessed a growing list of systems for distributed data-parallel training. Existing systems largely fit into two paradigms, i.e., parameter server and MPI-style collective operations. On the algorithmic side, researchers have proposed a wide range of techniques to lower the communication via system relaxations: quantization, decentralization, and communication delay. However, most, if not all, existing systems only rely on standard synchronous and asynchronous stochastic gradient (SG) based optimization, therefore, cannot take advantage of all possible optimizations that the machine learning community has been developing recently. Given this emerging gap between the current landscapes of systems and theory, we build BAGUA, a MPI-style communication library, providing a collection of primitives, that is both flexible and modular to support state-of-the-art system relaxation techniques of distributed training. Powered by this design, BAGUA has a great ability to implement and extend various state-of-the-art distributed learning algorithms. In a production cluster with up to 16 machines (128 GPUs), BAGUA can outperform PyTorch-DDP, Horovod and BytePS in the end-to-end training time by a significant margin (up to 2 times) across a diverse range of tasks. Moreover, we conduct a rigorous tradeoff exploration showing that different algorithms and system relaxations achieve the best performance over different network conditions.

研究の動機と目的

• 分散学習アルゴリズムにおける理論的進展と、依然として標準的な同期/非同期SGDに依存する既存のシステムとの間のギャップを埋めること。
• 量子化、分散化、通信遅延などの多様なシステム緩和技術をネイティブにサポートする柔軟でモジュラーな通信ライブラリを設計すること。
• 統一された最適化フレームワークを通じて、最先端の分散学習アルゴリズムの効率的かつ拡張可能な実装を可能にすること。
• 異なるネットワーク環境とワークロード下での、さまざまなアルゴリズムとシステム緩和のトレードオフを実証的に評価すること。

提案手法

• 低レベルの集合的演算を抽象化するモジュラーでMPIスタイルの通信ライブラリの設計により、多様な通信パターンをサポートすること。
• 3つのコアなシステム最適化の統合：計算・通信の重ね合わせ（O）、テンソルの統合とフラットニング（F）、階層的GPU通信（H）。
• 複数のシステム緩和技術のサポート：低精度勾配（QSGD、1-bit Adam）、分散型学習（Decen）、非同期実行（Async）。
• システムレベルの最適化に直接対応するプリミティブを使用して、ユーザーがアルゴリズムを組み合わせて拡張できる統一フレームワークの提供。
• 最適化スタックを介した自動パフォーマンスチューニングの実装により、モデルおよびネットワーク特性に動的に適応すること。

実験結果

リサーチクエスチョン

• RQ1 アルゴリズム固有の論理をハードコードせずに、幅広いシステム緩和技術をネイティブにサポートできる通信ライブラリを設計できるか？
• RQ2 量子化、分散化、非同期性などの異なるシステム緩和が、多様なワークロードとネットワーク環境下でエンドツーエンドの学習パフォーマンスにどのように影響を与えるか？
• RQ3 オーバーラップ、統合、階層的通信といった主要なシステム最適化が、全体の学習効率に与える相対的な影響は何か？
• RQ4 特定のモデルとネットワーク環境において、最良のパフォーマンスを達成するアルゴリズム設定は何か？

主な発見

• BAGUAは、VGG16、BERT、Transformer、LSTM+AlexNetのワークロードにおいて、PyTorch-DDP、Horovod、BytePSを最大2倍速く、エンドツーエンドの学習時間を短縮する。
• 低帯域幅のネットワークでは、QSGD や 1-bit Adam などの圧縮アルゴリズムが通信オーバーヘッドを大幅に低減し、パフォーマンスを向上させる。
• 高レイテンシーネットワークでは、同期のボトルネックが減少するため、分散型アルゴリズム（Decen-32bits/8bits）が集中型のものよりも優れる。
• アブレーションスタディの結果、オーバーラップ、統合、階層的通信の3つの最適化がすべて不可欠であることが確認され、ワークロードによってその影響が異なる：Hは通信集約的タスクに有効、Fは多数の小さなテンソルを持つモデルに有利、Oは計算集約的シナリオで最も効果的である。
• ストレグラーが存在する環境では、非同期学習（Async）によりエポック時間が30–50％短縮され、非同期クラスタにおける有効性が裏付けられる。
• 実証的ガイドラインが確立された：SGDベースの最適化手法にはQSGDを、Adamには1-bit Adamを、通信対計算比が低い環境では非同期手法を用いること。