Smoothing DiLoCo with Primal Averaging for Faster Training of LLMs

本文提出了一种名为广义原始平均(GPA)的新优化器,它通过解耦 Nesterov 动量插值常数来简化 DiLoCo 的双循环结构,在降低内存开销的同时,显著提升了 LLM 及 ViT 模型的训练收敛速度并保持了理论收敛保证。

Aaron Defazio, Konstantin Mishchenko, Parameswaran Raman, Hao-Jun Michael Shi, Lin Xiao

发布于 2026-03-02
📖 1 分钟阅读☕ 轻松阅读

Each language version is independently generated for its own context, not a direct translation.

这篇论文介绍了一种名为 GPA (Generalized Primal Averaging,广义原对偶平均) 的新优化算法。它的目标是让训练大型人工智能模型(如 LLM)变得更快、更省内存,而且更稳定。

为了让你轻松理解,我们可以把训练 AI 模型想象成教一个学生(模型)通过做练习题(数据)来掌握一门复杂的技能

1. 背景:现在的“老师”是怎么教学生的?

在训练 AI 时,我们需要一个“优化器”(Optimizer),它就像老师,负责告诉学生:“你刚才做错了,下次往这个方向改一点。”

  • AdamW(目前的行业标准): 就像一位循规蹈矩的老师。他每次只根据学生上一道题的反馈,小心翼翼地调整一步。这种方法很稳,但有时候走得太慢,或者在复杂的迷宫里容易绕圈子。
  • DiLoCo(之前的改进版): 就像一位搞“特训营”的老师。他让学生先自己埋头做很多道题(比如 32 道),然后老师再根据这 32 道题的总结果,给学生做一个大的总结和调整。
    • 优点: 这种“先积累后爆发”的方法确实比 AdamW 快。
    • 缺点: 这种方法很笨重。老师需要同时盯着两套笔记(一套是学生的草稿,一套是老师的总结),而且必须等学生做完整整 32 道题才能给一次反馈。这就像学生做了一半题,老师却还在等,导致信息传递是“断断续续”的,不够流畅。

2. GPA 的核心理念:让“特训”变得“平滑”

GPA 的作者发现,DiLoCo 之所以有效,是因为它利用了**动量(Momentum)平均(Averaging)**的概念。但 DiLoCo 把这两个概念绑死在了“做 32 道题”这个死板的规则上。

GPA 做了什么?
GPA 就像一位更聪明的导师。他保留了“特训营”的精髓,但去掉了那些繁琐的规矩。

  • 比喻:从“断奏”到“连奏”
    • DiLoCo 像是在弹钢琴时,每弹 32 个音符就停下来,深吸一口气,然后突然用力弹下一个和弦。节奏是跳跃的(断断续续)。
    • GPA 则像是把这种跳跃变成了平滑的滑音。它不再等待学生做完 32 道题才给反馈,而是每一步都在微调。它把“做很多题”和“给反馈”这两个动作解开了,让它们可以独立控制。

3. GPA 的三个超能力

① 拆掉“双层结构”,变轻变快

DiLoCo 需要维护两套模型状态(就像老师手里要拿两本厚厚的笔记),这非常占内存。

  • GPA 的做法: 它只需要一本笔记。它通过一种巧妙的数学技巧(把两个控制参数分开),在每一步都平滑地更新模型,不需要额外存那么多数据。
  • 结果: 训练大模型时,内存占用更少,机器能跑得更快。

② 像“老练的向导”一样平滑

DiLoCo 在更新模型时,信息是“咔嚓”一下跳过去的,这会导致训练曲线忽高忽低(不稳定)。

  • GPA 的做法: 它引入了指数移动平均(Exponential Moving Average)。想象一下,学生每走一步,老师不仅看刚才的一步,还温柔地回顾一下之前的路,把方向平滑地融合进去。
  • 结果: 训练过程非常平稳,不容易“翻车”(发散),而且能使用更大的学习率(走得更快)。

③ 速度更快,效果更稳

论文在多个模型上做了测试(从 1.6 亿参数的小模型到 80 亿参数的大模型,还有图像识别任务):

  • 速度提升: 相比传统的 AdamW,GPA 能节省 8% 到 25% 的训练时间(或者说用更少的步骤达到同样的效果)。
  • 稳定性: 在训练过程中,GPA 的曲线比 DiLoCo 更平滑,不像 DiLoCo 那样随着“内循环步数”增加而变得难以控制。

4. 总结:GPA 到底好在哪里?

如果把训练 AI 比作开车

  • AdamW新手司机:每走一步都小心翼翼,开得稳但慢。
  • DiLoCo赛车手:偶尔猛踩油门冲一段,然后急刹车调整方向。虽然快,但操作复杂,容易失控,而且需要很多额外的装备(内存)。
  • GPAF1 赛车手:它继承了赛车手“猛冲”的直觉,但把“急刹车”变成了流畅的过弯。它不需要额外的装备,操作更简单,而且全程保持高速且平稳。

一句话总结:
GPA 通过一种聪明的数学方法,把原本笨重、断断续续的“特训模式”(DiLoCo)变成了平滑、高效、省内存的“连续加速模式”,让训练大模型变得更快、更简单。

在收件箱中获取类似论文

根据您的兴趣定制的每日或每周摘要。Gist或技术摘要,使用您的语言。

试用 Digest →