Smoothing DiLoCo with Primal Averaging for Faster Training of LLMs

TL;DR

This paper introduces Generalized Primal Averaging (GPA), a new optimizer that accelerates training of large language models and vision transformers by unifying and improving upon existing averaging-based methods, reducing memory use and increasing speed.

Contribution

GPA extends Nesterov's method to unify recent averaging optimizers, simplifies implementation, and provides theoretical convergence guarantees while improving empirical training speed.

Findings

GPA outperforms DiLoCo and AdamW in training speed for LLMs.

GPA achieves significant speedups on ImageNet ViT workloads.

GPA maintains convergence guarantees for base optimizers with $O(\sqrt{T})$ regret.

Abstract

We propose Generalized Primal Averaging (GPA), an extension of Nesterov's method that unifies and generalizes recent averaging-based optimizers like single-worker DiLoCo and Schedule-Free, within a non-distributed setting. While DiLoCo relies on a memory-intensive two-loop structure to periodically aggregate pseudo-gradients using Nesterov momentum, GPA eliminates this complexity by decoupling Nesterov's interpolation constants to enable smooth iterate averaging at every step. Structurally, GPA resembles Schedule-Free but replaces uniform averaging with exponential moving averaging. Empirically, GPA consistently outperforms single-worker DiLoCo and AdamW with reduced memory overhead. GPA achieves speedups of 8.71%, 10.13%, and 9.58% over the AdamW baseline in terms of steps to reach target validation loss for Llama-160M, 1B, and 8B models, respectively. Similarly, on the ImageNet ViT workload, GPA achieves speedups of 7% and 25.5% in the small and large batch settings respectively. Furthermore, we prove that for any base optimizer with $O(\sqrt{T})$ regret, where $T$ is the number of iterations, GPA matches or exceeds the original convergence guarantees depending on the interpolation constants.