LSH-MoE: Communication-efficient MoE Training via Locality-Sensitive Hashing

Summary

The authors focus on the communication overhead in large-scale MoE training, specifically under the expert parallel plus the data parallel regime. They reduce the communication workload by transmitting only the clustering centroids, which are calculated on the fly using LSH functions. To reduce the compression-induced error, the authors also propose a residual-based error compensation scheme. The authors perform evaluations using 4 different architectures in both language and vision, and for both pertaining and finetuning tasks.

Strengths

1. This paper is primarily well-written, and the algorithm is presented clearly. 2. As far as I know, the idea of compress activation using cluster to reduce communication for MoE training is rather original. 3. The authors focus on a timely and important problem, especially because there has been a trend of adopting MoE in large-scale models recently.

Weaknesses

1. To convince the audience of this problem's importance and increase the work's significance, the authors may provide a more detailed analysis of the communication overhead. One example is given in Section 2.2. However, it would be better if the authors could study what factors influence the percentage, perhaps the relationship between the communication overhead and the scale of the training servers, the scale of models, etc. 2. Lack of background and related work. I would suggest the authors provide background on LSH algorithms, such as how to calculate centroid using LSH. Further, as far as I know, there are more works that improve the efficiency of MoE training, such as DeepSpeed MoE, DeepSpeed-TED, and SCoMoE, which should be discussed and compared.

Questions

1. Why did T5-MoE achieve a much faster convergence than RoBERTa-MoE? 2. Figure 5 shows a difference in validation perplexity( 2.37 vs. 2.5), while LSH-MoE seems not quite flattened out. Would continuing training LSH-MoE reach the same perplexity? 3. For the pertaining setting, what is the accuracy difference on zero-shot or few-shot downstream tasks between origin and LSH-MoE? The

Rating

6

Confidence

3

Soundness

2

Presentation

3

Contribution

2

Limitations

The authors claim there is no limitations and no societal impact of the work performed.

Summary

The paper introduces LSH-MoE, a communication-efficient training framework for Mixture-of-Experts (MoE) models using Locality-Sensitive Hashing (LSH). The authors identify the inefficiencies in existing MoE training methods, particularly the high communication costs due to all-to-all communications among GPUs. The proposed method leverages token similarity for data compression, significantly reducing communication overhead and achieving substantial speedups in training time while maintaining model quality.

Strengths

1. The use of Locality-Sensitive Hashing for compressing communication data in MoE training is novel and addresses a significant bottleneck in distributed training systems. 2. The authors conduct extensive experiments on various language and vision models, demonstrating the effectiveness of their method across different tasks and datasets. 3. The experimental results show impressive speedups (1.28-2.2×) in training time, making the approach highly beneficial for large-scale model training.

Weaknesses

Maybe it would be better to include more larger MoE results.

Questions

How does the authors view the potential speedup of LSH-MOE over larger MoE models like Mixtral 8x7B?

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

NA

Summary

This paper presents a method to speed up MoE large model training with locality-sensitive hashing. It conducts experiments on both language models and vision models for both pre-training and fine-tuning tasks, and achieves 1.28-2.2× of speedup.

Strengths

1. This paper introduces an efficient LSH-based compression technique to apply tokens on similar expert model, which is both effective and fast. 2. This paper is well written for understanding. 3. This paper conducts experiment on different types of large model and training tasks, which shows the generality of proposed method.

Weaknesses

1. There is a lack of the whole framework figure for the proposed LSH-MoE method for better understanding of reader. 2. More deeper analisis should be added for ablation study, the experiment result shows the effect of different quantity and type of hash functions. But we still don't have even a guess of reason. 3. Some expressions need to be improved for example, what is the difference between compression rates and compression ratio in Fig 6.

Questions

1. What's the difference between compression rates and compression ratio in Fig 6, and why we choose 20%,15%,10% compression rates for accuracy comparation. 2. We only try cross-polytope and spherical hashing for comparation, are there any other hashing types could be adopted? 3. Do we consider to adopt some learning to hash method for MoE model training speedup? Because as far as I see, learning to hash method usually keep semantic similarity better than LSH method.

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors have not addressed the limitations and negative societal impact of their work. This work gives a good training speedup solution for MoE large models. But the proposed method could not work well on none-MoE models. Maybe the author can consider some speedup method for general large models.

Rating after reading the rebuttal

I have read the rebuttal，and I would like to keep my rating for this paper.