Nesterov Method for Asynchronous Pipeline Parallel Optimization
This addresses the challenge of maximizing pipeline utilization for large neural network training on interconnected devices, representing an incremental improvement over existing asynchronous methods.
The paper tackles the problem of stale gradients in asynchronous pipeline parallel optimization by introducing a variant of Nesterov Accelerated Gradient, proving sublinear convergence with fixed delay and showing it outperforms existing methods and even synchronous baselines in large-scale language modeling tasks with up to 1B parameters.
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.