[論文レビュー] Decentralized Training of Foundation Models in Heterogeneous Environments

Binhang Yuan, Yongjun He+7 more|arXiv (Cornell University)|Jun 2, 2022

Stochastic Gradient Optimization Techniques被引用数 23

ひとこと要約

この論文は、地理的に分散した異種GPU上でGPT風ファウンデーションモデルのトレーニング用のスケジューラとシステム最適化を提案し、集中型ベースラインより大幅なスピードアップを達成します。

ABSTRACT

Training foundation models, such as GPT-3 and PaLM, can be extremely expensive, often involving tens of thousands of GPUs running continuously for months. These models are typically trained in specialized clusters featuring fast, homogeneous interconnects and using carefully designed software systems that support both data parallelism and model/pipeline parallelism. Such dedicated clusters can be costly and difficult to obtain. Can we instead leverage the much greater amount of decentralized, heterogeneous, and lower-bandwidth interconnected compute? Previous works examining the heterogeneous, decentralized setting focus on relatively small models that can be trained in a purely data parallel manner. State-of-the-art schemes for model parallel foundation model training, such as Megatron, only consider the homogeneous data center setting. In this paper, we present the first study of training large foundation models with model parallelism in a decentralized regime over a heterogeneous network. Our key technical contribution is a scheduling algorithm that allocates different computational "tasklets" in the training of foundation models to a group of decentralized GPU devices connected by a slow heterogeneous network. We provide a formal cost model and further propose an efficient evolutionary algorithm to find the optimal allocation strategy. We conduct extensive experiments that represent different scenarios for learning over geo-distributed devices simulated using real-world network measurements. In the most extreme case, across 8 different cities spanning 3 continents, our approach is 4.8X faster than prior state-of-the-art training systems (Megatron).

研究の動機と目的

• 高価なデータセンタークラスターへの依存を減らすために、大規模ファウンデーションモデルの分散・異種ネットワークトレーニングを動機づける。
• 異種設定におけるパイプラインおよびデータ並列性のためのタスクレット-to-デバイス割り当て問題を定式化する。
• 分散トレーニングにおける通信コストを最小化するためのコストモデルと最適化アルゴリズムを開発する。
• 計算と通信を重畳させてスループットを向上させるシステムレベルの最適化を実証する。
• 現実のネットワーク測定を模した地理分散GPUクラスタでアプローチを評価する。

提案手法

• トレーニングを、デバイスに割り当てられたタスクレット（マイクロバッチ、レイヤー段階）の集合としてモデル化する。
• 2段階のコストモデルを用いて通信コストをデータ並列とパイプライン並列の成分に分解する。
• 粗化されたグラフとオープンループ巡回セールスマン問題を用いてDataP-CostとPipelineP-Costを定式化する。
• 特殊化された局所探索を備えたハイブリッド遺伝アルゴリズムを用いてほぼ最適な割り当てを見つける。
• CUDAストリーム全体で通信と計算をパイプライン化するシステムレベルの最適化を導入する。

実験結果

リサーチクエスチョン

• RQ1 異種ネットワークで通信コストを最小化するために分散GPUタスクレットをどのように割り当てるか？
• RQ2 異種性下でデータ並列およびパイプライン並列通信コストをどのようにモデル化するか？
• RQ3 進化的スケジューリング手法がランダムや Kernighan-Lin 戦略を上回る割り当てを見つけられるか？
• RQ4 地理分散設定におけるエンドツーエンドのトレーニングスループットに対するスケジューラとシステム最適化の影響は何か？
• RQ5 分散トレーニングはパフォーマンスでデータセンターのベースラインにどれだけ近づけるか？

主な発見

• 提案されたスケジューラとシステム最適化は、世界規模の地理分散設定で最先端の集中型システムよりエンドツーエンドのトレーニングを3.8〜4.8x高速化する。
• 世界規模の地理分散シナリオでは、GPT3-1.3BトレーニングでMegatron-1.3Bより4.8x、Deepspeedより3.6x高速である。
• スケジューラなしでは性能向上が低下し、アブレーションでは最大で2.7x遅くなることが示される。
• 分散アプローチはデータセンター条件でMegatron/Deepspeedより1.7–3.5x遅い程度であり、ネットワークは最大100x遅くても。
• 通信と計算の重ね合わせを可能にするシステム最適化はスループットをさらに向上させ、スケジューラはランダム割り当てに対する速度向上に大きく寄与する。
• アブレーション研究はコストモデルに合わせた局所探索が、標準的な Kernighan-Lin ベースの探索より優れることを示す。