言語モデリングのための非同期Local-SGDトレーニング

要旨

局所確率的勾配降下法（Local-SGD）、別名フェデレーテッド・アベレージングは、各デバイスが通信ごとに複数のSGD更新を実行する分散最適化のアプローチです。本研究では、言語モデルの訓練における非同期型Local-SGDの実証的研究を提示します。つまり、各ワーカーはSGDステップを完了次第、即座にグローバルパラメータを更新します。ワーカーのハードウェアの異質性、モデルサイズ、ワーカー数、オプティマイザが学習性能に与える影響を包括的に調査しました。素朴な実装では、非同期Local-SGDは同期型と比較して、グローバルモデルパラメータをより頻繁に更新するにもかかわらず、収束までにより多くの反復を要することがわかりました。ワーカーの勾配が古くなった際のグローバルパラメータに対するモーメンタム加速が主要な課題であることを特定しました。我々は、遅延ネステロフ・モーメンタム更新を活用し、ワーカーの計算速度に基づいてローカル訓練ステップを調整する新規手法を提案します。このアプローチは、C4データセット上で最大1億5000万パラメータのモデルを用いて評価され、更新ステップあたりのパープレキシティにおいて同期型Local-SGDと同等の性能を示し、実時間においてはそれを大幅に上回りました。

English

Local stochastic gradient descent (Local-SGD), also referred to as federated averaging, is an approach to distributed optimization where each device performs more than one SGD update per communication. This work presents an empirical study of {\it asynchronous} Local-SGD for training language models; that is, each worker updates the global parameters as soon as it has finished its SGD steps. We conduct a comprehensive investigation by examining how worker hardware heterogeneity, model size, number of workers, and optimizer could impact the learning performance. We find that with naive implementations, asynchronous Local-SGD takes more iterations to converge than its synchronous counterpart despite updating the (global) model parameters more frequently. We identify momentum acceleration on the global parameters when worker gradients are stale as a key challenge. We propose a novel method that utilizes a delayed Nesterov momentum update and adjusts the workers' local training steps based on their computation speed. This approach, evaluated with models up to 150M parameters on the C4 dataset, matches the performance of synchronous Local-SGD in terms of perplexity per update step, and significantly surpasses it in terms of wall clock time.

言語モデリングのための非同期Local-SGDトレーニング

Asynchronous Local-SGD Training for Language Modeling

要旨

Support