FlexDeMo: Decoupled Momentum Optimization for Hybrid Sharded Data Parallel Training

要約

大規模なニューラルネットワークモデルのトレーニングには、多くの場合、いくつかのノードとアクセラレータに分布する広範な計算リソースが必要です。
最近の発見は、勾配の速い移動コンポーネントのみを交換し、局所的に勢いを蓄積するのに十分である可能性があることを示唆しています（分離された運動量、またはデモ）。
ただし、単一のアクセラレータに適合しない大きなモデルを考慮する場合、勾配情報の交換とデモの統合を再考する必要があります。
ここでは、ハイブリッドシャードデータパラレルトレーニング戦略であるFlexDemoを採用することを提案します。これにより、異なる加速器間でノードモデルモデルパラメーターを完全にシャードモデルパラメーターを使用しますが、ノード間通信帯域幅要件は、完全な勾配ではなく高速変動コンポーネントのみを同期することで削減されます。
これにより、以前のハイブリッドシェード戦略とデカップレッドモーメントの利点を効果的に組み合わせます。
私たちの実験結果は、FlexDemoがADAMWを採用したハイブリッドシャードデータ並列トレーニングと検証損失の観点から完全な勾配同期と同等であり、その生存率を示していることを示しています。
さらに、FlexDemoは、ノード間の完全な勾配同期と比較して、トレーニング速度の改善を達成します。
帯域幅が制約された2ノードのセットアップでは、FlexDemoは、完全なグラデーション同期を備えたハイブリッドシャードデータパラレルトレーニングよりも速く、望ましいレベルの検証損失に到達することができます。

要約(オリジナル)

Training large neural network models requires extensive computational resources, often distributed across several nodes and accelerators. Recent findings suggest that it may be sufficient to only exchange the fast moving components of the gradients, while accumulating momentum locally (Decoupled Momentum, or DeMo). However, when considering larger models that do not fit on a single accelerator, the exchange of gradient information and the integration of DeMo needs to be reconsidered. Here, we propose employing a hybrid sharded data parallel training strategy, FlexDeMo, whereby nodes fully shard model parameters locally between different accelerators, while inter-node communication bandwidth requirements are reduced by synchronizing only fast-moving components instead of the full gradients. This effectively combines previous hybrid sharded strategies with the advantages of decoupled momentum. Our experimental results show that FlexDeMo is on par with hybrid sharded data parallel training employing AdamW and full gradient synchronization in terms of validation loss, demonstrating its viability. Furthermore, FlexDeMo achieves improved training speed compared to full gradient synchronization across nodes. In a bandwidth-constrained 2-node setup, FlexDeMo allows reaching desired levels of validation loss faster than hybrid sharded data parallel training with full gradient synchronization.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google