Paper: AMPipe: Accelerating MoE Model Training with Intra-Block Pipelining (ICLR-2024)

Reviewer comment (Novelty):
The technical novelty of the work is limited. Particularly, the idea of breaking GeMM computation into smaller chunks and overlapping those chunks with communication collectives has been previously studied, such as in Tutle. The main contribution of the work is to extend that to the Attention layer in an MoE model.

Author response:
We appreciate your insights regarding the technical novelty of our work. Indeed, while the MoE communication and computation pipeline have been well explored in numerous previous studies (e.g., Tutel[1], FasterMoE[2], Lina[3], Mpipemoe[4]), our work uniquely tackles the MoE-attention bottleneck within the entire transformer block. This perspective, we argue, transcends mere incremental advancement, offering novel contributions to the field of Mlsys research.

Reviewer comment (Experiments):
There is a big misconnection between the motivation of the work and the evaluation. Pipeline parallelism for MoE models is usually used for massive MoEs (e.g. hundreds of billions or even trillion-scale models) when the model cannot fit on aggregated GPUs with data parallelism, expert parallelism, and tensor parallelism. However, the tested MoE models are relatively small, with the largest one having only 12 layers and a hidden dimension 1600. This raises questions in terms of whether pipeline parallelism is even needed for this scale of models, i.e., the authors might have compared with a very sub-optimal configurations that are unnecessarily slow or have compared with configs that do not need pipeline parallelism. Indeed, from Figure 8, we see that as the model size increases, the gap between AMPIPE and Tutle decreases. To be more convincing, the paper should either test larger-scale models where pipeline parallelism is actually needed or use stronger baselines that sweep over different parallelism combinations.

Author response:
We would like to clarify that our experimental setup **does not employ a pipeline parallelism**. We categorize parallelism technologies into two categories: DP+EP+TP, which is used for optimizing single transformer blocks, and PP, which is used to schedule across multiple transformer blocks along the model depth. In this context, our proposed method, AMPipe, falls within the scope of DP+EP+TP, specifically targeting the optimization of a single transformer block.

Due to resource constraints, our experiments were conducted without the implementation of pipeline parallelism. This allowed us to specifically assess the speed performance of our method in transformer block optimization. Regarding the **pipeline degree** mentioned in Figures 7 and 8, it refers to the **internal pipeline degree within AMPipe and Tutel** (i.e., intra-block pipeline as the paper title) in Figure 6 rather than an inter-block pipeline degree as in Figure 4. We intend to provide a more precise definition of 'pipeline degree' in our revised manuscript. For a pipeline degree of n and an original sequence length of N, each sequence chunk in our model is designed to have a length of N/n. This refined explanation will be included in the revised manuscript to enhance clarity and understanding.

Reviewer comment (Reproducibility):
The paper lacks many important implementation details. In particular, while the projection of Q, K, V can be chunked along the sequence dimension, the softmax operation in the attention calculation requires all tokens in a sentence to calculate the normalized attention score. Therefore, it does not seem to be correct to claim that chunking Q does not affect the output of attention calculation unless additional synchronization is added. Also, the paper overlooks the normalization operations in its analysis, which are crucial for Transformer effectiveness and efficiency (as they may introduce additional synchronization that affects the pipeline's schedule). Unfortunately, the paper has been very vague on how these important operators are handled, e.g., in terms of how those operations are handled, the paper simply says "these operations are also amendable to chunking", but how and how would they affect the pipeline efficiency?

Author response:
Your question regarding the implementation details is important. Indeed, the results presented in Figure 8 encompass all elements within the transformer blocks, including both softmax and layernorm operations.
**Normalization operations will not affect the integrity** of AMPipe. AMPipe's algorithm should be exactly the same as the original PyTorch implementation. We assure you of the correctness of our algorithm and will provide more comprehensive details in our revised manuscript. Regarding the concerns raised:

(1) **Normalization Outside Softmax**: In modern LLM architectures like [GPT-2](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py#L367C24-L367C33), [OPT](https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py#L260C40-L260C49), [LLaMA](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L635C32-L635C44), and [Falcon](https://github.com/huggingface/transformers/blob/main/src/transformers/models/falcon/modeling_falcon.py#L732C28-L732C37), normalization operations (e.g., RMSNorm[5] and LayerNorm[6]) are typically applied **token-wise**. Consequently, when we chunk the sequence tokens into shards while retaining the integrity of each single token, the output remains unaffected.

(2) **Normalization Inside Softmax**: These operations are conducted on a **query-wise** basis, meaning that each query token undergoes normalization across the entire kv length. Thus, dividing the whole query sequence into chunks also does not affect the output as we maintain the integrity of each query token's computation.

Thus, both the normalization inside and outside softmax will not introduce additional synchronization that affects the pipeline's schedule.

Additionally, the implementation of FlashAttention-v2[7], which includes two loops — the outer loop iterating over the query length and the inner loop handling softmax normalization across the kv length — demonstrates that chunking the outer loop does not impact the results. This method maintains the integrity of the output compared to the original attention implementation, providing further evidence for point (2).

Furthermore, the Blockwise Parallel Transformer's[8] approach of splitting sequence along the query length also substantiates that such an operation does not compromise the model output, even with normalization operations. This approach offers validation for both points (1) and (2).

Reviewer comment (Reproducibility):
Do the results in Figure 8 include normalization operations, such as softmax, layer norm, etc.

Author response:
Your question regarding the implementation details is important. Indeed, the results presented in Figure 8 encompass all elements within the transformer blocks, including both softmax and layernorm operations.
**Normalization operations will not affect the integrity** of AMPipe. AMPipe's algorithm should be exactly the same as the original PyTorch implementation. We assure you of the correctness of our algorithm and will provide more comprehensive details in our revised manuscript. Regarding the concerns raised:

(1) **Normalization Outside Softmax**: In modern LLM architectures like [GPT-2](https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py#L367C24-L367C33), [OPT](https://github.com/huggingface/transformers/blob/main/src/transformers/models/opt/modeling_opt.py#L260C40-L260C49), [LLaMA](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L635C32-L635C44), and [Falcon](https://github.com/huggingface/transformers/blob/main/src/transformers/models/falcon/modeling_falcon.py#L732C28-L732C37), normalization operations (e.g., RMSNorm[5] and LayerNorm[6]) are typically applied **token-wise**. Consequently, when we chunk the sequence tokens into shards while retaining the integrity of each single token, the output remains unaffected.

(2) **Normalization Inside Softmax**: These operations are conducted on a **query-wise** basis, meaning that each query token undergoes normalization across the entire kv length. Thus, dividing the whole query sequence into chunks also does not affect the output as we maintain the integrity of each query token's computation.

Thus, both the normalization inside and outside softmax will not introduce additional synchronization that affects the pipeline's schedule.

Additionally, the implementation of FlashAttention-v2[7], which includes two loops — the outer loop iterating over the query length and the inner loop handling softmax normalization across the kv length — demonstrates that chunking the outer loop does not impact the results. This method maintains the integrity of the output compared to the original attention implementation, providing further evidence for point (2).

Furthermore, the Blockwise Parallel Transformer's[8] approach of splitting sequence along the query length also substantiates that such an operation does not compromise the model output, even with normalization operations. This approach offers validation for both points (1) and (2).

Reviewer comment (Evaluation):
To support training on extremely long sequence length, memory is also important to consider. I would like to know whether AMPipe saves memory.

Author response:
Your question regarding memory efficiency is pertinent. Building on the FlashAttention [1][2], AMPipe achieves linear memory usage with increasing sequence length. In our empirical evaluations, we observed no additional memory usage compared to current solutions, thus efficiently supporting training on extremely long sequences.

Reviewer comment (Experiments):
Sequence parallel and tensor parallel are also important techniques to support lengthy sequences by amortizing the memory onto different GPUs. It would be interesting to see how to integrate AMPipe with them.

Author response:
AMPipe is theoretically compatible with data, tensor, and sequence parallelism[3]. We have already integrated it with tensor parallelism in our latest version, and details of this integration will be included in the appendix of a forthcoming revision. Regarding sequence parallelism, its integration is non-trivial. A more dedicated design is necessary to effectively combine current sequence parallelism computation and communication with the MoE computation and communication, aiming to achieve enhanced performance.

Reviewer comment (Reproducibility):
Table 1 does not report the number of layers, the number of experts, and the size (number of parameters) of each model.

Author response:
Regarding the suggestion to add the number of layers and parameters in Table 1, we would like to clarify that our experiments in Sections 4.1 and 4.2 focus on speedup comparison within a transformer block instead of the whole model, emphasizing head and dimension settings. In this experiment, we use data and expert parallelism, and each GPU has one expert (varying expert numbers will only change local MLP, which is not a bottleneck in our evaluations). The base, medium, large, and xl transformer blocks have 7M, 13M, 20M, and 31M parameters, respectively. Additionally, we conducted an end-to-end experiment, as detailed in Section 4.3 and Table 2, using a cluster with eight A800 GPUs. We believe that our setting can reflect its real-world training potential to some extent.

Reviewer comment (Experiments):
If I understand well, AMPipe requires token-level communication in MoE layers, so I would like to ask whether it can be applied to sequence- or task-level MoEs?

Author response:
AMPipe is adaptable to both sequence-level and task-level MoEs. Given that AMPipe effectively serves the finer-grained token-level communication, it can also be applied to these larger-grained scenarios. While it is particularly well-suited to token-level communication in MoEs, which involves shuffling tokens across different devices, it can also accommodate larger-grained communication in sequence or task-level MoEs, where tokens from the entire sequences or tasks are sent to the same experts/devices. Many LLM-based MoE models and systems (e.g., Tutel[4], FasterMoE[7], Lina[8]) utilize token-level communication, and our empirical experiments demonstrate AMPipe's effectiveness in this token-level scope.

Reviewer comment (Evaluation):
When taking into account the tutel baseline, AMPipe provides a 7% improvement at best? (1.4x vs 1.3x speedup). This is effectively low impact.

Author response:
Your observation about AMPipe's 7% improvement over the Tutel baseline is appreciated. However, it's important to note that the speedups in Figure 7 vary with different model/sequence settings. The maximum speedup compared to Tutel is indeed higher, as illustrated by the 1.4x vs 1.23x difference. Furthermore, in the context of training large LLMs on thousands of GPUs, which require months of training time and millions of dollars[1], even a 10% speedup can translate into substantial financial benefits, underscoring the practical significance of the AMPipe.

Reviewer comment (Experiments):
The work does not show how this scales to large gpu clusters (where the majority of MoE training happens given MoE networks are relatively large).

Author response:
We appreciate your insight regarding scalability. We acknowledge that our evaluation, conducted on a smaller cluster, may not fully reflect the cases of training LLMs on real-world, large GPU clusters. We recognize that the communication bandwidth in larger clusters differs significantly from that in smaller setups. However, our present approach, encompassing diverse models, sequence settings, and low-bandwidth emulation, aims to replicate a range of real-world scenarios. We believe that performance in large GPU cluster environments will align with our experimental results to some extent, although perfect emulation may not be feasible.

Reviewer comment (Novelty):
While novel, the novelty is limited to combining tutel based pipelineing with pipelining based on chunked attention computation (an already existing concept).

Author response:
Thank you for pointing out our work's conceptual scope. While there are numerous studies on MoE communication and computation pipeline with various methods(e.g., Tutel[2], FasterMoE[3], Lina[4], Mpipemoe[5]), our work distinctively addresses the MoE-attention bottleneck from a **whole transformer block perspective**. Our approach, we believe, is more than incremental work, offering valuable insights for future Mlsys research.

Reviewer comment (Experiments):
This is generally more useful for long context length settings. But the majority of pretraining happens at seq len = 2k or 4k and long context length training only happens for a small subset of the total training time. This meta-point lowers the overall impact of works targeting improvements at really long context lengths.

Author response:
We acknowledge the prevalent trend of pretraining at shorter sequence lengths, a focus of prior MoE system works (e.g., Tutel, FasterMoE, Lina, Mpipemoe). However, AMPipe shows benefits in training with longer sequences, aligning with the emerging trend of longer sequences in LLM designs. We need to admit that addressing long context lengths is obtaining the importance, even though current methods often allocate a moderate portion of training time to them.

Reviewer comment (Presentation):
Can a variant of figure 8 be creates which combines the fwd and bwd pass? This allows the reader to see the total end to end speedup at different degrees instead of just seeing it in table 2 in tabular form.

Author response:
We appreciate your suggestion to revise Figure 8 to combine forward and backward passes. We will incorporate this enhancement in our revised version. However, as discussed in the paper, Figure 8 is a variant of Figure 7 across 56 benchmarks. Pipeline degree is a configuration for AMPipe, and the case with the highest speedup is more important (i.e., pipeline degree=4 case in Figure 7).

Reviewer comment (Presentation):
GPT-3 is not an open source model, I believe you mean GPT-2, considering the description of the model in Table 1.

Author response:
Thanks for pointing out our mistake. Our baseline is constructed on the open-sourced GPT-2, not GPT-3. This will be duly corrected in our revised manuscript.

Reviewer comment (Evaluation):
The fact that you provided code is great, but lines of code is not an informative metric. The optimum number of lines of code is the minimum necessary to have performant and readable code.

Author response:
Thanks for your constructive advice. We agree with your point on the metric of lines of code. Our revision will better reflect the emphasis on the performance and readability of the code rather than the quantity of code.

Reviewer comment (Writing):
It is worth explicitly defining what "pipe degree" is rather than keeping it implicit. Perhaps say that pipe degree = n and chunk length = N / n, etc.

Author response:
Thank you for highlighting the need to define "pipe degree" more explicitly. We confirm that your understanding is correct: with a pipe degree of n and an original sequence length of N, each sequence chunk will have a length of N/n. This definition will be included in the revised manuscript for clarity.

Reviewer comment (Writing):
Typo: “unofficially reported to use MoE paradigm.” should be “unofficially reported to use the MoE paradigm.”

Author response:
Thanks for your diligence in identifying typographical errors. These will be corrected to ensure clarity and professionalism in our presentation.

Reviewer comment (Writing):
Typo: “For MoE, sharding the input batch along seqeunce dimension” should be “For MoE, sharding the input batch along sequence dimension.”

Author response:
Thanks for your diligence in identifying typographical errors. These will be corrected to ensure clarity and professionalism in our presentation.