DeMo: Decoupled Momentum Optimization
Authors: name: Bowen Peng; name: Jeffrey Quesnelle; name: Diederik P. Kingma
Published: 2024-11-29 17:31:47

Abstract: 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