Nesterov Method for Asynchronous Pipeline Parallel Optimization

We thank the reviewer for encouraging comments and address the specific points below.

Results on the validation set

We 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.

Theoretical analysis in the stochastic setting

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.

We 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.

• Lyu, Kaifeng, and Jian Li. "Gradient descent maximizes the margin of homogeneous neural networks." ICLR (2019).
• Soudry, Daniel, et al. "The implicit bias of gradient descent on separable data." JMLR (2018)
• Zhu, Zhihui, et al. "A geometric analysis of neural collapse with unconstrained features." NeurIPS (2021)
• Wu, Yongtao, et al. "On the convergence of encoder-only shallow transformers." NeurIPS (2023).
• Arora, Sanjeev, et al. "A convergence analysis of gradient descent for deep linear neural networks." ICLR (2018).

Code

As 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.

Typo in Eq. 5

Thanks for pointing out the extra closing parentheses. We will correct it.

We thank the reviewer for constructive feedback and address the specific concerns below.

Ablations for NAG

We 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.

How can an asynchronous method be better than GPipe?

This 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.

There 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.

Here, 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.

At how many workers does this modification start making sense?

Since 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:

This further confirms the effectiveness of our method even when the number of pipeline stages is small.

NAG modification might not hold in different settings

We 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.

We thank the reviewer for encouraging comments and address the specific points below.

Scalability and sensitivity to multiple GPU servers

Our 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.

Speed-up compared to its synchronous alternative

Asynchronous 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.

Training time compared to synchronous methods

As 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.

Nevertheless, 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.

Inference results

We 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:

Note, 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.

Convolutional networks

To 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:

This 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).

Low-precision training

We 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.