DeMo: Entkoppelte Momentum-Optimierung

DeMo: Decoupled Momentum Optimization

November 29, 2024

Autoren: Bowen Peng, Jeffrey Quesnelle, Diederik P. Kingma

cs.AI

Zusammenfassung

Das Training großer neuronaler Netzwerke erfordert in der Regel den Austausch von Gradienten zwischen Beschleunigern über spezialisierte Hochgeschwindigkeitsverbindungen. Unter Bezugnahme auf die Signalverarbeitungsprinzipien der Frequenzzerlegung und Energiekompaktierung zeigen wir, dass eine Synchronisierung vollständiger Optimiererzustände und Modellparameter während des Trainings nicht erforderlich ist. Durch die Entkopplung von Momentum-Updates und die Zulassung kontrollierter Divergenz in den Optimiererzuständen zwischen den Beschleunigern erzielen wir eine verbesserte Konvergenz im Vergleich zu modernsten Optimierern. Wir stellen {De}coupled {Mo}mentum (DeMo) vor, einen fusionierten Optimierer und datenparallelen Algorithmus, der die Anforderungen an die Kommunikation zwischen den Beschleunigern um mehrere Größenordnungen reduziert. Dies ermöglicht das Training großer neuronaler Netzwerke auch bei begrenzter Netzwerkbandbreite und heterogener Hardware. Unsere Methode ist topologieagnostisch und architekturunabhängig und unterstützt skalierbares, taktsynchrones verteiltes Training mit vernachlässigbarem Rechenaufwand und Speicheroverhead. Empirische Ergebnisse zeigen, dass mit DeMo trainierte Modelle die Leistung äquivalenter Modelle, die mit AdamW trainiert wurden, erreichen oder übertreffen, während die Notwendigkeit für Hochgeschwindigkeitsverbindungen beim Vor-Training großer Grundlagenmodelle entfällt. Eine Open-Source-Referenzimplementierung in PyTorch wurde auf GitHub veröffentlicht unter 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