DeMo: 分離されたモーメンタム最適化

DeMo: Decoupled Momentum Optimization

November 29, 2024

著者: Bowen Peng, Jeffrey Quesnelle, Diederik P. Kingma

cs.AI

要旨

大規模なニューラルネットワークのトレーニングには、専用の高速インターコネクトを介してアクセラレータ間で勾配を共有する必要があります。周波数分解とエネルギー圧縮の信号処理原則に基づいて、トレーニング中に完全なオプティマイザーの状態とモデルパラメーターを同期させる必要はないことを示します。モーメンタムの更新を分離し、アクセラレータ間でオプティマイザーの状態に制御された発散を許可することで、最新のオプティマイザーよりも改善された収束を達成します。私たちは、{De}coupled {Mo}mentum（DeMo）という融合オプティマイザーおよびデータ並列アルゴリズムを導入し、インターアクセラレータ間の通信要件を数桁削減します。これにより、限られたネットワーク帯域幅や異種ハードウェアでも大規模なニューラルネットワークのトレーニングが可能となります。私たちの手法はトポロジに依存せず、アーキテクチャに依存せず、計算およびメモリのオーバーヘッドが無視できるスケーラブルなクロック同期分散トレーニングをサポートします。実証結果では、DeMoでトレーニングされたモデルがAdamWでトレーニングされた同等のモデルと同等以上の性能を発揮することが示されており、大規模な基盤モデルの事前トレーニング時に高速インターコネクトが不要となります。GitHubにオープンソースの参照PyTorch実装が公開されています：https://github.com/bloc97/DeMo

English

Training large neural networks typically requires sharing gradients between accelerators through specialized high-speed interconnects. Drawing from the signal processing principles of frequency decomposition and energy compaction, we demonstrate that synchronizing full optimizer states and model parameters during training is unnecessary. By decoupling momentum updates and allowing controlled divergence in optimizer states across accelerators, we achieve improved convergence compared to state-of-the-art optimizers. We introduce {De}coupled {Mo}mentum (DeMo), a fused optimizer and data parallel algorithm that reduces inter-accelerator communication requirements by several orders of magnitude. This enables training of large neural networks even with limited network bandwidth and heterogeneous hardware. Our method is topology-agnostic and architecture-independent and supports scalable clock-synchronous distributed training with negligible compute and memory overhead. Empirical results show that models trained with DeMo match or exceed the performance of equivalent models trained with AdamW, while eliminating the need for high-speed interconnects when pre-training large scale foundation models. An open source reference PyTorch implementation is published on GitHub at https://github.com/bloc97/DeMo