NoLoCo: 大規模モデル向けのNo-all-reduce低通信トレーニング手法

要旨

大規模言語モデルのトレーニングは、一般的に、高帯域幅の相互接続を介して通信する数万のアクセラレータを含むクラスタ上での最適化手法によって行われます。これらのクラスタをスケールアップすることは高コストであり、実用的でなくなる可能性があり、トレーニング可能なモデルのサイズに制限を課します。最近のいくつかの研究では、高度に接続された計算クラスタを必要としない、通信量の少ないトレーニング手法が提案されています。これらの最先端の低通信トレーニング手法では、依然としてモデルパラメータの同期ステップが採用されていますが、すべてのモデルレプリカに対して行われる場合、低帯域幅ネットワーク上ではコストがかかることがあります。本研究では、トレーニング中にすべてのモデルパラメータを明示的に同期せず、その結果、集団通信を一切必要としない新しい最適化手法、NoLoCoを提案します。NoLoCoは、Nesterovモメンタムオプティマイザの新しい変種を介して、ランダムに選択された他のモデルと部分的に平均化することで、モデルの重みを暗黙的に同期します。提案するオプティマイザの理論的な収束解析と、言語モデルトレーニングからの実証結果を提供します。 NoLoCoを、125Mから6.8Bパラメータまでの広範なアクセラレータ数とモデルサイズでベンチマークしました。我々の手法は、完全にシャード化されたデータ並列トレーニングや、広く使用されている低通信トレーニング手法であるDiLoCoよりも、大幅に少ない通信オーバーヘッドを必要とします。同期ステップ自体は、数百のアクセラレータがインターネット上でトレーニングする際にDiLoCoで使用されるall-reduceよりも一桁高速であると推定されます。また、アクセラレータのアイドル時間を削減するグローバルなブロッキング通信もありません。DiLoCoと比較して、広範なモデルサイズとアクセラレータ数において、最大4%の高速な収束率も観察されています。

English

Training large language models is generally done via optimization methods on clusters containing tens of thousands of accelerators, communicating over a high-bandwidth interconnect. Scaling up these clusters is expensive and can become impractical, imposing limits on the size of models that can be trained. Several recent studies have proposed training methods that are less communication intensive, avoiding the need for a highly connected compute cluster. These state-of-the-art low communication training methods still employ a synchronization step for model parameters, which, when performed over all model replicas, can become costly on a low-bandwidth network. In this work, we propose a novel optimization method, NoLoCo, that does not explicitly synchronize all model parameters during training and, as a result, does not require any collective communication. NoLoCo implicitly synchronizes model weights via a novel variant of the Nesterov momentum optimizer by partially averaging model weights with a randomly selected other one. We provide both a theoretical convergence analysis for our proposed optimizer as well as empirical results from language model training. We benchmark NoLoCo on a wide range of accelerator counts and model sizes, between 125M to 6.8B parameters. Our method requires significantly less communication overhead than fully sharded data parallel training or even widely used low communication training method, DiLoCo. The synchronization step itself is estimated to be one magnitude faster than the all-reduce used in DiLoCo for few hundred accelerators training over the internet. We also do not have any global blocking communication that reduces accelerator idling time. Compared to DiLoCo, we also observe up to 4% faster convergence rate with wide range of model sizes and accelerator counts.

NoLoCo: 大規模モデル向けのNo-all-reduce低通信トレーニング手法

NoLoCo: No-all-reduce Low Communication Training Method for Large Models

要旨

Support