DeMo大幅降低训练通信量，让大模型高效跨数据中心训练。
DeMo: Decoupled Momentum Optimization
- 分离本地动量更新，用快速正交变换+顶K稀疏化压缩梯度
- 每步通信量减少85倍，10亿参数模型训练精度不降
- 无需改架构，适配多机多卡、跨数据中心场景
大规模神经网络训练依赖同步数据并行，但全精度梯度归约带来严重通信瓶颈。我们提出解耦动量优化（DeMo），作为任意基于动量优化器的即插即用替代方案，显著降低通信带宽且保持收敛性。DeMo通过三步实现：(i) 分离本地动量更新，(ii) 使用快速正交变换（如DCT）后进行顶K稀疏化，(iii) 利用动量缓冲区作为误差反馈进行动量减法。该设计使每步通信量减少达两个数量级，计算开销极低。在3亿和10亿参数的DeMo语言模型上实验表明，DeMo相比AdamW-DDP每GPU通信量减少最高85倍，同时达到相近损失与准确率。DeMo与拓扑无关，支持跨多数据中心或以太网环境训练。代码已开源于https://github.com/bloc97/DeMo。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。