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

arXiv:2411.19870v2 Announce Type: replace 摘要：神経ネットワークのトレーニングのスケールアップは、同期データ並列化に依存度が高まりつつありますが、完全精度の勾配全-AllReduceは深刻な通信ブottleneckを招きます。当論文では、通信帯域幅を劇的に削減 tout converge の特性を維持する、どのモーメンタムベースの最適化器にも即応として導入可能な「Decoupled Momentum Optimization (DeMo)」を提案します。DeMo は (i) 局所的なモーメンタム更新を分断し、(ii) 高速正直交変換（例：DCT）を適用した後に top-k スパース化を行い、(iii) モーメンタムバッファをモーメンタム減法を通じて誤差フィードバックとして再利用する、この 3 つの要素から構成されています。この設計により、単一ステップあたりの通信量は最大で 2 桁削減でき、計算オーバーヘッドは最小限です。3 億パラメータおよび 10 億パラメータの DeMo LLM での実験は、DeMo が AdamW-DDP に比べて各 GPU ごとに最大 85 倍もデータ転送量を削減しつつ、同等の損失と精度を達成することを示しています。DeMo はトポロジーに依存せず、複数データセンター間、またはイーthern 基の環境でのトレーニングを可能にします。コードは https://github.com/bloc97/DeMo で利用可能です。

arXiv:2411.19870v2 Announce Type: replace Abstract: Scaling neural network training increasingly depends on synchronous data-parallelism, yet full-precision gradient all-reduce imposes a severe communication bottleneck. We propose Decoupled Momentum Optimization (DeMo), a drop-in replacement for any momentum-based optimizers that significantly reduces the communication bandwidth while maintaining convergence. DeMo (i) decouples local momentum updates, (ii) applies a fast orthonormal transform (e.g., DCT) followed by top-k sparsification, and (iii) reuses the momentum buffer as error feedback via momentum subtraction. This design reduces per-step communication by up to two orders of magnitude with minimal computational overhead. Experiments on 300M and 1B-parameter DeMo language models show DeMo transmits up to 85x less data per GPU than AdamW-DDP while achieving comparable loss and accuracy. DeMo is topology-agnostic and enables training across multi-datacenter or Ethernet-based setups. Code is available at https://github.com/bloc97/DeMo