Smoothing DiLoCo with Primal Averaging for Faster Training of LLMs

2025/12/18 by Aaron Defazio, Defazio, Aaron, Konstantin Mishchenko +7
Computer Science · Decision Sciences · #Advanced Bandit Algorithms Research #Advanced Neural Network Applications #Artificial Intelligence (cs.AI) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Stochastic Gradient Optimization Techniques

paper · doi:10.48550/arxiv.2512.17131

openalex publication_date 2025/12/18 · openalex created_date 2025/12/23 · openalex updated_date 2026/07/28

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(√(T)) regret, where T is the number of iterations, GPA matches or exceeds the original convergence guarantees depending on the interpolation constants.

Citations

• SNOO: Step-K Nesterov Outer Optimizer - The Surprising Effectiveness of Nesterov Momentum Applied to Pseudo-Gradients
• Understanding Outer Optimizers in Local SGD: Learning Rates, Momentum, and Acceleration
• Purifying Shampoo: Investigating Shampoo's Heuristics by Decomposing its Preconditioner
• MuLoCo: Muon is a practical inner optimizer for DiLoCo
• Communication-Efficient Language Model Training Scales Reliably and Robustly: Scaling Laws for DiLoCo
• Training Deep Learning Models with Norm-Constrained LMOs
• Streaming DiLoCo with overlapping communication: Towards a Distributed Free Lunch
• Modular Duality in Deep Learning
• The AdEMAMix Optimizer: Better, Faster, Older
• The Road Less Scheduled
• Scalable Optimization in the Modular Norm
• Asynchronous Local-SGD Training for Language Modeling
• DiLoCo: Distributed Low-Communication Training of Language Models
• Optimal Linear Decay Learning Rate Schedules and Further Refinements
• A Distributed Data-Parallel PyTorch Implementation of the Distributed Shampoo Optimizer for Training Neural Networks At-Scale
• GPT-4 Technical Report
• Training trajectories, mini-batch losses and the curious role of the learning rate
• Training Compute-Optimal Large Language Models
• Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer
• Momentum via Primal Averaging: Theoretical Insights and Learning Rate Schedules for Non-Convex Optimization
• LaProp: Separating Momentum and Adaptivity in Adam
• Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer
• Lookahead Optimizer: k steps forward, 1 step back
• On the Curved Geometry of Accelerated Optimization
• On First-Order Meta-Learning Algorithms
• Shampoo: Preconditioned Stochastic Tensor Optimization
• Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks
• Acceleration of Stochastic Approximation by Averaging
• A Stochastic Approximation Method

Related

• When, Where and Why to Average Weights?
• Stochastic Weight Averaging in Parallel: Large-Batch Training that Generalizes Well
• Averaging Weights Leads to Wider Optima and Better Generalization
• Parallel Restarted SGD with Faster Convergence and Less Communication: Demystifying Why Model Averaging Works for Deep Learning
• Trainable Weight Averaging: Accelerating Training and Improving Generalization