论文

余弦门控 Adam-Decay:解耦 DiLoCo 的直接陈旧感知外部优化

Cosine-Gated Adam-Decay: Drop-In Staleness-Aware Outer Optimization for Decoupled DiLoCo

AI 基础设施分布式训练基础设施

摘要

异步 DiLoCo 系统可能会接收较早几轮外轮计算的伪梯度,但标准 Nesterov 外优化器不会明确根据每次更新年龄来调整其更新。这可能会使外部动量缓冲器在大的受控延迟下变得脆弱。我们提出了 Cosine Gated Adam Decay (CGAD),一个简单的、插入式的、年龄感知的外部优化器,它在进入 Adam 的第一时刻和第二时刻缓冲区之前将每个传入的伪梯度缩放 $σ(τ) = γ(τ) e^{-ατ}$ ;指数模型信息衰减和余弦门 $γ(τ)$ 平滑地将超过选定截止值的贡献归零。 CGAD 在 $τ=0$ 时简化为普通 Adam,添加两个超参数,其默认值跨尺度传输,并通过每个片段年龄感知变体 (PA-CGAD) 扩展到部分同步调度程序。对于平滑非凸目标的理想化门控自适应更新,我们证明了一个非渐进收敛界限,其陈旧偏差项仅取决于 $α$,而不是实现的最大延迟 $τ_{\max}$;异步动量 SGD 的标准分析带有 $τ_{\max}^2$ 因子。根据经验,在 25M、1B 和 7B 参数下的 Llama 风格语言模型预训练中,CGAD 在我们扫描的受控延迟范围内稳定地进行训练。余弦截止值起到了规模保险的作用:最接近的基线 Adam Decay(没有截止值的 CGAD)在 25M 时具有竞争力,但其在 $τ=8$ 时的种子到种子 $σ$ 从 25M 增长到 7B,增长了 27 倍,将其单次风险(平均值 + $σ$)推至机会水平损失之上,而 CGAD 的损失则远低于机会水平损失。已发布的 Nesterov 配方是全面扫描中最不稳定的方法。