一种新优化器让大模型训练更快更省内存。
Smoothing DiLoCo with Primal Averaging for Faster Training of LLMs
- 用平滑迭代平均替代复杂梯度聚合，降低内存占用。
- 在不同规模模型上比AdamW快8.7%至10.1%，图像任务提速达25.5%。
- 适合追求训练效率的大模型研究者和工业部署场景。
我们提出广义原始平均（GPA），这是Nesterov方法的扩展，在非分布式设置下统一并泛化了近期基于平均的优化器，如单工作节点的DiLoCo和Schedule-Free。传统DiLoCo依赖内存密集的双循环结构，周期性使用Nesterov动量聚合伪梯度，而GPA通过解耦Nesterov插值常数，实现每步平滑迭代平均，消除该复杂性。结构上，GPA类似Schedule-Free，但将均匀平均替换为指数移动平均。实验表明，GPA在保持更低内存开销的同时，持续优于单工作节点DiLoCo和AdamW。在达到目标验证损失时，对Llama-160M、1B和8B模型分别实现8.71%、10.13%和9.58%的加速。在ImageNet ViT任务中，小批量和大批量设置下分别提速7%和25.5%。此外，我们证明：对于任意基优化器具有$O(\sqrt{T})$遗憾的场景，其中$T$为迭代次数，GPA的收敛保证可匹配或超越原算法，具体取决于插值常数。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。