Nesterov Method for Asynchronous Pipeline Parallel Optimization

TL;DR

This work tackles gradient staleness in asynchronous pipeline parallel optimization by introducing a Nesterov Accelerated Gradient variant with delay correction in the weight space. The method modifies the look-ahead in NAG to discount gradient contributions by $(1-\gamma_t)$ and evaluates gradients at delayed look-ahead points, yielding sublinear convergence for convex, smooth objectives with fixed delay. Empirically, it substantially outperforms existing asynchronous methods and even surpasses synchronous GPipe on decoders-only language models up to 1B parameters, including decentralized SWARM training; a memory-efficient no-weight-stash variant also shows competitive performance. The results establish asynchronous optimization as a viable option for large-scale language modeling and provide a practical, convergent approach for 100% pipeline utilization in both centralized and decentralized settings.

Abstract

Pipeline Parallelism (PP) enables large neural network training on small, interconnected devices by splitting the model into multiple stages. To maximize pipeline utilization, asynchronous optimization is appealing as it offers 100% pipeline utilization by construction. However, it is inherently challenging as the weights and gradients are no longer synchronized, leading to stale (or delayed) gradients. To alleviate this, we introduce a variant of Nesterov Accelerated Gradient (NAG) for asynchronous optimization in PP. Specifically, we modify the look-ahead step in NAG to effectively address the staleness in gradients. We theoretically prove that our approach converges at a sublinear rate in the presence of fixed delay in gradients. Our experiments on large-scale language modelling tasks using decoder-only architectures with up to 1B parameters, demonstrate that our approach significantly outperforms existing asynchronous methods, even surpassing the synchronous baseline.

Figures (13)

• Figure 1: Original NAG (left) and our modified version (right) for delayed gradients (denoted with $\bar{{\mathbf{g}}}_t$). Our method discounts the gradient term by $(1-\gamma_t)$. When $\gamma_t\to 1$, the angle $\alpha\to 0$, making the weight trajectory smoother. Consequently, the look-ahead ${\mathbf{d}}_t$ can be shown to act as delay correction, alleviating gradient staleness.
• Figure 2: Training trajectory comparison on three language modelling datasets. In all scenarios, our method significantly outperforms the asynchronous methods while surpassing the synchronous GPipe method throughout training. Our memory efficient version clearly outperforms the asynchronous methods while being competitive to GPipe in two out of three datasets.
• Figure 3: Training and validation trajectory for the 1B parameter model. Similar to the base model, our method outperforms GPipe while the memory efficient version is competitive with GPipe.
• Figure 4: Comparison with other delay correction methods on WT. Our method, outperforms all other delay correction methods in terms of training loss and weight discrepancy. Additionally, NAG improves all previous delay correction methods, while NAG alone yields the best performance.
• Figure 5: Performance with respect to the number of stages. Even though, performance slightly degrades for our method compared to GPipe, the training time increase is exponentially larger for GPipe.

Theorems & Definitions (8)

• Proposition 1
• proof
• Theorem 1
• proof
• Proposition 2
• proof
• Theorem 2
• proof