DeMo: Ongelinkte Momentumoptimalisatie

DeMo: Decoupled Momentum Optimization

November 29, 2024

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

cs.AI

Samenvatting

Het trainen van grote neurale netwerken vereist doorgaans het delen van gradiënten tussen versnellers via gespecialiseerde hoogwaardige interconnects. Geïnspireerd door de signaalverwerkingsprincipes van frequentiedecompositie en energiecompactie, tonen we aan dat het synchroniseren van volledige optimizerstaten en modelparameters tijdens training overbodig is. Door het loskoppelen van momentumupdates en het toestaan van gecontroleerde divergentie in optimizerstaten over versnellers heen, bereiken we een verbeterde convergentie in vergelijking met state-of-the-art optimalisatoren. We introduceren {De}coupled {Mo}mentum (DeMo), een samengevoegde optimizer en datageparallellisatie-algoritme dat de communicatievereisten tussen versnellers met meerdere ordes van grootte vermindert. Dit maakt training van grote neurale netwerken mogelijk, zelfs met beperkte netwerkbandbreedte en heterogene hardware. Onze methode is topologie-agnostisch en architectuur-onafhankelijk en ondersteunt schaalbare klok-synchrone gedistribueerde training met verwaarloosbare reken- en geheugenoverhead. Empirische resultaten tonen aan dat modellen getraind met DeMo overeenkomen met of zelfs beter presteren dan equivalente modellen getraind met AdamW, terwijl de noodzaak voor hoogwaardige interconnects bij het vooraf trainen van grootschalige basismodellen wordt geëlimineerd. Een open source referentie-implementatie in PyTorch is gepubliceerd op GitHub op 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