[論文レビュー] DiLoCo: Distributed Low-Communication Training of Language Models

Arthur Douillard, Qixuan Feng+7 more|arXiv (Cornell University)|Nov 14, 2023

Topic Modeling被引用数 7

ひとこと要約

DiLoCoは、大規模言語モデルの分散学習フレームワークを提案する。AdamWを用いた大規模な内部最適化ステップ（500）と外部のネステロフモーメンタム更新を組み合わせることで、通信量を500倍削減しながら、完全同期学習と同等の性能を達成する。C4データセット上での実験では、データ分布、ワーカーの可用性の変化、異種デバイスに対しても頑健であることが示された。

ABSTRACT

Large language models (LLM) have become a critical component in many applications of machine learning. However, standard approaches to training LLM require a large number of tightly interconnected accelerators, with devices exchanging gradients and other intermediate states at each optimization step. While it is difficult to build and maintain a single computing cluster hosting many accelerators, it might be easier to find several computing clusters each hosting a smaller number of devices. In this work, we propose a distributed optimization algorithm, Distributed Low-Communication (DiLoCo), that enables training of language models on islands of devices that are poorly connected. The approach is a variant of federated averaging, where the number of inner steps is large, the inner optimizer is AdamW, and the outer optimizer is Nesterov momentum. On the widely used C4 dataset, we show that DiLoCo on 8 workers performs as well as fully synchronous optimization while communicating 500 times less. DiLoCo exhibits great robustness to the data distribution of each worker. It is also robust to resources becoming unavailable over time, and vice versa, it can seamlessly leverage resources that become available during training.

研究の動機と目的

• アクセラレータが同じ場所に配置されておらず、通信帯域が限られている状況での大規模言語モデルの学習の課題に対処すること。
• 密に同期された高帯域のクラスタを必要とする完全同期分散学習のスケーラビリティと信頼性の問題を克服すること。
• 相互接続性が低い複数の隔離されたクラスタ（「島」）間での学習を可能にすること。
• データ分布の変化、ワーカーの障害、計算リソースの動的な可用性変化に対する耐性を高めること。
• 通信オーバーヘッドを最小限に抑えつつ、完全同期学習と同等のモデル性能と効率を維持すること。

提案手法

• 通信頻度を低減するため、フェデレーテッド・アveragingに類似したフレームワークを採用し、大きな内部ステップ数（H=500）を用いる。
• 各ワーカーのローカルモデル更新に、標準のLLM学習実務に合致するAdamWを内部最適化子として採用する。
• 全ワーカーからの勾配を集約した後、グローバルモデルパラメータを更新するためにネステロフモーメンタムを外部最適化子として適用する。
• モデル更新を局所的学習フェーズと稀なグローバル同期に分離し、ノード間通信を1回/500ステップにまで削減する。
• 異なるハードウェアタイプや場所に配置されたワーカーを許容し、定期的なパラメータ同期のみを実施する。
• 各外部最適化ステップの後、更新されたグローバルモデルを全ワーカーに再配布し、反復的改善を可能にする。

実験結果

リサーチクエスチョン

• RQ1 通信頻度を著しく減らした分散学習手法が、大規模言語モデル学習において完全同期学習と同等の性能を達成できるか？
• RQ2 異なるワーカーシェアにおけるデータ分布のばらつきに対して、この手法はどれほど頑健か？
• RQ3 障害や新規リソースの参加といったワーカーの可用性の動的な変化に対しても、この手法は適応可能か？
• RQ4 8ワーカーを超えてスケーリングできるか？通信削減の限界は何か？
• RQ5 クラスタ間で非一様または異種のデバイスを用いても、このアプローチは学習の安定性と収束性を維持できるか？

主な発見

• C4データセット上では、通信量を500倍削減したにもかかわらず、DiLoCoは完全同期学習と同等またはそれを上回るパープレキシティ性能を達成した。
• ワーカーシェア間での多様なデータ分布に対しても、DiLoCoは安定した学習と収束性を維持した。
• DiLoCoは動的なワーカー可用性にうまく対応でき、新規ワーカーの統合や障害からの回復をスムーズに行える。
• 低帯域の接続を持つ地理的に分散したクラスタ間でも学習が可能で、通信頻度を数分にまで低下させ、数百ミリ秒ごとの同期から大幅に改善した。
• DiLoCoは完全同期モデルと同等のモデル効率と推論速度を維持しており、推論時におけるパフォーマンスペナルティが一切ない。
• DiLoCoは8ワーカーを超えても効果が薄れる傾向を示すが、異種で緩く接続された環境においても、強力なウォールクロック時間効率と頑健性を示した。