2024/11/26 by Shaoqing Yu, Yu, Shuhua, Ding Zhou +11 · 1 citation
Computer Science · #FOS: Computer and information sciences #Machine Learning (cs.LG) #Neural Networks and Applications
paper · pdf · doi:10.48550/arxiv.2411.17866
openalex publication_date 2024/11/26 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28
Pre-training Transformer models is resource-intensive, and recent studies have shown that sign momentum is an efficient technique for training large-scale deep learning models, particularly Transformers. However, its application in distributed training remains underexplored. This paper investigates a novel communication-efficient distributed sign momentum method with multiple local steps, to cope with the scenarios where communicating at every step is prohibitive. Our proposed method allows for a broad class of base optimizers for local steps, and uses sign momentum in the global step, where momentum is generated from differences accumulated during local steps. For generic base optimizers, by approximating the sign operator with a randomized version that acts as a continuous analog in expectation, we present a general convergence analysis, which specializes to an O(1/√(T)) rate for a particular instance. When local step is stochastic gradient descent, we show an optimal O(1/T1/4) rate in terms of ℓ1 gradient norm for nonconvex smooth cost functions. We extensively evaluate our method on the pre-training of various sized GPT-2 models from scratch, and the empirical results show significant improvement compared to other distributed methods with multiple local steps.