🤖 AI Summary
This work addresses momentum instability in asynchronous DiLoCo systems, where the outer-loop optimizer receives pseudo-gradients delayed by multiple rounds. Standard Nesterov momentum fails under large delays because it ignores gradient age. To resolve this, the authors propose CGAD, a plug-and-play, age-aware optimizer that introduces delay awareness into outer-loop optimization for the first time. CGAD weights delayed gradients via exponential decay and cosine gating, then incorporates these into Adam’s first- and second-moment estimates. Theoretically, its bias depends only on the decay coefficient, not the maximum delay. The method provides default hyperparameters transferable across model scales and a variant, PA-CGAD, suited for partially synchronous schedules. Experiments on Llama-style models (25M–7B parameters) demonstrate that CGAD enables stable training across diverse delays, significantly outperforming Adam Decay and Nesterov baselines, with failure probability far below random chance.
📝 Abstract
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.