Nesterov Method for Asynchronous Pipeline Parallel Optimization
Authors: Thalaiyasingam Ajanthan
Sameera Ramasinghe
Yan Zuo
Gil Avraham
Alexander Long
Venue: ICML 2025 poster (2025)
Keywords: Asynchronous Optimization
Pipeline Parallelism
Nesterov Method
Convergence Analysis
Decentralized Training
Protocol Learning
Primary area: optimization->large_scale_parallel_and_distributed
TL;DR: We introduce a variant of Nesterov accelerated gradient method to effectively address gradient staleness in asynchronous pipeline parallel optimization, demonstrating feasibility in large-scale language model training.

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.

Nesterov Method for Asynchronous Pipeline Parallel Optimization
Thalaiyasingam Ajanthan'
Sameera Ramasinghe"
Yan Zuo'
Avraham Alexander Long'
Abstract
The main objective of PP methods is mask the communi- cation overhead and improve device utilization to min- imize training time. To this end, many pipeline scheduling strategies (i.e., the order of processing forward and back- ward passes of microbatches) have been developed (Huang et al., 2019; Narayanan et al., 2021b; Qi et al., 2023). How- ever, the main bottleneck in these methods is the require- ment to synchronize the weights and gradients across stages at each update step, hindering 100% pipeline utilization.
Pipeline Parallelism (PP) enables large neural net- work training on small, interconnected devices by splitting the model into multiple stages. To maximize pipeline utilization, asynchronous opti- mization is appealing it offers 100% pipeline utilization by construction. However, it is inher- ently challenging as the weights and gradients are no longer synchronized, Icading to stale (or de- layed) 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 effec- tively address the staleness in gradients. We the- oretically prove that our approach converges at a sublinear rate in the presence of fixed delay in gradients. Our experiments on large-scale lan- guage modelling tasks using decoder-only archi- tectures with up 1B parameters, demonstrate that approach significantly outperforms exist- ing asynchronous methods, even surpassing the synchronous baseline.'
To enable full pipeline utilization, we consider asynchronous optimization in PP, where the weight updates are performed independently at each stage without waiting for the corre- sponding backward pass to complete. This approach intro- duces gradient staleness, as weights are updated multiple times between the forward and backward passes of a micro- batch, resulting in outdated gradients being used for weight updates. This gradient staleness presents a significant opti- mization challenge, necessitating sophisticated delay correc- tion mechanisms to ensure convergence, even in traditional distributed settings (Agarwal & Duchi, 2011; 1; Zheng et al., 2017; Stich & Karimireddy, 2019; Mishchenko et al., 2022).
Typically, delay correction methods directly estimate the gradients via forecasting approaches (Zheng et al., 2017; Liu et al., 2024b) or extrapolate the previous update step in the weight space (Hakimi et al., 2019; Guan et al., 2019). We follow the latter approach as delay correction in the weight space does not make any assumptions on the loss landscape (as in gradient forecasting methods), but rather assumes that the update directions change slowly. Noting that the smoothness of update directions can be controlled in momentum based optimizers, we derive a look-ahead based optimization method to alleviate gradient staleness.
1. Introduction
Pipeline Parallelism (PP) is a standard parallelism technique for training large-scale frontier foundational models (Dubey et al., 2024; Liu al., 2024a). PP partitions neural net- works into sequential stages, enabling the training of models that exceed the memory capacity of any single device by distributing computations across multiple interconnected devices. The devices can be co-located in datacenters or connected via the internet (i.e., low bandwidth connections) in a fully decentralised setting. Each device processes a stage (i.e., a set of consecutive layers) and communicates the activations and gradients with adjacent stages often via interconnects.
Specifically, we introduce a variant of the Nesterov Accel- erated Gradient (NAG) method (Nesterov, 1983; 2013) for asynchronous PP optimization. Our idea stems from the observation that NAG has look-ahead step which can be repurposed as a delay correction in the weight space, by carefully modifying the update formula (refer to Fig. 1). Our approach is a simple. yet elegant modification to NAG that does not introduce any hyperparameters. We theoreti- cally prove that our approach converges at a sublinear rate for convex, smooth functions with fixed delay in gradients.
bandwidth-constrall
Pluralis Research. Correspondence to: Thalaiyasingam Ajan- than <aj @lluralis.ai.>
Proceedings of the 4gnd International Conference on Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
'Our code is available https://githun com/ Plurali sesseerchhhhhhcoee AhsypppP.
We demonstrate the merits of our approach on large-scale

---

Nesterov Method for Asynchronous Pipeline Parallel Optimization
language modelling tasks with decoder-only transformer architectures (Vaswani et al., 2017; Karpathy, 2022). In short, our approach significantly outperforms all existing asynchronous PP methods including sophisticated delay correction mechanisms, even surpassing the synchronous baseline. Notably, for the first time, we train a 1B param- eter model to convergence the asynchronous PP setup outperforming the synchronous baseline. Moreover, we show the effectiveness of our method in a realistic decen- tralized training framework, namely SWARM (Ryabinin et al., 2023), where our approach significantly outperforms both synchronous and asynchronous methods. Our experi- ments clearly demonstrate the feasibility of asynchronous PP optimization in the large-scale setting.
stage. Formally, the gradient with respect to the weights and the error signal the previous stage can be written as:
V(((w)) h,,(w, ei - [ e(),
where h; corresponds to the chain rule. Then at each stage, the computed gradients V[((w() used to perform the optimization step with learning rate 7) >> 0:
w w 7 V f((())
(4)
We omit optimizer specific updates and stochasticity for simplified notation.
Note that, in the asynchronous setting, the weights w/ and gradients of (and f.()(()) are not synchronized. In other words, the weights are updated without waiting for the back- ward passes (i.e., gradient computation) of all active micro- batches to complete, leading to delayed gradients. There- fore, asynchrony would affect both the backpropagation step, Eq. (3), and the optimization step, Eq. (4).
Our contributions can be summarized as follows:
For the first time, we show an asynchronous PP optimiza- tion method can surpass the synchronous alternative for large-scale language modelling tasks.
Our approach is an elegant variant of NAG, and pro- vide theoretical and empirical justification of convergence in the presence of gradient staleness.
Furthermore, we test our method in a realistic decentral- ized training framework (SWARM), proving its empirical effectiveness beyond doubt.
2.2. PipeDream
PipeDream (Narayanan et al., 2019) uses One-Forward- One-Backward (1F1B) schedule with weight stashing for asynchronous PP optimization. At steady state, each stage alternates between forward and backward passes. Since weight stashing retains a copy of weights used in the for- ward pass until the microbatch completes, correct backprop- agation can be computed by loading the stashed weights, This ensures synchronous backpropagation. Eq. (3), while weight updates, Eq. (4), remain asynchronous, i.e., delayed gradients are used to update the most recent weights.
2. Preliminaries
We first define the problem setup and briefly review PipeDream (Narayanan et al., 2019) and NAG (Nesterov, 2013; Bubeck et al., 2015), upon which we build our work. We refer the interested reader the respective papers for more details.
2.1. Problem Setup
Precisely, let T, be the delay at stage i, meaning the weights at stage i are updated т, times between the forward and backward passes of particular microbatch. Assuming constant delay each stage, this delay can be written as:
Let us consider single pipeline without data parallelism for simplicity. Let P be the number of pipeline stages. Let f((w; x;-1) (w,,;;1)) be the forward function in stage i with weights w; (correspond to all learnalee parameters) and input x;-1, respectively. We may simply write ft for brevity, when the context is clear. Now, the full neural network forward function can be written as:
2(P - 1) +1
T
2K
where i E 2 P} and K is the update interval, e.g., K = 1 if updated for every microbatch. Note that earlier stages incur larger delays.
F(W, xo) = fp SP- 0.......o fx()),,
where W = (wp...., }1} and xy is the input data point. Analogously the backward function can be written as:
Now the backward function, and optimizer step at time for PieeDream can be written as:
G(W,ep) = 91 o g2 o ..... gp(ep).
Ti (w( T211) w! (w:
(6)
where ep is the error signal, and gi( (w:, e;) is the backward function at stage i corresponding to ft. Specifically, is the gradient of the loss with respect to the output activations x; = fw,,, (w,,x;11), x;-1), and it is backpropagated through the
To reiterate, the forward and the backward passes are com- puted on the weights at time step t - T., and these outdated gradients are used update the latest weights time step
network using the weights (and non- -ineariitiii of each
2

---

Nesterov Method for Asynchronous Pipeline Parallel Optimization
In short, PipeDream uses more memory in each stage to store old weights, ensuring correct gradient computation, but the optimization is performed asynchronously without any delay correction.
0
70
g8
d
W++1
d
W(+1
Wr
Additionally, due to asynchrony, PipeDream does not guar- antee that weights are synchronized across stages. Specifi- cally, since the delay is stage dependent, earlier stages will use older weights compared to the later stages. Precisely, the forward function at time step & takes the following form:
W
We-1
Figure 1: Original NAG (left) and our modified version (right) for delayed gradients (denoted with Bz). Our method discounts the gradient term by (1-7). When "n + 1, the angle a - 0, making the weight trajectory smoother. Consequently, the look-ahead di can be shown to act as delay correction, alleviating gradient staleness.
F(W xo) = l30 f ....o (X-P+((x)))
where / = J.(w J.( The gradients are computed for this function. Note, for t 2 P. all stages are updated at each time step following the IFIB schedule. Thus, the set of weights W can be interpreted as stage dependent initialization, and in practice, this weight discrepancy has not shown to cause any convergence issues (Narayanan et al., 2019).
the ideal one. This is appealing as it does not make any as- sumptions about the loss function or gradients as in gradient forecasting methods. The only assumption is that the update directions change slowly with respect to iterations, which is valid for momentum based optimizers.
As discussed previously, NAG is an ideal candidate for this as performs a look-ahead step by extrapolating the last update step. We now introduce our modified update formula for delay correction in the weight space.
2.3. Nesterov Accelerated Gradient
Nesterov Accelerated Gradient (NAG) (Nesterov, 1983; 2013; Bubeck et al., 2015) is an accelerated gradient method that has the optimal (++) convergence rate for smooth con- vex functions in the non-stochastic setting. The main idea is to perform look-ahead step in the previous update di- rection, combined with carefully selected sequence of step-sizes to ensure accelerated convergence.
3.1. Nesterov Method for Delayed Gradients
Let us consider particular stage and drop the stage index for brevity.' Let T be the delay as defined in Eq, (5), which is constant for each stage, and A, be the corresponding delay in the weight space. Then, the delayed versions of w: and d, take the following form:
Let f Rm + R be the objective function. Then, NAG performs the following iterations starting from an initial point W E pm:
W. = W = W; A.,
de = d
(9)
de (WW Wr-1), Wa+1= = W. + d: - -7V + dd),
(8)
Now, our variant of NAG with delayed gradients perform the following iterations:
where 17 > 0 is the learning rate. Here, the momentum coef- ficient 7/ satisfies, 71 = 0, <<< < < 1, and the sequence of Tt is derived as part of the convergence proof (Bubeck et al., 2015). Note that, d, corresponds to the look-ahead step which extrapolates the update (w, - we-1) by Y and the gradients are computed at the extrapolated point (w, + d.).
d: = 7 (w: Wi-1),
(10)
W++1 = w. + d: 77(1- (1(1-%)Vfe V) + +d.)
Note that compared to Eq. (8), our update formula discounts the gradient term by (1 - ).) This is illustrated in Fig. 1. This subtle but important difference, allows us to show that our look-ahead step approximates A, (i.e., acts as delay correction) and our algorithm converges at sublinear rate when the gradients are delayed.
NAG has been incorporated into popular deep learning optimizers such as SGD (Sutskever et al., 2013) and Adam (Dozat, 2016), although it often slightly underper- forms in synchronous settings. However, we show the supe- riority of a variant of NAG in asynchronous optimization.
Look-ahead as delay correction. The momentum coef- ficient is usually chosen to be a constant close to 1 or chosen as an increasing sequence satisfying lim 7t- = 1. Assuming this, we can intuitively see that the influence of the gradient term decreases when %/e increases, due to the discount factor (1 - 7)). This translates into d, dominating the updates, resulting in a smooth trajectory in the weight
3. Method
We address the gradient staleness in asynchronous PP op- timization by incorporating a delay correction mechanism. Specifically, our idea is to do the delay correction in the weight space by extrapolating the last update step, so that the gradients can be computed at a point that is closer to
We also use a subscript for the time step to reduce clutter.
3

---

Nesterov Method for Asynchronous Pipeline Parallel Optimization
space. Consequently, it can be shown that the vector direc- tions of look-ahead time & t- ie, d, d and the delay A are aligned. This means that taking step in the direction of d, reduces the delay, effectively acting as delay correction in the weight space. We formally state this helow.
Implementation details. For transformer architectures, AdamW optimizer (Loshchilov, 2017) is known to be su- perior over SGD. To this end, we adopt the NAdam opti- mizer (Dozat, 2016) that incorporates the Nesterov method in Adam (Kingma, 2014) with decoupled weight decay for our language model training. Interestingly, both Adam and NAdam discount the gradient term by (1 (1-7)), as they treat momentum as an exponential moving average of gradients. Even though the motivation of Adam is different, in this pa- per, we theoretically and empirically show the effectiveness of this discount factor for asynchronous optimization.
Proposition 1. Let "Ve be increasing sequence such that lim 7/- - 1 then, lim cos (4,,, de) = 1, where cos(..) is the cosine similarity.
Proof. By algebraic manipulations, we can write A, in terms of d, as follows:
In the Py Torch implementation of NAdam (PyTorch Con- tributors, 2025), the momentum coefficient is warmed up to the B1 value passed to the algorithm, which satisfies our as- sumption in Proposition I when 31 set close 1. To this end, we use the NAdam optimizer as is, and use large value for 81 (0.99 in our experiments). It is remarkable that, with almost no modifications to an existing implementation, we show significant improvement over existing asynchronous optimization methods and provide theoretical justification. In fact, as we show in our experiments, simply changing the optimizer to NAdam improves other delay correction meth- ods such as the learning rate discounting approach (Yang et al., 2021) and gradient forecasting methods (Zheng et al., 2017). However, such delay correction methods deteriorate the performance of NAdam as is, validating our insight that correcting delays in the weight space is more effective.
4 = Σg
(+))))))7>) de
(11)
Σ
mT= 11)) (1 1 (-))))) gk
7
k=t-r
where g: = +d.) When, 7t- + L, the gradient term vanishes, and cos( A,, d,) - 1. Refer to Appendix A.2 for the detailed proof.
D
Here, since the last update step aligns with the delay direc- tion, one may wonder what if we extrapolate in the direction of (w, wr-1) to completely compensate for the delay rather than taking a fractional step d, = Ye(,, w.-1). The answer is, convergence may be disrupted if one naively extrapolates further than "/- Note that by definition of NAG, Tt. cannot be larger than 1, and it is usually derived as part of the convergence proof. Nevertheless, our experiments show that our delay correction in the weight space is superior compared to gradient forecasting based correction methods.
Since we build our method using the PipeDream framework (refer Sec. 2.2), at each stage we store copy of weights for each active microbatch ensure correct gradient com- putation. This amounts to storing T; copies of weights in stage i, therefore, the memory requirement of our method grows linearly with the delay (or number of stages). Even though, in practice, these weight stashes can be offloaded to the CPU effectively masking the memory requirement, for completeness, we also discuss a no- of our method below. Despite backpropagating through different sets of weights compared to the forward pass, sur- prisingly, this method is competitive to the synchronous method without any additional memory requirements.
Convergence analysis. We now state our convergence theorem for a convex, 3-smooth function.
Theorem 1. Let f be a convex, B-smooth function with bounded gradients, then the iterates in Eq. (10) with 3 converges at a rate of (+))
version
Proof. We largely follow the proof of (Bubeck et al., 2015) while adopting delayed gradients. As part of the proof, the momentum coefficient is set to 7/- ش and we show that || O(t) using the bounded gradient assumption. Detailed proof can be found in Appendix A.3.
3.2. Memory Efficient Version
As noted before, without weight stashing, the backpropaga- tion is incorrect, i.e., Eq. (3) takes the following form:
D
We do not claim the sublinear rate we derived is tight, and the rate or the constants in the bound may be improved. We leave any such analysis for future work. To our knowledge, this is the first time a variant of the Nesterov method is shown to converge in the presence of delayed gradients.
(( ) ha (w/, of I
(12)
o T. g (w/, ey 11))
Here, since w! is used to backpropagate (instead of w! " the gradients are altered for the following stages. Hence, we denote the delayed error signal by and the weight gradients by f to indicate that they are altered. To compensate for this, we make two modifications: 1) stage-dependent learn-
Note, at time f t-7. r, the weights w.-r are extrapolated by de- and the gradients are computed at the point we-r + de-r, to compensate for the "future" delay A.
4

---

Nesterov Method for Asynchronous Pipeline Parallel Optimization
ing rate, and 2) stage-dependent momentum coefficient.
and tested on small-scale image classification language translation tasks. In contrast, for the first time, we demon- strate that asynchronous methods can surpass synchronous methods in 1B parameter scale language modelling tasks, in addition to providing theoretical convergence guarantees.
Here, the idea is that earlier stages incur larger delays and also larger error accumulation due to incorrect backpropaga- tion. Even if our variant of NAG can alleviate the issue with larger delays, larger error accumulation is still a problem. To this end, we further decrease the learning rate for earlier stages following the idea of (Yang et al., 2021). Addition- ally, the momentum coefficient linearly increased from 0.9 to 0.99 from the last stage the first one. Precisely, the learning rate y and the momentum coefficient ? for stage E 21, P) are set as follows:
5. Experiments
5.1. Experimental Setup
We evaluate our method on the language modelling task using decoder-only architectures. We use three large- scale datasets: WikiText (WT) (Merity et al., 2016), BookCorpus (BC) (Zhu et al., 2015), and Open Tex (OWT) (Gokaslan et al., 2019) datasets. For WikiText, we utilize the predefined training and validation splits, for the other datasets, we randomly select 10% of the training set as the held-out validation set. Our model architecture is based on NanoGPT (Karpathy, 2022) with dropout. The base configuration includes context length of 512, an em- bedding dimension of 768, 12 attention heads, and 8 layers, with approximately 134M parameters, and each layer is treated as a stage in our PP framework. We use the GP2 tokenizer (Radford et al., 2019) and train the model from scratch. Across all experiments, we maintain microbatch size of 8, learning rate nj of 3e-4, and a weight decay of 0.01, unless otherwise specified. The update interval K = 1 for the asynchronous methods. The learning rate is obtained by tuning the performance of GPipe on the Wiki- Text dataset. All baseline methods employ the AdamW optimizer (Loshchilov,, 2017). Each experiment run for 50k iterations, with a linear warmup of 3k iterations starting from a learning rate of 1e-7. Then, it is decayed 3e-5 following a cosine decay schedule.
7 n/ = T
where Pt = 1 - min ():1),
(13)
P-i + 0.09
7 = 0.9 +
P
where P is the number of stages. Note, the learning rate correction is only applied for the first T iterations to stabilize training as in (Yang et al., 2021).
4. Related Work
Asynchronous data parallel methods. Data Parallelism (DP) a traditional distributed training setting, where each device optimizes the full model and periodically synchro- nizes the model parameters. Asynchronous DP methods are well-studied within the theoretical framework and many gradient delay correction mechanisms have been devel- oped (Agarwal& & Duchi, 2011; Stich & Karimireddy, 2019; Assran et al., 2020). Notable methods that improve over the simple asynchronous SGD (Recht et al., 2011) include delay dependent learning rate (Barkai et al., 2019; Mishchenko et al., 2022), gradient forecasting with second-order infor- mation (Zheng et al., 2017), and look-ahead in the weight space (Hakimi et al., 2019). Apart from this, training dy- namics of asynchronous DP methods have also been ana- lyzed (Mitliagkas et al., 2016; Liu et al., 2024b) and some of these observations may be useful in the PP setting as well.
We evaluate two existing asynchronous PP methods, PipeDream (Narayanan et al., 2019) and PipeMare ( (Yang et al., 2021), together with the synchronous GPipe (Huang et al., 2019) method. Unlike PipeDream, which stashes old weights, PipeMare estimates these weights using the velocity of weight updates. Moreover, PipeMare incorpo- rates a learning rate discounting mechanism described in Eq. (13), with T set to 6k iterations. We implement PipeMare within the PipeDream framework. For GPipe, we use the torch. di attribuudd. pipeliiiinn millinn ininn package, setting the number of microbatches to 4, which is limited by the GPU memory. Due to differences in the underlying implementations, the absolute wall-clock time is not compa- rable between methods. We instead compare performance based on training iterations, or equivalently the amount of data processed.
Pipeline parallel methods. The main objective of PP methods (Guan et al., 2024) is to improve pipeline uti- lization which led to the development of many pipeline scheduling strategies including GPipe (Huang et al., 2019), IFIB (Narayanan et al., 2021b), and Zero Bubble (Qi et al., 2023). However, these methods suffer from synchronization bottlenecks. Asynchronous methods alleviate this bottle- neck to achieve 100% pipeline utilization at the cost of gradient staleness. Notable gradient delay correction mech- anisms for PP include weight stashing (Narayanan et al., 2019; 2021a), learning rate discounting (Yang et al., 2021), and direct weight prediction (Chen et al., 2018; Guan et al., 2019). Moreover, gradient forecasting methods developed for DP (Zheng et al., 2017) can also be employed. However, existing asynchronous PP methods are mainly empirical
'In our experiments, GPipe accumulates 4 microbatches for each weight update reduce pipeline and therefore performs 4x less updates for the same amount of data than asynchronous methods.
bubbles (Huang et al., 2019).
5

---

Nesterov Method for Asynchronous Pipeline Parallel Optimization
Geipe
GPipe
Moetream
Rieeore
7.0
Miceroam
Roeearr
7.
Rieearre
Noeerramm
Durs Curs No WS
x1000 ateiiions
x1000 Baralisrn
XIGDC Berationn
WikiText
(b) BookCorpus
(c) OpenWehText
Figure 2: Training trujectory 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.
Method
WT
BC
OWT
Memory
ofiec PpoOmem
GPipe
30.63
42.39
65.17
O(N)
Mooitarc
Our
Ourr
0KP NC-V
PipeDream
99.48
52.98
224.30
O(PN)
PipeMare 71.38
76.93
239.13
O(N)
Ours 27.72
39.85
62.86
O(PN)
x1003
12
x2000 terstion
Ours-No-WS
29.90
42.61
108.20
O(N)
Figure 3: Training and validation trajectories for the IB parameter model. Similar the base model, our method outperforms GPipe while the memory efficient version is competitive with GPipe.
Table 1: Perplexity scores the validation set at 50k itera- tions. The memory requirement of each method is shown in the last column, where N is the number of parameters and P is the number of pipeline stages. Ours outperforms all meth- ods, including the synchronous GPipe method. Our memory efficient version (Ours-No- -o--S) on par with GPipe on per- plexity, while outperforming other asynchronous methods.
methods. Unlike PipeMare, which estimates old weights to facilitate correct backpropagation, our no-weight-stash ver- sion achieves superior performance without such estimation, relying instead on the Nesterov-based delay correction.
Although we primarily report training loss, the trends are consistent with validation loss as well. Validation loss com- parison is provided in Fig. 9 in the appendix. For complete- ness, we include the perplexity scores on the validation set for all methods at 50k iterations in Table 1 along with the memory requirement. In this, we only consider the mem- ory requirement for storing weights and do not consider activation memory as it is the same for all the methods.
Our proposed method is denoted as Ours, which em- ploys the Nadam optimizer (Dozat, 2016) with decoupled weight decay and momentum coefficient 91 of 0.99. The mfforr-eeffcient variant of our method, Ours-No-WS. -SS in- corporates the same learning rate discounting PipeMare (see Eq. (13)), with T also set to 6k iterations. Unless oth- erwise specified, all experiments use the base architecture described above and are performed on the WikiText dataset with the aforementioned hyperparameters. All experiments are performed on a system equipped with A10G GPUs.
5.3. Increasing the Model Size
To show the scalability of our approach, we train a 1B parameter model in the asynchronous setting. We maintain the number of stages at 8 but increase the context length to 1024 and the embedding dimension to 2688, with 24 attention heads. A learning rate of le-4 is used for all methods. These experiments are performed on system equipped with 8 A GPUs.
5.2. Main Results
We analyze the training trajectories of our method against the baselines for the base architecture with stages, il- lustrated in Fig. 2. The results demonstrate that method significantly outperforms existing asynchronous approaches and even surpasses GPipe across all three datasets. Our memory-efficient variant is competitive with GPipe, match- ing its performance on two out of the three datasets. Notably, the training trajectories of PipeDream and PipeMare reveal the optimization challenges inherent in asynchronous se- tups,* while Nesterov-based delay correction effectively bridges the gap between asynchronous and synchronous
As shown Fig. 3, the results are consistent with those of the base model. Specifically, our method significantly out- performs all baselines, including GPipe, throughout training, Notably,
no-weight-stash variant matches the perfor- mance of GPipe. Aside from learning rate adjustment, no changes were made to the method or its hyperparameters. This large-scale experiment demonstrates the merits of our method and the practicality of asynchronous PP optimiza- tion for language model training.
"To the hest of our knowledge, this is the first time PipeDream and PipeMare are tested large-scale language modelling tasks.
6

---

Nesterov Method for Asynchronous Pipeline Parallel Optimization
Riperremm
Pip3rremm
ApeDream-ULt
RocDream Moctroanuu
Pipcrrram-Me (Duns)
Tuitrry
Infoumiadert
UR- Sccoonddörder
Ours
0,003
8,004
Polyrommal PPT
Uhf Secondbroer
CurS
Tairin Loos
UR SecondOrder
PipCOream
Saconec der-AA
0,000
x1000 aterotoon
x1000 baratcons
x3000 Ecrations
(a) Training loss
(b) Weight discrepancy for Stage-1
(c) NAG with other delay corrections
Figure 4: Comparison with other delay correction methods on WikiText. Our method, outperforms all other delay correction methods in terms of training loss and weight discrepancy. Additionally, NAG improves previous delay correction methods, while NAG alone yields the best performance.
5.4. Other Delay Correction Methods
Taining Lops at S0k Itens.
GPipe
Chrrs
We compare our method with the following delay correc- tion mechanisms which originally developed for asyn- chronous optimization in the DP setting.
Thaining Time
2.5
2.0
2.5
3.0
PipeDream-LR. The learning rate discounting method (Yang et al., 2021; Mishchenko et al., 2022)
No
Stages
which employs delay dependent discounting of the learning rate as noted in Eq. (13) where T is set to 6k.
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 exponen- tially larger for Gpipe.
L.R-SecondOrder. On top of the learning rate correction above, we forecast the gradients to the current step using the second-order Taylor expansion of the loss as in (Zheng et al., 2017). The implementation is similar to (Zheng et al., 2017) where the diagonal of the Fisher matrix is used to approximate the Hessian.
5.5. Increasing the Number of Stages
To evaluate scalability with respect the number of stages, increase the number of layers in the base model whilst maintaining the same embedding dimension. We compare the results against GPipe in terms of training loss and the percentage increase in training time. For configurations with 20 and 24 stages, the learning rate reduced to le-4 for our method, while all other hyperparameters remain unchanged. The results are presented in Fig. 5.
Polynomial+FFT. In this approach, we frame gradient forecasting as a time series prediction problem, leveraging historical gradient data to predict future gradients. Specifi- cally, we employ a second-order polynomial to model trends and utilize Fast Fourier Transform (FFT) to capture any periodic signals. The history size is set to 8. This is well- known method in time series forecasting literature (Bloom- field, 2004), which we adopt for gradient delay correction.
Since the delay grows with the number of stages, the perfor- mance of our method degrades slightly as expected. How- ever, due to 100% pipeline utilization, the percentage in- crease in runtime for our approach is significantly lower compared to GPipe. Concretely, for GPipe, 24-stage model takes 8.5x more time compared to the 4-stage model, how- ever, our model is only 2.5x slower. This highlights the trade-off between performance and runtime efficiency, un- derscoring the advantages of asynchronous optimization. This runtime discrepancy is more pronounced when a faster GPU is used, such as A100, as observed in our IB model experiments (refer to Fig. 10 in the appendix).
In addition to monitoring training loss, we also evaluate the Root-Mean- Square Error (RMSE) of the weight discrep- ancy A, (as in Eq. (9)) the first stage, which experiences the largest delay. This metric is named gap in (Hakimi et al., 2019), which directly measures the effectiveness of delay correction, where smaller gap indicates better delay correction. The results are presented in Fig. 4.
Among the above delay correction methods, polynomial fitting is the most effective. However, our simple Nesterov- based approach outperforms all other sophisticated tech- niques in both training loss and weight discrepancy. Notably, our method is complementary to those strategies and further enhances their performance. Although combining Nesterov with other delay correction mechanisms diminishes its effec- tiveness, supporting our hypothesis that correcting delays in the weight space is more impactful than other approaches.
5.6. Ablation Study
To understand the effect of momentum coefficient, we vary the momentum coefficient and report the performance of our base model on the WikiText dataset. In addition to training loss, we also measure the cosine similarity be-

---

Nesterov Method for Asynchronous Pipeline Parallel Optimization
RseBream-NAC NAC-mo.9
NAG-MD.
RocDream NAG a
coood. for Staze-
No WS HAG a 74:-W5- NAN--M SP
Woetheem NAG-mD up (Dtr))
Toining Loos
(40: WS Lt
Pac-00ti-4 E--NAC- (Oum))
pcDrcim-MAGo N4G
petiream WD 99 (Dun)
x1000 aterito
X:00D hanaaionn
x3000 Ecrations
(a) Ablation of our main method
(b) Look-ahead and delay alignment
(c) Ablation of Ours-No-WS
Figure 6: Ablation study of our methods. Constant momentum coefficient of 0.99 performs slightly better than the adaptive version (denoted with '-a' ) and it also shows the best alignment in (b). For the memory efficient version, in addition to adaptive momentum, delay dependent learning rate discounting also helps as shown in (c).
My
RiatOnna RAG
SWARR
10
SWARM-AsynE
Ours IWS
Thayk
PipDDrum--e NAG
Taining Lass
<1000 Bodiar
12
x1000 toeettor
30
6
Figure 7: Our approach with and without the gradient dis- counting term for the Nesterov method. Without the dis- counting term, training is significantly disrupted due to gradient staleness, validating our insight.
10
x1000 Iterations
Figure 8: Training trajectory comparison in SWARM. Our method significantly outperforms both the synchronous and asynchronous versions of SWARM.
tween the look-ahead direction d, and the weight difference A, Wt Wr-r the first stage, where the discrepancy is most pronounced. The results are presented in Fig. 6.
5.7. Realistic Decentralized Training
Finally, stress test our method, we evaluate our approach in a realistic decentralized setting using SWARM (Ryabinin et al., 2023). SWARM is built using the Hivemind frame- work (Ryabinin & Gusev, 2020; Ryabinin et al., 2020) and it supports fault-tolerant, pipeline- -arralllll training with multi- ple worker nodes per stage (i.e., DP at each stage) connected over the internet. Natively, it employs gradient accumula- tion akin to synchronous training and the workers at each stage are synchronized periodically. We refer the interested reader to Ryabinin et al. (2023) for more details.
As expected, increasing the momentum coefficient from 0.9 to 0.99 leads to improved performance. However, using an adaptive momentum coefficient, as defined in Eq. (13), results in slightly worse performance compared to a fixed value of 0.99 for our main method. The influence of the momentum coefficient on the alignment between the look- ahead step and the weight difference is also empirically demonstrated, matching our theoretical insight. On the other hand, for the memory-efficient version, a delay-adaptive momentum coefficient performs better, and incorporating learning rate discounting further enhances performance.
Our architecture and the hyperparameters are similar to the base model experiments but adapted to accommodate SWARM and train on the WikiText dataset. More details can be found in Appendix B.1. We compare three variants: 1) the standard, synchronous setting (SWARM); 2) an asyn- chronous setting with local updates for every microbatch and a periodic stage-wise weight synchronization (SWARM- Asyne); and 3) SWARM-Asyne with our no-weight-stase version (Ours-No- WoW.S). Note that weight stashing is not applicable in SWARM. Results are reported in Fig.
Efffect of gradient discounting. As outlined in Sec. 3.1, the standard implementation of NAdam incorporates dis- count factor of (1 (1-x) >) for the gradient term. To demonstrate the necessity of this discounting factor for our approach, we train a model using a modified optimizer where this term removed, denoted as PipeDream-NAG: -AG--Baa-- The results are presented Fig. 7.
Without this discount factor, the method struggles due to gradient staleness and fails to achieve training loss compa- rable to our method. Notably, the weight discrepancy at the first stage is an order of magnitude larger than that observed with our gradient discounting approach. Empirically, this strongly validates the significance of the discounting factor.
Our method clearly outperforms both the synchronous and asynchronous versions of SWARM throughout training. No- tably, SWARM-Asyne shows training instability even with a lower learning rate (refer to Appendix B. whereas our method shows stable convergence. Validation loss also fol- lows similar trend shown in Fig. 13 in the appendix.
8

---

Nesterov Method for Asynchronous Pipeline Parallel Optimization
6. Conclusion
plexity. Foundations and Trends@ in Machine Learning. 8(3-4):231-357, 2015.
introduce a novel variant of NAG to alleviate gradient staleness in asynchronous PP optimization. Theoretically, we show that our algorithm converges a sublinear rate for convex, smooth functions in the non-stochastic setting with fixed delay in gradients. Practically, adopting our method is as simple as switching the optimizer and changing the value of an existing hyperparameter. To show the merits of our approach, we performed large-scale experiments on lan- guage modelling tasks using models up to 1B parameters. In all our experiments, our method consistently outperformed previous asynchronous methods well as the synchronous GPipe method. The behaviour is consistent even in decen- tralized experiments in SWARM. In the future, we intend to investigate the tightness of our convergence rate and ex- tend our approach to PP with stage-wise DP setting. in a heterogenous decentralized training environment.
Chen, C.-C., Yang, C.-L., and Cheng, H.-Y. Efficient and robust parallel dnn training through model parallelism on multi-gpu platform. arXiv preprint arXiv:1809.02889 2018.
Dozat, T. Incorporating nesterov momentum into adam. ICLR Workshop, 2016.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A., et al. The llama 3 herd of models. arXiv preprint arXiv:2407.22 22..224777.173., 2024.
Gokaslan, A., Cohen, V., Pavlick, E., and Tellex, S. Open- webtext corpus.
http://yyylooooo..oooo io/OpenwesTexttCoppsss/ 2019.
Impact Statement
Guan, L., Yin, W., Li, D., and Lu, X. Xpipe: Efficient pipeline model parallelism for multi-gpu dnn training. arXiv preprint 1rXiv:1911.044 1,64100, 2019.
Our method establishes the feasibility of asynchronous PP optimization for billion-scale language model training. im- proving efficiency in distributed and decentralized settings. By enabling more effective parallelization, it has the poten- tial reduce energy consumption and training costs while increasing accessibility to large-scale model training via decentralized infrastructures. As a theoretical contribution, its societal impact is application-dependent. with no direct consequences attributable to the method itself.
Guan, L., Li, D.-S., Liang, J.-Y., Wang, W.-J, Ge, K.-S., and Lu, X.-C. Advances of pipeline model parallelism for deep learning training: overview. Journal of Computer Science and Technology, 39(3):567-584, 2024.
Hakimi, I., Barkai, S., Gabel, M., and Schuster, A. Taming momentum in a distributed asynchronous environment. arXiv preprint arXiv: 1162,, 2019.
References
Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, D., Chen,
Agarwal, A. and Duchi, J. C. Distributed delayed stochastic optimization. Advances in neural information processing systems, 24, 2011.
M., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., al. Gpipe: Efficient training of giant neural networks using pipeline parallelism. Advances in neural information processing systems, 32, 2019.
Assran, M., Aytekin, A., Feyzmahdavian, H. R., Johansson, M., and Rabbat, M. G. Advances asynchronous parallel and distributed optimization. Proceedings of the IEEE, 2020.
Karpathy, A. NanoGPT.
kar karpat hy http///nnnPp/00000000
Kingma, D. P. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6990, 2014.
Barkai, S., Hakimi, I., and Schuster, A. Gap aware mitigation of gradient staleness. arXiv preprint 1909. 2019.
Liu, A., Feng, B., Xue, B., Wang, B., Wu, B., Lu, C., Zhao, C., Deng, C., Zhang, C., Ruan, C., et al. Deepseek- v3 technical report. arXiv preprint arXiv:2412.1944 19437, 2024a.
Bloomfield, Fourier analysis of time series: an introduc- tion. John Wiley& & Sons, 2004.
Brown, T., Mann, B., Ryder, N., Subbiah, M.. Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
Liu, B., Chhaparia, R., Douillard, A., Kale, S., Rusu, A. A., Shen, J., Szlam, A., and Ranzato, M. Asynchronous local-sgd training for language modeling, arXiv preprint arXiv:2401.09133 2024b.
Loshchilov, I. Decoupled weight decay regularization. arXiv preprint arxiv: 2017.
Bubeck, S. et al. Convex optimization: Algorithms and com-

---

Nesterov Method for Asynchronous Pipeline Parallel Optimization
Merity, S., Xiong, C., Bradbury, J., and Socher, R. Pointer sentinel mixture models, 2016.
Ryabinin, M. and Gusev, A. Towards crowd- sourced training of large neural networks using decentralized mixture-of-experts in Neural Information Processing Systems, vol- ume 33, 2020. URI. courapps.ccppppe//22//////////// 254c0FfoccB822223338775333833833 Paper. pdf.
In Advances
Mishchenko, K., Bach, F., Even, M., and Woodworth, B. Asynchronous sgd beats minibatch sgd under arbitrary delays. URL htlps://arxiv. orglabw/2206.07633 2(6):7, 2022.
Mitliagkas, I., Zhang, C., Hadjis, S., and Ré, C. Asynchrony begets momentum, with application to deep learning. In 2016 54th Annual Allerton Conference Communi- cation, Control, and Computing (Allerton), pp. 997-1004. IEEE, 2016.
Ryabinin, M., Borzunov, A., Diskin, M., Gusev, A., Mazur, D., Plokhotnyuk, V., Bukhtiyarov, A., Samygin, P., Simitsin, A., and Chumachenko, A. Hivemind: Decentralized Deep Learning in Py- Torch, April 2020. URL. https : / thb. learning-ar atmomeeeettteeeeeeeeeeeoe /hi heviiddd.
Narayanan, D., Harlap, A., Phanishayee, A., Seshadri, V., Devanur, N. R., Ganger, G. R., Gibbons, B., and Za-
haria, M. Pipedream: Generalized pipeline parallelism for dnn training. In Proceedings of the 27th ACM symposium operating systems principles, pp. 1-15, 2019.
Ryabinin, M., Dettmers, T., Diskin, M., and Borzunov, A. Swarm parallelism: Training large models can be surpris- ingly communication-efficeeen In International Confer- ence Machine Learning, pp. 29416-29440, PMLR. 2023.
Narayanan, D., Phanishayee, A., Shi, K., Chen, X., and Za- haria, M. Memory-efficient pipeline-paralle dnn training. In International Conference on Machine Learning, pp. 7937-7947. PMLR, 2021a.
Stich, S. U. and Karimireddy, S. The error-feedback framework: Better rates for sgd with delayed gradi- ents and compressed communication. arXiv preprint arXiv: arXiv: 1909.0530, 2019.
Narayanan, D., Shoeybi, M., Casper, J., LeGresley, P., Pat- wary, M., Korthikanti, V., Vainbrand, D., Kashinkunti, P., Bernauer, J., Catanzaro, B., et al. Efficient large-scall scale language model training on gpu clusters using megatron- Im. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-15,20215 2021b.
Sutskever, I., Martens, J., Dahl, G., and Hinton, G. On the importance of initialization and momentum in deep learn- ing. In International conference on machine learning, pp. 1 1139-1147. PMLR, 2013.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L.. Gomez, A. N., Kaiser, L., and Polosukhin, I. At- tention is all you need. Advances in neural information processing systems, 30, 2017.
Nesterov, Y. A method for solving the convex programming problem with convergence rate (1/k2). In Dokl akad nauk Sssr, volume 269, pp. 543, 1983.
Wang, B. and Komatsuzaki, A. Gpt-j-6b: A billion param- eter autoregressive language model, 2021.
Nesterov, Y. Introductory lectures on convex optimization: A basic course, volume 87. Springer Science & Business Media, 2013.
Yang, B., Zhang, J., J., Ré, C., Aberger, C., and De Sa, C. Pipemare: Asynchronous pipeline parallel dnn training. Proceedings of Machine Learning and Systems, 3:269- 296, 2021.
PyTorch Contributors. Nadam optimizer pytorch 2.5.0 documentation. https://pytoz: org/ docs/ stable Staalle/generate/tooreee im. NAdam.
html, 2025. Accessed: 2025-01-16.
Zheng, S., Meng, Q., Wang. T., Chen, W., Yu, N., Ma, Z.-M., and Liu, T.-Y. Asynchronous stochastic gradient descent with delay compensation. In International conference on machine learning, pp. 4120 4129. PMLR. 2017.
Qi, P., Wan, X., Huang, G., and Lin, M. Zero bubble pipeline parallelism. arXiv preprint arXiv:2401.10044 10241, 2023.
Radford, A., Wu, J., Child, R. Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Zhu, Y., Kiros, R., Zemel, R., Salakhutdinov, R., Urta- sun, R., Torralba, A., and Fidler, S. Aligning books and movies: Towards story-like visual explanations by watch- ing movies and reading books. In The IEEE International Conference on Computer Vision (ICCV), December 2015.
Recht, B., Re, C., Wright, S., and Niu, F. Hogwild!: A lock- free approach to parallelizing stochastic gradient descent. Advances in neural information processing systems, 24, 2011.
10

---

Nesterov Method for Asynchronous Pipeline Parallel Optimization
A. Theoretical Analysis
first restate our NAG updates, and then turn proofs.
A.1. Nesterov Method for Delayed Gradients
Our variant of NAG performs the following iterations starting from an initial point W1:
d2 (w(: Wi-1),
(14)
W++1 Wt + d; 77(1 7h) V)V + dz). Here, 711 = and w, and d, denote the delayed versions weights and look-ahead, respectively, i.e.
W (-T We A
(15)
d d
where 2 0 the delay, and for simplicity we assume be fixed. The main difference the widely used NAG version the (1 - 7x) term for the gradients.
Let us introduce some notations that might be helpful later:
g: V.(w, + de)
(16)
h; -n(1 7
Å = w, + d, - (w, + de) A. + d de
g: and h, are analogously defined. From above, we may note the following identities:
W++1 = W; + d2 + h4 d = 7 (d, + h:),
(17)
We are now ready to prove that the look-ahead step acts as delay correction when the momentum coefficient % is chosen appropriately. Then, we prove the convergence rate.
A.2. Look-ahead as Delay Correction
Proposition 2. Let y be an increasing sequence such that lim 1 then, lim cos(A, a.) = 1, where cos(,,) the cosine similarity.
Proof. Let us expand the delay term:
A Wt Wi-r= Σ w. -
W
Σ
d:
Σ i=1-7+1
d + h Σ
d + he
(18)
76
i=(-++
(=t-++1
1=1
Note, all d;-; + hi- can be written in terms of d, di-;i
d + h,-, = - 7/ d + Σ
j= hi
(19)
h=t-r
(IT - 71) d
Σ (1i
(1-76) B
k=t-7
Now can write A, terms of d, follows:
2-1 d. - Σ
A
(u
(nr- (m²=)+?7) (1 - --x) Bk
(20)
ket-rr
When, 7/ + 1, the gradient term vanishes, and cos(Ar, d.) - 1.
Similarly, by writing A in terms of d,, one can show that cos( Ae, dl) + 1. This effectively shows that when the momentum coefficient 7/ - 1, the optimization trajectory becomes smooth and the delay and the look-ahead directions align.

---

Nesterov Method for Asynchronous Pipeline Parallel Optimization
A.3. Convergence Analysis
In this section, we will prove that NAG with delayed gradients converges at a rate O similar the standard gradient descent. The sequence "/t will be set part of the convergence proof. Even though, due to the delay the faster ((z)) rate is
not achieved, in practice NAG is effective as the synchronous alternatives despite the staleness in gradients.
Theorem 2. Let be convex, 3-smooth function with bounded gradients. then the iterates in Eq. (14) with n) 3 converges at a rate of (
Proof. Our proof largely follows the proof of (Bubeck et al., 2015) while adopting the delayed gradients. Let us first introduce a few useful inequalities:
- (x -) -11x 111x-10 yl?
J is 3-smooth
(21)
f(x) - f(y) (x(x)"(x->)) (x-x)) y) [V - ff(y)) - -1l
convex
3-smooth
Applying the first inequality above to the update equation, while using Eq. (16) and B
S(we+1) s(w, + de) S Sg
(A.
1 -
g+
+
8
A.
1
Y
g
(22)
8
2
3
1 - Vtt
3
2(1
(1 - )1
9t A,
gr
|A
g:
A
g+
B
8
g2
3 ||l
2 + 7:g: A,
7
|
2
23
Here, denotes the dot product. Similarly, using convexity:
f(w. + de) f(wi) S -dd A.) -g (d, Ae)
(23)
Now summing Eq. (22) and Eq. (23):
8,41 - 6, = [(W++1)) - s(w.)) 5 1 (7. - 11). A. + 7
Br 1 + gr d2.
(24)
23
where 61+1 = f(W(+1) f(w*). Similarly,
8 / 111..11 + (7c - 1)g: A + 7/
81+1 = f(w++1) - f(w") -
|
12
+
g:
-w: + dj -w).
(25)
23
Now, suppose d₂ > 1 be sequence. By multiplying Eq. (24) by (Ae 1) and adding to Eq. (25):
N 2
A. + Ae(1? -
Arde+1 - 1)); -
A + Ae(71 - 1)g,
|g.ll + (w: + Aed,- ww),
(26)
23
\1 + 7
B
= 24
hi (w; + drde w"),
2(1 - 7c)
7t
< Ωr
h, (w; + X,de w*),
B
2(1 7t)
7t
B
= 2 2 (1 76)
(11,, + w. + X.d, - -2 || + X,d, - -w"1l2)
C
where, A substitutes Q2
2 A 2 + Xe(7 1)g: A., B is due to < 7t < 0, > 0, > 0, and C follows from
a² + 2ab = (a + b)2
Now, to enable telescoping sum, we want,
Xph. + w, + A,d; w" wet1 + A411441 w" w, + d, ht + A41 + +h)) - w"
(27)
12

---

Nesterov Method for Asynchronous Pipeline Parallel Optimization
To this end, we set,
1 + 14417441 de-
(28)
Now, let the sequence d then,
2
72
and
1 7
(29)
Now, let
2 w. Xd; w'1l Then,
1
Xed1+1 (de -1); S Ω +
(u: ue+1),
(30)
Ae 71)
1
Aedt+1 x--10= S S4
+
2 [u: u(+1),
redc+1 - 20d1 $ Σ S4 + (n) u441):
Σ Lkmi for both sides
2
k=1
1101+1 S
S4 11,
Ao 0, 01+11 2 0
2
k=1
B w- w w
81+1
71 0
2t
To have the rate ) it remains to show that Σ k= Se grows much slower than O(t). To this end, using the bounded
gradients assumption we show || -well = (
). We prove this by induction. The base case for t < T can be enforced using an appropriate warmup phase. Suppose ||w, w = ( Then,
Wi+1 w. 2| l|w, de h w.
(31)
= || + hrll?
= >// - 1 + (1 72)
2yx(1 - 71)
||:
gi (w: - WL-1),
g
B
$ 7/11 - w + (1 )2
27x(1
71)
g
+
|g wt - we-1ll
g
B
0
Tt. ||
0(1), 7t 0()), and, || - w.-1 -
12
Now, it is easy to see that |d, - d.ll = ( as = 011) and ||w, 1oy-ww,pp w. o+t). Consequently, || = 0(+) due to bounded delay т. Consider.
X || AAe 1)gr A.
Ω
(32)
1,8
2|
|.l.
1
0
1
Therefore,
1 = O(It),
(33)
k=1
1=1
which completes the proof.
Note here that A, depends on the delay r. Despite it being a constant, it may influence the convergence rate it is sufficiently large. The proof largely follows the existing proof of NAG and analogously may be extended to non-convex functions. We leave any such analysis future work.
13

---

Nesterov Method for Asynchronous Pipeline Parallel Optimization
Furthermore, it may be intuitive to think of discounting term (1 -x) as learning discounting mechanism by considering t)t - O())- Analogously, relationships may be drawn from the convergence proof of such methods (Mishchenko et al., 2022). Nevertheless, in practical deep learning optimization, having separate provides greater control over the learning rate schedules and our approach is consistently better than the learning rate discounting mechanism as shown in our experiments.
B. Experiments
B.1. SWARM Training Configuration
We the SWARM baseline from Ryabinin et al. (2023) for large-scale decentralized training framework. For all baselines, the model used is Transformer language model with architecture similar to that in prior work (Brown et al., 2020; Wang & Komatsuzaki, 2021). Our SWARM configuration consists of 3 worker nodes per stage, for total of 24 worker nodes. Each worker node is assigned an NVIDIA L4 GPU, We assign 24 trainer nodes to serve the entire pipeline, where each trainer node has a 4-core Intel Cascade Lake CPU with a base clock of 2.2 GHz and 32 GB of RAM.
We use the following layer configuration for baselines: embedding dimension 768 with attention heads, the hidden layer dimension of the FFN layer of 3072 with layers. Each layer is assigned its own stage in the pipeline. The microbatch size used is 8 and sequence length is 2048. We employ learning rate of 2e-4 for the synchronous SWARM setting and our Nesterov- adapted approach. For the asynchronous variant of SWARM, we use reduced learning rate of 5e-5, due to training instability (and divergence) observed at higher learning rates.
A linear warmup was used for all baselines to Ik steps and our Nesterov-adapted approach employs a stage-dependent learning rate up to 2k steps (T in Eq. (13)) and the momentum coefficient 31 is set as per Eq. (13). For the other methods, a default value of 31 0.9 used. All methods were trained for total of 10k iterations, using stage-wise all-reduce batch size of 256.
B.2. Additional Results
provide validation loss trajectories for the main experiments in Fig. 9, training loss vs time for the model in Fig. 10, and more ablation results in Fig. 11. Furthermore, a comparison with the recent XPipe method (Guan al., 2019) is provided in Fig. 12. Finally, validation loss trajectories for the SWARM experiments provided in Fig. 13.
peppe
PipDrrcam
Mootrroam
7.0
Msccrromm
PiuMaae
Roeeare
1000 Wot0ttonn Lo55
Roeearm
Ours
Our
Ours
Valoutoon
OuTS No WS
Curs-No-WWe
Ours No-Wy
farations
20
20
20
(a) Wexi
(b) BookCorpus
(c) OpenWehText
Figure 9: Validation loss trajectory for the base model. The behaviour is the same as training loss where our method consistently outperforms all methods including GPipe.
14

---

Nesterov Method for Asynchronous Pipeline Parallel Optimization
GPipe
PipeDream
PipeMare
Training Loss
Ours
Ours-No-WS
0.0
2.5
5.0
7.5
10.0
12.5
15.0
17.5
Time (hr)
Figure 10: Training loss vs. wall-clock time for the IB model, where all models were trained for 50k iterations. For faster GPUs (A100 in this case), the runtime discrepancy between GPipe and other asynchronous methods is more pronounced. Note, overhead of our method over PipeDream is negligible.
RoeDream
10
PesDream
Porthram N4G-n03
cooldb Ae for Staap-
Yeirung toee
Treamream-Kee NAG-mD 99 (Ours)
PirDrramm N6G-a
PlocDraam 140 mc.39 (Oui))
8,033
AX032
puDrccmmMAAG
peDrremm-MGGGme (Dun)
x1.00
(a) Training Loss
(b) Weight Discrepancy for Stage-1
(c) Look-ahead and delay alignment
Figure 11: Ablation study of our main methods additionally showing weight discrepancy at Stage-1 in (b). Constant momentum coefficient of 0.99 performs slightly better than the adaptive version (denoted with 'aa ). It also shows the best delay correction and alignment between ahead and delay.
GPipe
PipeDream
PipeMare
Tranning Loss
XPipe
Ours
10
20
30
40 50
x1000 Iterations
Figure 12: WikiText results for the base model for the XPipe method (Guan et al., 2019) for completeness. XPipe is direct weight prediction method, that extrapolates the previous AdamW W step based the delay. We implemented XPipe following the description from the paper, however, we were unable to reproduce its reported performance on our large-scale language model training. Note, only small-scale image classification experiments were reported in the paper and in our experiments all methods including PreaD perform similarly (within pps) on those datasets.
15

---

Nesterov Method for Asynchronous Pipeline Parallel Optimization
SWARM Curs No WS
Validation Loss
10
X1000 Iterations
Figure 13: Validation loss on WikiText for the SWARM experiment. The observed performance follows similar trend the training performance, where method significantly outperforms the other methods.
16


Reviews:
[{"review_id": "fp9Nf3ctQo", "reviewer": "Reviewer_tQWW", "summary": "The gradient staleness caused by existing asynchronous pipeline parallelism mechanisms hinders the practical usage in contrast to synchronous pipeline parallelims. This paper aims to solve tackle the stale gradients when using asynchronous pipeline parallelisms by nesterov accelerated gradients. The paper has made the following contributions: (1) the paper has demonstrated that the asynchoronous pipeline parallelism can outperform the synchronous alternative on large language model workloads with partical setups; and (2) the paper proposes both theoretical and empricial justification of the proposed methods.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for encouraging comments and address the specific points below.\n\n## Scalability and sensitivity to multiple GPU servers\n\nOur **SWARM experiments were conducted on multiple GPU servers** (24 to be exact) in GCP, confirming that our method works seamlessly in such scenarios (Fig. 8 and Appendix B.1). We will include this information in the main paper. Additionally, in the paper, we have tested the scalability of our method by increasing the parameter count (Fig. 3) and by increasing the number of pipeline stages (Fig. 5). In both cases, our method **shows superior scalability**.\n\n## Speed-up compared to its synchronous alternative\n\nAsynchronous methods offer 100% pipeline utilization by construction. In contrast, pipeline utilization of synchronous methods depends on the bubble size (Huang et al., 2019). For GPipe, the pipeline utilization is estimated as $\\frac{N}{N+P-1}$ (Yang et al., 2021), where $N$ is the number of microbatches and $P$ is the number of pipeline stages. Typically, $N = P$ and therefore, the pipeline utilization of GPipe is about 50%. Based on this, **asynchronous methods would be twice as fast compared to GPipe** in a homogeneous environment with full GPU utilization. In a heterogenous, bandwidth constrained setup, the speed-up can be higher as asynchronous methods completely mask the communication bottleneck.\n\n## Training time compared to synchronous methods\n\nAs noted by the reviewer, we have **already provided the training time comparison in Fig. 5**, where it is clearly visible that *GPipe becomes exponentially slower* when increasing the number of stages compared to our method. The reviewer’s confusion might be due to misinterpreting Pipedream as a synchronous method. However, we would like to clarify that **Pipedream is an asynchronous method** and the *training time is the same for all asynchronous methods* including Ours, Pipedream, and Pipemare. The training curves for the 1B parameter model in Fig. 10 in the appendix illustrates this, and shows that the synchronous **GPipe is about 1.5x slower than the asynchronous methods** in this case.\n\nNevertheless, as mentioned in the paper, the absolute wall-clock times are not comparable between asynchronous methods (ours, pipedream and pipemare) and the synchronous method (GPipe), due to the differences in the underlying implementation – asynchronous methods are based on the 6 years old third-party pipedream codebase, whereas GPipe is a pytorch official implementation. \n\n## Inference results\n\nWe have **provided perplexity scores on the validation set in Table 1** as a quantitative measure of inference results. Additionally, the validation loss curves were provided in Fig. 9 and also Fig. 3 right. As requested by the reviewer, we have now provided some generation results for GPipe and Our method below as qualitative results: \n\n|Prompt|GPipe Generation|Ours Generation|\n|-|-|-|\n|Cyclone|Iloʿ Trespass was a little slower organized than initially, but rapidly strengthened to become a hurricane … |Tropical Storm Bonnie (Harak 10) made landfall on the Outer Banks of North Carolina, dropping moderate rainfall … |\n|Largest country in the world|Hair bars at Festival Day Student Association events ...|France ...|\n|Wikipedia has a variety of articles in topics like|The regime constitutes a more moral dependence on Wikipedia than … |vernacular romances and classical history …|\n|Some of the popular sports in the world include|[Wurgar Canoe] (meaning football game) Wurgar Canoe, the city hall, … |Cricket and ice hockey (football, soccer, ice hockey, ice hockey, men's and women's hockey ) are the sports …|\n\nNote, this is the output of pretraining 1B model on the WikiText dataset which is relatively small and contrived to train useful LLMs compared to larger datasets like C4 or FineWebText. Furthermore, post-training plays a major role in incorporating instruction following capabilities. Nevertheless, both methods generate plausible English sentences relevant to the prompt, whereas our method seems to be better for some prompts. Also, as seen in Figs. 2, and 9, the models are not yet saturated and can be trained further to improve the generation quality.\n\n## Convolutional networks\n\nTo test the generality of our method, we trained a Resnet50 model on the TinyImageNet image classification dataset with 8 pipeline stages. In short, the conclusions in the paper hold: \n|Method|Top-1 Validation|\n|-|-|\n|GPipe|55.70|\n|Pipedream|55.02|\n|Ours|**55.79**|\n\nThis is a small network and dataset in which Pipedream yields similar results to GPipe, even without any delay correction, highlighting that asynchrony is not a major pain point in these tasks. This observation aligns with previously reported results (Yang et al., 2021).\n\n## Low-precision training\n\nWe see no restrictions in using low-precision training in our method. We expect it to affect asynchronous and synchronous methods in the same way, however, experiments are needed to validate this.", "claims_and_evidence": "While the paper has shown great accuracy improvement by the proposed methods, the performance evaluation is limited (Figure 5 right). It is thus obscure how the proposed methods will work in practice. The proposed methods should also be compared against other synchronous methods (e.g. PipeDream) in terms of performance to claim it is better than other synchronous methods.", "methods_and_evaluation_criteria": "The paper is evaluated with commonly used language models and the experimental setups does align with the problem settings by comparing the loss curves of training against other baselines.", "theoretical_claims": "I do not find any problems in the theoretical analysis in the paper but my expertise is limited.", "experimental_designs_or_analyses": "While it is great to see the training curves of the proposed methods, it is obscure how the method could affect the generated outputs in practice. I would recommend the authors to add the inference results (e.g. the generated text from langauge models) of the trained models with different baselines.", "relation_to_broader_scientific_literature": "The paper may also include the discussion about how the proposed methods could be applied to training with low-precision floating point numbers, since they are becoming more and more prevalent in training modern DNNs.", "essential_references_not_discussed": "I think the paper has already discussed all essential related works.", "other_strengths_and_weaknesses": "I do find the evaluation of the paper robust since it have evaluated on multiple widely used datasets. However, it would also be interesting to see how the proposed methods perform on other types of networks (e.g. convolutional networks).", "other_comments_or_suggestions": "The paper's writing is in good shape. However, I would suggest the author to consider moving the related work section (section 4) as the second-to-last section of the paper.", "questions_for_authors": "* How is the proposed method scalable & sensitive to multiple GPU servers?\n* How much speedup can the proposed method achieve compared to its synchronous alternative?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kSZ5vbhlHQ", "reviewer": "Reviewer_wivu", "summary": "This paper tells the story of overcoming a major challenge in training huge neural networks with pipeline parallelism. When models are split into stages running on different devices, asynchronous updates keep the pipeline full but introduce the problem of stale gradients, updates based on outdated information. To fix this, the authors reimagine Nesterov’s Accelerated Gradient method: they modify its look-ahead step to serve as a delay correction, effectively aligning the updates with the current gradient direction despite the delay. Their theoretical analysis shows that this approach converges reliably, and experiments on large-scale language models, even one with 1B parameters, demonstrate that it outperforms existing asynchronous methods and even beats synchronous training.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for constructive feedback and address the specific concerns below.\n\n## Ablations for NAG\n\nWe have **already provided the ablation experiments for our modifications in Sec. 5.6**. Specifically, the effect of the momentum coefficient for NAG is reported in Fig. 6 and the effect of the gradient discounting term (our main modification to NAG in Eq. 9) is illustrated in Fig. 7. We believe we have clearly justified our modifications to NAG, however, if the reviewer believes any specific experiment is missing, we are happy to include them.\n\n## How can an asynchronous method be better than GPipe?\n\nThis is an intriguing question and we believe further study is required to completely understand the effects of asynchronous updates. Nevertheless, we answer this from a practical point of view below. \n\nThere are two points to consider when comparing two algorithms: 1) the frequency of the weight updates, and 2) the effectiveness of the weight updates. Typically, asynchronous methods perform more frequent weight updates than synchronous methods (such as GPipe) when processing the **same amount of data**. Specifically, GPipe accumulates gradients for a particular number of steps (4 in our experiments) to increase the pipeline utilization (Huang et al., 2019)) and update the weights, whereas asynchronous methods update the weights for every microbatch. Since more frequent updates (if the updates are beneficial) translate into a faster reduction in loss, an asynchronous method can be better than GPipe.\n\nHere, the key point is “if the updates are beneficial” – as not all asynchronous methods are better than GPipe. In fact, Pipedream and Pipemare are significantly worse than GPipe despite doing 4 times more weight updates (Fig. 2). This takes us to the second point, the effectiveness of the updates. To test this, we trained our base model (134M model on WikiText, same as in the paper) with 2 pipeline stages and updated the weights at every microbatch for GPipe and compared against our method (same number of weight updates for both methods). Even in this case, **our method is better than GPipe**: training loss at 50k iterations is **3.275 vs 3.323**. This confirms that *weight updates of our method are more effective than GPipe updates*, and overall, it is an effective asynchronous PP optimization method.\n\n## At how many workers does this modification start making sense?\n\nSince we consider pipeline parallel optimization, single worker setting does not make sense, and at minimum, we require 2 pipeline stages. In the paper, we have tested our method by varying the number of pipeline stages from 4 to 24 (Fig. 5). We now tested the same base model (134M model on WikiText) for 2 pipeline stages, and even in this case **our method outperforms GPipe in terms of both training and validation losses**: \n\n|Method|Train Loss| Validation Loss|\n|-|-|-|\n|GPipe|3.374|3.410|\n|Ours|**3.275**|**3.318**|\n\nThis further confirms the effectiveness of our method even when the number of pipeline stages is small.\n\n### NAG modification might not hold in different settings\n\nWe agree with the intuition that our gradient discounting term can be seen as an interplay between learning rate and momentum coefficient. However, we are unclear about the reviewer’s point that “the modification might not hold in different settings”. In case our response above misses any setting that the reviewer is interested in, if the reviewer can be more specific, we can try to better clarify them.", "claims_and_evidence": "The claims are backed by both empirical and theoretical results.", "methods_and_evaluation_criteria": "The evaluation includes several datasets for LLM evaluation.", "theoretical_claims": "The proposed algorithm is backed by a theoretical convergence proof.", "experimental_designs_or_analyses": "The experiments are conducted on several datasets and the proposed algorithm is compared to several previous works. The result include real-world scenario and convergence speed, both in terms of time in iterations. However, the proposed modification to NAG is not well justified empirically and the paper does not provide ablations for this modification.", "relation_to_broader_scientific_literature": "This work extends the broader scientific literature on distributed optimization and parallel training by addressing the longstanding challenge of gradient staleness. Building on prior methods that either forecast gradients or use weight stashing and learning rate adjustments to mitigate delays, the paper uniquely repurposes the look-ahead step of Nesterov’s Accelerated Gradient method to serve as an intrinsic delay corrector. In doing so, it offers a theoretically sound and practically effective solution that outperforms earlier approaches in large-scale, asynchronous pipeline parallel training.", "essential_references_not_discussed": "n/a", "other_strengths_and_weaknesses": "n/a", "other_comments_or_suggestions": "n/a", "questions_for_authors": "* **Q.1.** How come the proposed asynchronous algorithm outperforms synchronous (GPipe) in terms of convergence speed in number iterations?\n* **Q.2.** The modification to NAG can be see as a dependency between learning rate and momentum coefficient which might not hold in different settings. Would it make sense to validate this on a single worker setting? If not, then at how many workers does this modification start making senes?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jZ1BdttEqa", "reviewer": "Reviewer_xheX", "summary": "In this paper, the authors proposed a Nesterov Accelerated Gradient algorithm for asynchronous pipeline parallel optimization. It use Nesterov acceleration to reduce the negative impact of delay in asynchronous iterations. The authors firstly provide the convergence guarantee that the Nesterov acceleration can help the asynchronous pipeline parallel algorithm achieve the $\\mathcal{O}(1/T)$ convergence rate. Experiments on LLMs also validate the benefits of the proposed algorithms.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for encouraging comments and address the specific points below.\n\n## Results on the validation set\n\nWe have **already provided the results on the validation set** in Table 1 (perplexity scores) and Fig. 3,9 (validation loss curves). Additionally, as requested by Reviewer tQWW, we have included some generated text as qualitative results in the response to tQWW.\n \n## Theoretical analysis in the stochastic setting\n\n**We followed the standard practice in machine learning literature to provide theoretical analysis in a simplified setting and validate it using realistic large-scale experiments**. Many seminal works can be pointed out as examples (Kaifeng and Li 2019, Soudry et al 2018, Zhihui, et al. 2021). Furthermore, there have been instances where convergence analysis derived for non-stochastic settings are shown to be validated by experiments in the stochastic settings (Wu et al 2023, Arora et al 2018). We believe ours is one such example. Rigorous theoretical analysis matching the practical setting of large-scale language models requires significant research effort and we believe it would warrant a standalone publication.\n\nWe emphasize that our empirical results including the results with the 1B parameter model, and the SWARM experiments validate the effectiveness and practical usefulness of our method beyond doubt. \n\n- Lyu, Kaifeng, and Jian Li. \"Gradient descent maximizes the margin of homogeneous neural networks.\" ICLR (2019).\n- Soudry, Daniel, et al. \"The implicit bias of gradient descent on separable data.\" JMLR (2018)\n- Zhu, Zhihui, et al. \"A geometric analysis of neural collapse with unconstrained features.\" NeurIPS (2021)\n- Wu, Yongtao, et al. \"On the convergence of encoder-only shallow transformers.\" NeurIPS (2023).\n- Arora, Sanjeev, et al. \"A convergence analysis of gradient descent for deep linear neural networks.\" ICLR (2018).\n\n## Code\n\nAs requested by the reviewer, we provide a minimal version of our code [here]. Upon publication, we intend to release the code for reproducibility and to enable future improvements.\n\n[here]: https://drive.google.com/file/d/10Alm3sgwSic3Hi3w7gbJrme5JgUcsP9j/view?usp=sharing\n\n## Typo in Eq. 5\n\nThanks for pointing out the extra closing parentheses. We will correct it.", "claims_and_evidence": "Yes.", "methods_and_evaluation_criteria": "Yes.", "theoretical_claims": "The reviewer has taken a check on the theoretical results in this paper and not found any fatal mistakes. However, one limitation of the theory is that the convergence analysis focus on deterministic scenarios. The theoretical claims will be better if the authors provide the analysis of the stochastic scenarios.", "experimental_designs_or_analyses": "The experiments have shown the benefits of the proposed algorithms. However, the reviewer suggests the authors present some results on the validation set so as to illustrate the value in practical application. Moreover, it is better for the authors to offer their core codes.", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "**Strengths**\n1. Good presentation and writing.\n2. Detailed analysis on the insight of algorithm design.\n\n**Weaknesses**\n1. There is a typo in Eq. (5). Please take a check and revision.\n2. The convergence analysis is based on gradient descent without any stochastic noise. However, the convergence on stochastic algorithms like SGD or Adam may be essential for LLMs due to their small batch size.", "other_comments_or_suggestions": "NA", "questions_for_authors": "1. Whether the author can provide the experimental results on validation?\n2. Whether the author can provide a theoretical analysis based on stochastic algorithms？\n\nThe reviewer may consider raising the score on overall recommendation according to the response to the weaknesses and questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents addresses gradient staleness in asynchronous pipeline parallelism by reimagining Nesterov's Accelerated Gradient method. The authors modify the look-ahead step to serve as a delay correction mechanism, enabling reliable convergence despite asynchronous updates. The work provides both theoretical guarantees and empirical validation on large language models, including a 1B parameter model.

The paper has received positive reviews from all three reviewers, with particular praise for its theoretical foundations and practical impact. The authors have demonstrated that their asynchronous method can outperform synchronous alternatives in both convergence speed and final accuracy, while maintaining full pipeline utilization.

The authors have thoroughly addressed reviewers' questions during the rebuttal phase. They provided additional experimental results showing the method's effectiveness with as few as two pipeline stages, clarified the relationship between update frequency and convergence speed, and included qualitative results demonstrating the practical impact on model outputs. The authors also openly acknowledged limitations regarding theoretical analysis in stochastic settings, positioning this as future work.

Some limitations were noted, including the lack of theoretical analysis in the stochastic setting and questions about the precise mechanisms enabling asynchronous methods to outperform synchronous ones in terms of iteration efficiency. However, these limitations do not significantly detract from the paper's contributions.
Decision: Accept (poster)