提出一种自适应延迟感知的优化器,提升异步训练稳定性。
Cosine-Gated Adam-Decay: Drop-In Staleness-Aware Outer Optimization for Decoupled DiLoCo
- 用余弦门控衰减机制动态调整旧梯度权重
- 在25M~7B参数模型上延迟控制下训练更稳定
- 适合大规模异步分布式训练场景
异步DiLoCo系统可能接收数个外层迭代前计算的伪梯度,但标准Nesterov外层优化器未显式考虑每条更新的年龄。这会使外层动量缓冲区在大延迟下变得脆弱。我们提出余弦门控Adam衰减(CGAD),一种简单、可直接替换的延迟感知外层优化器,在伪梯度进入Adam的一阶和二阶矩缓冲前,按σ(τ) = γ(τ) e^{-ατ}缩放;指数项建模信息衰减,余弦门γ(τ)平滑截断超过设定阈值的贡献。CGAD在τ=0时退化为普通Adam,新增两个超参数,其默认值可跨规模迁移,并通过片段级延迟感知变体(PA-CGAD)扩展至部分同步调度。对于光滑非凸目标的理想化门控自适应更新,我们证明了非渐近收敛界,其延迟偏差项仅依赖α而非实际最大延迟τ_max;而标准异步动量SGD分析中该因子为τ_max²。实验显示,在25M、1B和7B参数的Llama风格语言模型预训练中,CGAD在所测延迟范围内均稳定训练。余弦截断起到规模保险作用:最接近的基线(无截断的Adam Decay)在25M时表现尚可,但在7B时τ=8下的种子间标准差σ增长27倍,导致单次运行风险(均值+σ)超过随机水平损失,而CGAD始终低于该阈值。发布的Nesterov方案在整个测试中稳定性最差。
原文摘要 · Abstract (English)
Asynchronous DiLoCo systems may receive pseudo-gradients computed several outer rounds earlier, yet the standard Nesterov outer optimizer does not explicitly condition its update on per-update age. This can make the outer momentum buffer brittle under large controlled delays. We propose Cosine Gated Adam Decay (CGAD), a simple, drop-in, age-aware outer optimizer that scales each incoming pseudo-gradient by $σ(τ) = γ(τ) e^{-ατ}$ before it enters Adam's first- and second-moment buffers; the exponential models information decay and the cosine gate $γ(τ)$ smoothly zeroes contributions past a chosen cutoff. CGAD reduces to plain Adam at $τ=0$, adds two hyperparameters whose defaults transfer across scales, and extends to partial-sync schedulers via a per-fragment age-aware variant (PA-CGAD). For an idealized gated-adaptive update on smooth non convex objectives, we prove a non-asymptotic convergence bound whose staleness-bias term depends on $α$ alone, rather than on the realized maximum delay $τ_{\max}$; standard analyses of asynchronous momentum-SGD instead carry a $τ_{\max}^2$ factor. Empirically, on Llama style language model pretraining at 25M, 1B, and 7B parameters, CGAD trains stably across the controlled delays we sweep. The cosine cutoff acts as scale insurance: the closest baseline, Adam Decay (CGAD without the cutoff), is competitive at 25M but its seed-to-seed $σ$ at $τ=8$ grows 27x from 25M to 7B, pushing its single-shot risk (mean + $σ$) above the chance-level loss while CGAD's stays well below. The published Nesterov recipe is the least stable method on the full sweep.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。