arXiv:2505.23725cs.LG2025-05被引 15

用Muon优化器提升DiLoCo训练效率,支持更大规模分布式训练。

MuLoCo: Muon is a practical inner optimizer for DiLoCo

  • 以Muon为内层优化器改进DiLoCo,改善梯度方向准确性。
  • 在1.76B到3.1B模型上,多卡训练时性能优于传统DiLoCo。
  • 兼容量化和长同步间隔,适合大规模语言模型训练场景。

DiLoCo是一种强大的大语言模型训练框架,可在网络受限条件下实现更大的最优批大小和更高的加速器利用率。然而,已有研究发现其性能随工作节点数(K)增加而下降。本文指出,内层优化器的选择是影响DiLoCo行为的关键但常被忽视因素,它决定了外层优化器使用的伪梯度质量。鉴于Muon在数据并行训练中表现优于AdamW,我们研究其归一化更新步长对伪梯度质量的影响。实验表明,随着工作节点数增加,Muon生成的伪梯度方向更准确。我们在150M、416M、914M、1.76B和3.1B模型上对DiLoCo、MuLoCo、AdamW DP和Muon DP进行了大规模超参数调优。结果一致显示:当K≥1时,MuLoCo绝对性能优于DiLoCo;当K>2时,相对数据并行基线也更优。该方法兼容量化、流式处理和长同步间隔。在K=1时,MuLoCo甚至超越数据并行标准。我们进一步将最优超参数外推至15B规模,在K=1和K=16下训练六组模型。结果显示,K=16的MuLoCo接近单工作节点性能,而K=1的MuLoCo以1600万令牌大批次达到最佳基线表现。

原文摘要 · Abstract (English)

DiLoCo is a powerful framework for training large language models (LLMs), enabling larger optimal batch sizes and increased accelerator utilization under networking constraints. However, DiLoCo's performance has been shown to degrade as the number of workers (K) increases (Charles et al., 2025). In this work, we posit that a related but often overlooked factor in DiLoCo's behavior is the choice of inner optimizer, which shapes the pseudogradient used by the outer optimizer. Given the recent success of Muon relative to AdamW for data parallel (DP) training, we examine how Muon's normalized optimizer steps can affect the pseudogradient's quality. We find that, relative to AdamW, Muon yields more directionally correct pseudogradients as the number of workers ($K$) increases. In our experiments pre-training language models, we conduct extensive hyperparameter tuning across 150M, 416M, 914M, 1.76B, and 3.1B models for DiLoCo, MuLoCo, AdamW DP, and Muon DP. Consistently across all scales, we find that with $K\geq1$ workers, MuLoCo (Muon inner optimizer DiLoCo) achieves superior performance to DiLoCo in absolute terms and for $K>2$ it outperforms DiLoCo relative to their data parallel baselines, while being compatible with quantization, streaming, and long synchronization intervals. At $K=1$, we find that MuLoCo can even outperform the data-parallel gold standard while having larger critical batch sizes. Finally, we extrapolate optimal hyperparameters to 15B scale and train a model with each method (six in total) using $K=1$ and $K=16$ workers. We find that $K=16$ MuLoCo nearly matches single-worker performance at this scale, while MuLoCo $K=1$ matches the best performing baseline while using a much larger $16$M token batch size.

大模型训练分布式优化高效训练

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。