NetMoE: Accelerating MoE Training through Dynamic Sample Placement

Weakness 3 & Question 3

As pointed out by the reviewer, the Kuhn-Munkres (KM) algorithm does have a cubic time complexity. However, it is much more efficient than solving the original target optimization problem (defined in Equation 6,7), and fits the scenario of large language model training well.

• Firstly, as discussed in Section 3.2 and evaluated in our experiments (Table 4), the target optimization problem (defined in Equation 6,7) is an integer linear programming (ILP) problem, which is NP-hard. Although existing libraries like PuLP support solving ILP problems, it takes a long time for PuLP to accomplish the problem-solving, making it infeasible to hide the solving cost by overlapping. To cope with this problem, we designed a method based on the Kuhn-Munkres (KM) algorithm that can be done in polynomial time. Experiment results in Table 4 demonstrate that our KM-based solving algorithm is much faster than directly solving with PuLP, and we are able to completely overlap the problem-solving.
  In addition, the KM algorithm is extremely widely used to solve assignment problems. Considering that our work focuses on how to reduce the communication cost by re-assigning the training samples among the GPU devices, the KM algorithm is suitable for our work.
• Secondly, the complexity of KM algorithm is related to the batch size per device during training (yet independent of model size). In distributed training of large language models, due to the constraint of GPU memory, it is infeasible to support a large batch size per device at once. Instead, it is common to leverage the gradient accumulation technique [2,3] to ensure that the batch size per device for each gradient accumulation step is small (while allowing the model to be updated with the gradients of a large batch). Thus, the batch size per device for each step is typically small.
  In Section 4.4 of the revised manuscript, we further conduct an experiment with a batch size per device of 24 (a higher value would lead to out-of-memory errors) to examine the effectiveness of our KM-based solving algorithm. The results show that the problem-solving can still be overlapped well.

In the revised manuscript, we have added more discussion about why we solve the problem via KM algorithm in Section 3.2 and provided the experimental results in Section 4.4.

[2] Pytorch, “Gradient accumulation pytorch,” https://gist.github.com/thomwolf/ac7a7da6b1888c2eeac8ac8b9b05d3d3.

[3] Tensorflow, “Gradient accumulation tensorflow,” https://github.com/tensorflow/tensorflow/pull/32576.

Weakness 4 & Question 4

In our original manuscript, we consider training over 2 and 4 nodes, with each node consisting of 8 GPUs, which is a typical configuration in distributed training of large language models. To address the reviewer's concern, we further conduct experiments with 2 and 4 GPUs per node, respectively. The results are provided in Appendix B of the revised manuscript. Overall, NetMoE still consistently achieves the best performance.

We wish to express our sincere gratitude for your valuable feedback and thoughtful critique. We recognize the opportunities for improvement you've identified and believe that your insights will guide significant enhancements to our work.

Weakness 1 & Question 1

Indeed, our work does not provide advantageous scenarios for the model perspective and the data perspective, which could improve our work. However, we believe it does not harm the significance of our work. To the best of our knowledge, NetMoE is the first approach to accelerate MoE model training from the data perspective, and our evaluation also validates the effectiveness of NetMoE by comparing it with the approaches based on dynamic expert placement. Consequently, our work is of great significance as it offers a completely new paradigm for accelerating large language model training.

Moreover, as noted in Footnote 1 on page 6 of our manuscript, our work can be integrated with the dynamic expert placement technique, combining the two perspectives together to achieve further acceleration. We would like to leave it as our future work.

Weakness 2 & Question 2

Our assumption (inter-node communication cost is higher than intra-node communication cost) is reasonable: as shown in Table 2, the intra-node bandwidth is 4 times that of inter-node bandwidth. Intuitively speaking, as long as the intra-node communication volume (denoted as $s_{intra}$ in Section 3 of our work) is less than 4 times that of the inter-node communication volume (denoted as $s_{inter}$), the inter-node communication dominates. To address the reviewer's concern, we measure the intra-node and inter-node communication volumes (i.e., $s_{intra},s_{inter}$) before and after applying NetMoE. The results are presented in Table 5 of Appendix C in the revised manuscript.

The results demonstrate that the inter-node communication volume reduces substantially after applying NetMoE. Although the intra-node communication volume accounts for a large proportion (i.e., $s_{intra}$ or $\frac{s_{intra}}{s_{intra} + s_{inter}}$ increases) after applying NetMoE, it will not become the performance bottleneck since the ratio $\frac{s_{intra}}{s_{inter}}$ remains significantly smaller than 4. This indicates that the All-to-All communication can be accelerated due to the reduction in inter-node communication volume. Consequently, it is feasible to minimize the inter-node communication first, and then minimize the intra-node communication with a fixed inter-node communication.

In fact, due to the difference in bandwidth, it is a common practice to prioritize the reduction in the inter-node communication. For instance, [1] also focuses on minimizing the inter-node communication volume rather than the intra-node one (detailed in Section 3.1 of their paper). As a result, we believe the two-stage problem-solving in our work is reasonable and practical.

[1] Zhuang et al. On Optimizing the Communication of Model Parallelism. https://arxiv.org/abs/2211.05322.

We deeply appreciate your acknowledgment of our work and your constructive feedback. We are confident that our work will be significantly improved by incorporating your insights.

Weakness 1 & Question 1

The data locality is a well-known characteristic in MoE models and has been motivating many works to accelerate MoE training or inference [1,2,3,4,5]. Although some studies have tried to investigate the data locality given specific pre-trained MoE models and datasets [6,7], we wish to clarify that given any dataset and MoE model, the routing distribution would dynamically change during the training, so it is hard to control the distribution to examine the speedup in end-to-end model training. Therefore, to address the reviewer's comment, we assess whether NetMoE can reduce the inter-node communication volume when facing the dynamicity in routing distributions.

To be specific, we follow prior works [1,2] to record the distribution of expert selection across different iterations in order to describe the routing distribution and record the reduction in inter-node communication as well. The results are provided in Figure 10 of Appendix C in the revised manuscript. It can be seen that the routing distribution changes during the model training process, while NetMoE consistently reduces the inter-node communication by adjusting the sample placement given the dynamic distributions. Consequently, the effectiveness of NetMoE is general to various data locality conditions.

[1] He et al. FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models. https://dl.acm.org/doi/10.1145/3503221.3508418.

[2] Nie et al. FlexMoE: Scaling Large-scale Sparse Pre-trained Model Training via Dynamic Device Placement. https://arxiv.org/abs/2304.03946.

[3] Zhai et al. SmartMoE: Efficiently Training Sparsely-Activated Models through Combining Offline and Online Parallelization. https://www.usenix.org/system/files/atc23-zhai.pdf.

[4] Li et al. Accelerating distributed MoE training and inference with lina. https://www.usenix.org/system/files/atc23-li-jiamin.pdf.

[5] Yao et al. Exploiting inter-layer expert affinity for accelerating mixture-of-experts model inference. https://arxiv.org/pdf/2401.08383.

[6] Jiang et al. Mixtral of experts. https://arxiv.org/abs/2401.04088.

[7] Xue et al. Openmoe: An early effort on open mixture-of-experts language models. https://arxiv.org/abs/2402.01739.

Weakness 2 & Question 2

Expert parallelism, tensor parallelism, data parallelism, and pipeline parallelism can be combined to achieve hybrid parallel training of large language models. Given the fact that tensor parallelism is usually applied within nodes due to its high communication volume [8], if we wish to avoid inter-node expert parallelism, there must be $TP \times EP \leq 8$, where $TP$ and $EP$ are the parallel degrees of expert parallelism and tensor parallelism, respectively. Undoubtedly, this would lead to a limited parallel configuration space. As a result, our work considers a more general case where expert parallelism involves both intra-node and inter-node communication for the experiments.

[8] Singh et al. A Hybrid Tensor-Expert-Data Parallelism Approach to Optimize Mixture-of-Experts Training. https://arxiv.org/abs/2303.06318.

Weakness 3 & Question 3

In our experiments, we have compared NetMoE with two state-of-the-art baselines that are based on dynamic expert placement, which are the FasterMoE and SmartMoE. The experimental results demonstrate that NetMoE consistently outperforms the baselines in terms of training efficiency. Besides, we did not incorporate any auxiliary loss in our experiments.

We are truly grateful for your careful review and thoughtful questions, which are valuable to our work.

Weaknesses

Due to the lack of GPU resources, we are not able to examine the scalability of NetMoE with more GPUs currently (e.g., 64 GPUs). However, we wish to highlight that the improvement of NetMoE does not diminish as the number of GPUs increases. To be specific, as shown in Figure 5 in the original manuscript (which is now Figure 6 in the revised manuscript), NetMoE achieves a slightly higher speedup with 32 GPUs compared to 16 GPUs for 4 of the 5 experimented models. For MoE-GPT-XXL, the speedup is very close: with 16 GPUs, the speedup is $1.67 \pm 0.039$, and with 32 GPUs, it is $1.65 \pm 0.030$. Across all experiments, the average speedup of NetMoE on 32 GPUs surpasses that on 16 GPUs. Consequently, the speedup delivered by NetMoE is robust to the increase in the number of GPUs.

To avoid ambiguity, in the revised manuscript, we have provided the standard deviation of end-to-end speedup in Figure 6.

Questions

To compute attention efficiently, all tokens of a training sample should reside on the same GPU device. If we exchange only part of the tokens, then the tokens of each training sample would be distributed across different GPU devices. In this case, it necessitates substantial extra communication to accomplish the attention computation, which is counterproductive. Therefore, we adjust the placement at the granularity of training samples.

Thank you for your acknowledgment and insightful comments. Your feedback is extremely helpful, and we are committed to addressing each question you have raised.

Weakness 1: Notations and problem formulation hard to follow.

To address the reviewer's concern, in the revised manuscript, we have added a figure (Figure 3 on page 4) to present the overview of Section 3, which briefly includes the problem formulation, two-stage dissection, polynomial-time solver, and implementation of our work. We hope it can improve the readability of our work.

Weakness 2: No comparison with methods using a modification in the model definition.

We agree that the approaches based on modification in model definition (which impact model convergence) could train with more iterations to reach a similar level of perplexity. However, NetMoE is orthogonal to such lossy approaches. In other words, they can be integrated with our work to achieve better efficiency. Since large language model training is time-consuming and costly, we cannot keep training models until convergence. As a result, our evaluation focuses on lossless approaches that do not affect model convergence.

In addition, we wish to highlight that developing lossless approaches is timely and important. To be specific, when applying lossy approaches, we usually need to run numerous trials to tune the hyper-parameters (e.g., we need to adjust the weight of the topology-aware routing loss [1], or need to tune the hyper-parameters for different communication channels [2]), which is impractical since each trial of large language model training may take days or even months.

In Section 2.2 of the revised manuscript, we have added more discussion to address the reviewer's concern.

[1] Chen et al. TA-MoE: Topology-Aware Large Scale Mixture-of-Expert Training. https://arxiv.org/abs/2302.09915.

[2] Zeng and Xiong. SCoMoE: Efficient Mixtures of Experts with Structured Communication. https://openreview.net/pdf?id=s-c96mSU0u5.

Question 1

We feed the current layer's input directly to the next layer's router to predict the expert routing of the next layer. Such a prediction method is also adopted in existing works [5,6]. The rationality is that, since residual connections are needed in transformer layers, the inputs to the routers of two consecutive layers should share certain similarities. This prediction doesn't need high accuracy and serves only as a guide during algorithmic optimization.

[5] Eliseev and Mazur. Fast inference of mixture-of-experts language models with offloading. https://arxiv.org/abs/2312.17238.

[6] Tang et al. HOBBIT: A Mixed Precision Expert Offloading System for Fast MoE Inference. https://arxiv.org/abs/2411.01433.

Question 2

In Appendix A of the revised manuscript, we have provided a detailed visual explanation of the residual inlining. Specifically, the original residual addition method adds the attention output to the result obtained from the gather operation. In NetMoE, however, it is added after the scatter operation but before the gather operation. Such an inlining facilitates the adjustment of sample placement, and meanwhile ensures the correctness of computation.

We sincerely appreciate your recognition of our work and thank you for your constructive suggestions. We believe that addressing your suggestions will substantially improve our work.

Weakness 1

To address the reviewer's concern, in Appendix C of the revised manuscript, we consider two kinds of statistics to assess the source of performance improvement of NetMoE:

• The proportions of training samples that are exchanged across nodes/devices by NetMoE. A higher proportion indicates more samples are adjusted across nodes/devices.
• The intra-node and inter-node communication volumes before and after applying NetMoE.

Firstly, we summarize the mean and standard deviation across all iterations in Table 5 of Appendix C in the revised manuscript. After applying NetMoE, a great proportion of training samples are exchanged across nodes, leading to the reduction in the inter-node communication volume. It is noteworthy that although the intra-node communication volume accounts for a large proportion (i.e., $s_{intra}$ or $\frac{s_{intra}}{s_{intra} + s_{inter}}$ increases) after applying NetMoE, it will not become the performance bottleneck since the inter-node communication bandwidth is much lower. As a result, the All-to-All communication can be accelerated due to the reduction in inter-node communication volume brought by sample placement adjustment.

Secondly, to discover the impact of router probability, in Figure 10 of Appendix C in the revised manuscript, we plot (1) the reduction in inter-node communication, and (2) the proportion of samples exchanged across nodes, across different iterations. Meanwhile, we follow prior works [1,2] to record the distribution of expert selection across different iterations in order to describe the routing distribution. It can be observed that the routing distribution changes during the model training process. However, NetMoE consistently reduces the inter-node communication by adjusting the sample placement given the dynamic distributions. Consequently, the effectiveness of NetMoE is robust to the router probability.

[1] He et al. FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models. https://dl.acm.org/doi/10.1145/3503221.3508418.

[2] Nie et al. FlexMoE: Scaling Large-scale Sparse Pre-trained Model Training via Dynamic Device Placement. https://arxiv.org/abs/2304.03946.

Weakness 2

We wish to clarify that for any dataset mixture, the routing distribution would be dynamic during the training of MoE models. Thus, it is more important to be robust w.r.t. routing distributions. As elaborated in our response to Weakness 1 and demonstrated in Figure 10 of Appendix C, the reduction in communication volume is consistent given various routing distributions, demonstrating the robustness of our work.

Weakness 3

In distributed training of large language models, due to the constraint of GPU memory, it is infeasible to support a large batch size per device at once. Instead, it is common to leverage the gradient accumulation technique [3,4] to ensure that the batch size per device for each gradient accumulation step is small (while allowing the model to be updated with the gradients of a large batch). Thus, the batch size per device for each step is typically small. In fact, for Table 4, when we try to increase the batch size per device to 32, the training task encounters an out-of-memory error.

In addition, in the original manuscript, Table 4 only compares the time cost of the problem-solving and the scatter operation for simplicity, while in practice, the problem-solving can be overlapped as long as its time cost is smaller than the summed time cost of the scatter operation and the expert computation. To further address the reviewer's concern, we conduct an experiment with a batch size per device of 24, which gives the result that the time cost of problem-solving is 31.09ms, while the summed time cost of the scatter operation and the expert computation is 41.65ms (33.82ms + 7.83ms), showing that the problem-solving process can still be overlapped well.

In the revised manuscript of our work, we have incorporated the results above in Table 4, and added more discussion about the batch size in Section 3.3.

[3] Pytorch, “Gradient accumulation pytorch,” https://gist.github.com/thomwolf/ac7a7da6b1888c2eeac8ac8b9b05d3d3.

[4] Tensorflow, “Gradient accumulation tensorflow,” https://github.com/tensorflow/tensorflow/pull/32576.

Thanks for the response.

The new results in Appendix A, B, C are great. It addresses my first and (at least) half of my second concern. My third concern is also half addressed and I would suggest to add a brief discussion on relevant MoE training recipes that do not need large batch size per device, and when we indeed need large batch size per device (and for GPU utilization purpose we sometimes need sufficiently large microbatch to medium-sized models), what is the other approximate solver available (this doesn't need to be perfect, but should still be better than random).

My 2 questions are well answered. Thanks for this clear response!