通过原问题平均平滑 DiLoCo 以加速大语言模型训练
机器学习
2026-03-02 v3 人工智能
机器学习
摘要
我们提出了广义原问题平均 (GPA)，这是 Nesterov 方法的扩展，它在非分布式设置中统一并推广了最近的基于平均的优化器，如单工作节点 DiLoCo 和 Schedule-Free。虽然 DiLoCo 依赖于内存密集的双循环结构来周期性地使用 Nesterov 动量聚合伪梯度，但 GPA 通过解耦 Nesterov 的插值常数来消除此复杂性，从而在每个步骤实现平滑的迭代平均。在结构上，GPA 类似于 Schedule-Free，但用指数移动平均取代了均匀平均。经验上，GPA 一致优于单工作节点 DiLoCo 和 AdamW，且内存开销更低。GPA 相对于 AdamW 基线在达到目标验证损失所需的步数上实现了 8.71%、10.13% 和 9.58% 的加速，分别针对 Llama-160M、1B 和 8B 模型。同样，在 ImageNet ViT 工作负载上，GPA 在小批量和大批量设置下分别实现了 7% 和 25.5% 的加速。此外，我们证明对于任何具有 遗憾的基线优化器，其中 为迭代次数，GPA 根据插值常数匹配或超过原始收敛保证。
引用
@article{arxiv.2512.17131,
title = {Smoothing DiLoCo with Primal Averaging for Faster Training of LLMs},
author = {Aaron Defazio and Konstantin Mishchenko and Parameswaran Raman and Hao-Jun Michael Shi and Lin Xiao},
journal= {arXiv preprint arXiv:2512.17131},
year = {2026}
}