Paper: UniAP: Unifying Inter- and Intra-Layer Automatic Parallelism by Mixed Integer Quadratic Programming (ICLR-2024)

Reviewer comment (Theory):
The work makes certain assumptions about performance that are not always true. For instance, in Section 3.3, it states that "For Transformer-like models that mainly consist of the MatMul operator, the computation time in the BP stages is roughly twice that of the FP stages", which may not always be true depending on the optimizer.

Author response:
The assumption "the computation time in the BP stages is roughly twice that of the FP stages" is for estimating the overhead for backward computation time, which has been commonly adopted in distributed learning (Narayanan et al., 2021b; Li & Hoefler, 2021; Miao et al., 2022). UniAP follows this approach. It is important to note that the assumption is a trade-off for accuracy and fast cost modeling. Techniques such as modeling backward computation time for specific optimizers or different operators are not the current focus and are left for our future work.

Reviewer comment (Experiments):
Regarding the comparison to Alpa, the evaluation is not apple-to-apple. Alpa builds on JAX which uses XLA for JIT optimization, while Gavlaltron uses `torch.compile` for JIT. It would be best to compare Alpa and UniAP on the same framework, possibly done by porting to JAX.

Author response:
We acknowledge your opinion but assure you that we have strived to ensure fairness between Alpa and UniAP. Porting UniAP to JAX required rebuilding the entire profiling process, TP, PP, and FSDP module wrapper, communication primitives for specific parallelism schemes, and extensive work in debugging, verifications, and alignments. We were unable to meet our expectations, and we hope for your understanding.

However, we have revised the MFU results for our case study in the last paragraph of Appendix E. The metric is hardware and implementation independent. The results indicate that UniAP can achieve an MFU of approximately 23.6%, while Galvatron and Alpa can only achieve 13.7% and 19.6%, respectively. This may help to explain that UniAP's superior performance is solely due to its globally optimal parallel strategy.

Reviewer comment (Experiments):
The pipeline parallelism formulation of UniAP was based on a GPipe-style schedule, which has a major drawback of large bubble (idle) time. Would it be possible to adapt to more advanced PP schedules like DAPPLE or Chimera, and how would that affect your formulation?

Author response:
A1: More advanced PP schedules are possible. For example, the synchronous 1F1B pipeline proposed in PipeDream-2BW and DAPPLE maintains the MIQP's objective function due to the same bubble ratio as GPipe but offers reduced memory usage. To accommodate this change, one needs to adjust the constant coefficient for the activation memory term in the memory cost model. This illustrates the MIQP algorithm's flexibility in adapting to various PP strategies with an if-else clause.

Reviewer comment (Reproducibility):
It states that UniAP "partitions stages and determines micro-batch size using naive greedy algorithm". Could you elaborate on the exact algorithm you used, and whether you have reported the time for searching the microbatch+#PP stage in the search time (Table 1&2)?

Author response:
A2.1 & A2.2: The method "partitions stages and determines micro-batch size using a naive greedy algorithm" is adopted by Galvatron rather than UniAP. UniAP exhaustively searches for micro-batch size + #PP stages, as described in section 3.5.

Yes, we have reported the time for searching the micro-batch + #PP stage in the overall search time by "measuring the time of the UOP" as described in the experimental section in section 4.

Reviewer comment (Theory):
Section 3.4 formalizes the MIQP; however, it does not incorporate constraints related to computing resources. More precisely, there is no restriction on the number of GPUs employed by the strategy, thereby allowing it to stay within the available GPU count.

Author response:
The number of GPUs has been implicitly included in the strategy matrix $S$. We enumerate all possible combinations of intra-layer strategies for each layer given the partitioned pipeline stages. Each pipeline stage has a designated number of GPUs, and thus the computing resources are well constrained.

Reviewer comment (Reproducibility):
What is the form of the inter-parallel strategy? Is it a simple selection from DP TP FSDP, or is it a more complex form?

Author response:
The form of the inter-layer parallel strategy is pipeline parallelism (PP), consistent with previous work (Zheng et al., 2022).

Reviewer comment (Experiments):
The maximum parameter size of the models used in the experimental section is 1.02 billion. For the current trend in the development of large models, these models are slightly smaller. We look forward to the authors presenting the parallel performance of larger models. Perhaps your available computing resources limit the use of larger models, and in that case, simulated results would also be acceptable.

Author response:
Due to limited computing resources, we are unable to test larger models. However, our evaluation of end-to-end (E2E) throughput estimation against actual results shows that the E2E estimation error is limited to 3.59% in our environments. This observation highlights the potential of UniAP in large-scale clusters. Please refer to C2 and R2 in General Responses for further elaborations.

Reviewer comment (Evaluation):
How is the accuracy of your cost model?

Author response:
We have included end-to-end (E2E) assessments of estimated throughput accuracy in Appendix G. In our experiments, the average relative estimation error is 3.59%, significantly lower than Galvatron's 11.17%. This outcome showcases the robustness of UniAP's performance modeling.

Reviewer comment (Theory):
Section 3.4 formalizes the MIQP; however, it does not incorporate constraints related to computing resources. More precisely, there is no restriction on the number of GPUs employed by the strategy, thereby allowing it to stay within the available GPU count.

Author response:
The number of GPUs has been implicitly included in the strategy matrix $S$. We enumerate all possible combinations of intra-layer strategies for each layer given the partitioned pipeline stages. Each pipeline stage has a designated number of GPUs, and thus the computing resources are well constrained.

Reviewer comment (Reproducibility):
What is the form of the inter-parallel strategy? Is it a simple selection from DP TP FSDP, or is it a more complex form?

Author response:
The form of the inter-layer parallel strategy is pipeline parallelism (PP), consistent with previous work (Zheng et al., 2022).

Reviewer comment (Experiments):
The maximum parameter size of the models used in the experimental section is 1.02 billion. For the current trend in the development of large models, these models are slightly smaller. We look forward to the authors presenting the parallel performance of larger models. Perhaps your available computing resources limit the use of larger models, and in that case, simulated results would also be acceptable.

Author response:
Due to limited computing resources, we are unable to test larger models. However, our evaluation of end-to-end (E2E) throughput estimation against actual results shows that the E2E estimation error is limited to 3.59% in our environments. This observation highlights the potential of UniAP in large-scale clusters. Please refer to C2 and R2 in General Responses for further elaborations.

Reviewer comment (Writing):
I found the paper somewhat hard to follow, and details are unclear. This is especially the case for the formulation of the actual optimization problem, which would benefit from an explanatory figure (and possibly a table of notation). There are also important details that are elided, such as the communication model (the paper states "UniAP estimates the communication time by dividing the size of transmitting tensors by the communication efficiency for different communication primitives"); what exactly is used as the "communication efficiency"? This seems especially unclear as common libraries (e.g., NCCL) implement different algorithms for the same operation, which can have very different performance (e.g., logarithmic versus linear terms in tree versus ring allreduces).

Author response:
W1: Hard to follow and details are unclear.

W1.1: The formulation for the actual optimization problem lacks explanatory figure.

R1.1: We have included an additional figure in the appendix to help explain the optimization problem. Please refer to it, along with the supplemental descriptions in Appendix F.

W1.2: Important details for communication model are elided.

R1.2: We appreciate your feedback regarding the term "communication efficiency." We will revise it to "profiled communication efficiency" for better clarity. In the context of DP, TP, and FSDP, there are three main CCs: all-reduce, all-gather, and reduce-scatter. Each of these algorithms has its own complexity under different circumstances, but they are all optimized to suit specific hardware or networks. Our paper currently focuses on the parallel strategy search process and treats the specific collective algorithm chosen by the user or backend as a black box. We view this as a trade-off between accuracy and complexity. We have conducted an end-to-end (E2E) evaluation on our performance modeling in Appendix G. In the current evaluation, the average relative estimation error for UniAP is approximately 3.59%, which is lower than that of Galvatron.

Reviewer comment (Novelty):
The paper claims to unify inter- and intra-layer parallelism, but it does not seem that way to me. In particular, UniAP appears to define a cost model which can be used to evaluate the performance of a configuration, but then exhaustively enumerates different pipeline configurations (Section 3.5 & Algorithm 1) in order to find the one with least cost. The paper should address this, and also discuss how this approach compares to those of prior works (e.g., the dynamic programming formulation in Alpa).

Author response:
W2: See W2.1 and W2.2.

W2.1: Misclaim on unifying inter- and intra-layer parallelism.

R2.1: Unifying inter- and intra-layer parallelism and finding one strategy with minimum cost are two sides of the same coin. To jointly optimize for a globally optimal inter- and intra-layer parallelism strategy, we set our goal to selecting an inter- and intra-layer strategy with maximum throughput, or in other words, minimum time-per-iteration (TPI). Thus, we have to model the time cost for each possible inter- and intra-layer parallelism strategy, as well as the memory cost to avoid out-of-memory (OOM) exceptions. Therefore, the cost model, the enumeration process in UOP, and so on, is our proposed solution to jointly optimize for a globally optimal strategy.

W2.2: How does UniAP compare to prior works?

R2.2: Existing methods do not optimize these strategies jointly from a global perspective. For instance, Alpa feeds the ILP results for intra-operator stages to inter-operator stages for an optimal solution. However, combining local optima doesn't always ensure a global optimum. In certain cases, two consecutive operators with different optimal intra-operator sharding specs may result in higher communication costs between them. In systems with PCIe interconnections, it may be globally optimal to compromise on one of them to keep them the same, as it avoids invoking any CCs. Unfortunately, Alpa is unable to identify such cases as it only returns a single solution rather than a series of candidates after each intra-operator pass. This is where UniAP provides a unique solution.

Apart from Alpa, existing methods either overlook potential opportunities for inter- or intra-layer parallelism or optimize hierarchically instead of globally. They do not unify the complete inter- and intra-layer parallel strategy space to jointly search for a globally optimal solution. Based on these motivations, we identify two key challenges: formulating the problem globally and optimizing the formulation efficiently. To address these challenges, we extensively evaluated various existing methods and ultimately developed MIQP. To the best of our knowledge, none of the methods considered such an extensive strategy space and resolved the aforementioned challenges at the same time. Hence, UniAP distinguishes itself from prior works.

Reviewer comment (Experiments):
A related point, given the search over pipeline configurations, it is not clear to me how UniAP's search would scale. For example, how would it perform if the cluster contained 10 000 GPUs? (This is not at all unrealistic for large training clusters.)

Author response:
The search process will scale in the time complexity of $\mathcal{O}(|V||g|\log(n))$ as analyzed in Appendix C. Take $n=10,000$ GPUs as an example, the number of hidden layers $|V|$ remains a constant as it is orthogonal to the cluster size. Since 10,000 is not an exponent of 2, we enumerate its 25 factors to achieve load balance, resulting in a total of 2690 intra-layer parallelism candidates. Hence, the time complexity of $n=10,000$ GPUs would be $\mathcal{O}(2690|V|)$.

Reviewer comment (Experiments):
The paper omits important techniques used in practice from its evaluation, in particular, mixed precision (Section 4). While I appreciate the claim that this is orthogonal, using FP32 can result in significantly lower computation performance (due to not using tensor cores), which in turn can significantly change the communication/computation ratio and consequently the best parallel strategy. This is especially strange because Alpa evaluated transformers in FP16.

Author response:
In our paper, our focus is on developing a method to find the best hybrid parallel strategy, rather than seeking the best throughput across all existing training strategies. Therefore, enabling FP16, BF16, INT8, as well as activation checkpointing, gradient checkpointing, and other training techniques are irrelevant to our claims and contributions. However, these techniques are quite common in today's big model training, as you mentioned, and we will leave them for our future work. Regarding the experiments conducted on Alpa, we chose to use FP32 for comparison because FP16 was disabled in Galvatron's original paper. To maintain consistency and provide convincing results across the two baseline systems, we chose to align the floating point precision to 32 bits.