DeMo: Decoupled Momentum Optimization
作者: Bowen Peng, Jeffrey Quesnelle, Diederik P. Kingma
分类: cs.LG, cs.AI
发布日期: 2024-11-29
🔗 代码/项目: GITHUB
💡 一句话要点
DeMo：解耦动量优化，降低大规模模型训练的通信开销
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 分布式训练 解耦动量优化 数据并行 大规模模型 通信优化
📋 核心要点
- 现有大规模神经网络训练依赖高速互连同步梯度，成本高昂且受限于硬件条件。
- DeMo通过解耦动量更新，允许优化器状态在不同加速器上适度发散，降低通信需求。
- 实验表明，DeMo在性能上与AdamW相当或更优，同时显著降低了对高速互连的需求。
📝 摘要（中文）
本文提出了一种名为解耦动量优化(DeMo)的优化器和数据并行算法，旨在减少加速器之间训练大型神经网络时的通信需求。DeMo借鉴了信号处理中频率分解和能量压缩的原理，证明了在训练过程中同步完整的优化器状态和模型参数是不必要的。通过解耦动量更新并允许优化器状态在加速器之间进行受控发散，DeMo实现了比现有优化器更好的收敛效果。DeMo显著降低了加速器间的通信需求，使得即使在有限的网络带宽和异构硬件条件下也能训练大型神经网络。该方法与拓扑结构和硬件架构无关，并支持可扩展的时钟同步分布式训练，且计算和内存开销可忽略不计。实验结果表明，使用DeMo训练的模型在性能上与使用AdamW训练的等效模型相匹配或超过，同时消除了预训练大规模基础模型时对高速互连的需求。开源PyTorch实现已发布在GitHub上。
🔬 方法详解
问题定义：现有分布式训练方法，如AdamW，需要频繁地在各个加速器之间同步梯度和优化器状态，这导致了巨大的通信开销，尤其是在训练大规模模型时。高速互连虽然可以缓解这个问题，但成本高昂且并非所有硬件环境都具备。因此，如何在保证模型性能的前提下，降低分布式训练的通信需求，是一个重要的挑战。
核心思路：DeMo的核心思路是解耦动量更新。传统优化器将动量信息紧密地耦合在优化器状态中，需要在每次迭代时同步。DeMo通过允许优化器状态在不同加速器上适度发散，减少了同步的频率和数据量。这种解耦基于信号处理的原理，认为同步所有信息是不必要的，只需要同步关键信息即可。
技术框架：DeMo是一个融合的优化器和数据并行算法。其整体流程如下：1) 将数据划分到不同的加速器上；2) 每个加速器独立计算梯度和更新模型参数；3) 解耦动量更新，允许优化器状态在一定程度上发散；4) 周期性地同步关键的优化器状态信息。这个过程在时钟同步的分布式环境中进行，无需复杂的异步通信机制。
关键创新：DeMo最重要的技术创新点在于解耦动量更新。与传统的同步梯度方法相比，DeMo允许优化器状态在不同加速器上存在差异，从而减少了通信需求。这种解耦是通过控制发散程度来实现的，以保证模型的收敛性能。
关键设计：DeMo的关键设计包括：1) 如何解耦动量更新，具体来说，就是如何控制优化器状态的发散程度。这可能涉及到一些超参数的设置，例如发散的阈值或频率。2) 如何选择需要同步的关键优化器状态信息。并非所有信息都需要同步，只需要同步对模型性能影响最大的信息即可。3) 如何保证在解耦动量更新的情况下，模型的收敛性能。这可能需要对损失函数或优化算法进行一些调整。
📊 实验亮点
DeMo在实验中表现出色，在保证模型性能的前提下，显著降低了通信开销。具体来说，使用DeMo训练的模型在性能上与使用AdamW训练的等效模型相匹配或超过，同时消除了预训练大规模基础模型时对高速互连的需求。这表明DeMo在降低通信成本方面具有显著优势。
🎯 应用场景
DeMo适用于大规模神经网络的分布式训练，尤其是在网络带宽有限或使用异构硬件的环境中。它可以降低训练成本，加速模型开发周期，并使得在资源受限的条件下训练大型模型成为可能。例如，在边缘设备上进行联邦学习，或在没有高速互连的集群上训练基础模型。
📄 摘要（原文）
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 {\textbf{De}}coupled {\textbf{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