Shortcut-connected Expert Parallelism for Accelerating Mixture of Experts

1. It's better to include more modern designs like the ones used in deepseek v2 (more smaller experts).

Given our constraints on available hardware, we cannot conduct experiments on large-scale MoE models like DeepSeek-V2 (236B). Nonetheless, we conduct experiments on the OLMoE model (7B), a recent MoE model that features more modern designs, incorporating the fine-grained expert architecture (more smaller experts) similar to DeepSeek-V2. We train both the original OLMoE and the OLMoE+ScMoE from scratch, using 40,000 samples from the TuluV3 dataset. As demonstrated in the following table, our proposed ScMoE architecture not only accelerates the training by 1.13x and inference by 1.19x, but also achieves a marginally higher average accuracy (43.77% vs 43.25%).

2. My primary concern is that the skip connection design lacks novelty. Additionally, the model quality requires further justification or experiments to substantiate that it's actually effective. Existing models such as GPT-J and GPT-NeoX already parallelize attention and feedforward networks, while architectures like Snowflake Arctic [1] employ a structure where the input to the self-attention layer is directly fed to the MoE layer - closely resembling this work’s proposal.

(1) GPT-J and GPT-NeoX utilize parallel Attention and Feed-Forward Layers to address the All-Reduce communication overhead inherent in Tensor Parallelism, which is notably different from the All-to-All communication overhead encountered in Expert Parallelism. Distinct from existing methods, our proposed ScMoE architecture is specifically designed to minimize the All-to-All communication overhead in Expert Parallelism.

(2) Regarding Snowflake Arctic, our work was conducted independently and concurrently with it, and it stands out in several distinct ways. Snowflake Arctic is primarily an industrial product and does not include any published papers or detailed analyses of their proposed Dense-MoE Hybrid Transformer architecture. In contrast, we thoroughly examine the effectiveness of Shortcut-connected MoE architectures and analyze the phenomena and reasons behind them, offering deeper insights for the research community. Furthermore, our proposed Shortcut-connected MoE demonstrates superior generalization across different MoE model designs (adapt to various MoE placement frequencies). It bears similarity to their Dense-MoE Hybrid Transformer architecture, specifically to our proposed ScMoE (Pos-1), which is more like a subset of our study.

3. However, current SOTA models (e.g., DeepSeek V3/R1, Llama 3.2) do not adopt this parallel structure, raising questions about whether the proposed architecture can achieve competitive accuracy/quality at scale. The experiments in this paper inadequately validate model quality, as the metrics in Table 2 fall significantly below those of publicly available models, including smaller ones. In general, altering model architecture for performance optimization but do not validate the final accuracy on large scale carry risks, as such changes may compromise final accuracy. The accuracy number of Snowflake Arctic [1] also shows that this parallel architecture might harm the model quality (when it compares with other models with similar size of parameters). While the paper is well-argued on the motivation and the system-side optimizations appear sound, I remain cautious about the final model quality of the new architecture.

(1) We propose the ScMoE architecture as a viable option for the backbone of MoE models. In line with the initial experiments conducted on previously proposed MoE architectures, such as standard MoE and shared-expert MoE, we train the ScMoE model from scratch to ensure a fair comparison. When trained under identical conditions, ScMoE demonstrates superior efficiency while achieving comparable accuracy, compared to both the standard MoE and shared-expert MoE methods.

(2) Due to our limited devices, we are unable to train with pre-training datasets comprising 1 trillion tokens, as is common with existing industrial-level models. Consequently, our experimental models exhibit relatively lower accuracy, as we only train with datasets that are 1/1000th the size. Nevertheless, we offer theoretical analysis upon our empirical results to support the effectiveness of our methods, as discussed in Section 5 and Appendix A.1.

(3) According to the provided document on Snowflake Arctic, this parallel architecture offers significant value for practical deployments, rather than solely impacting model quality negatively. The document states that "Arctic is more capable than other open source models trained with a similar compute budget."

1. The claim that model quality is always maintained or improved is only shown on small models. Sometimes the accuracy improvements are very small, and the results might not work for larger models (like OLMoE) or different tasks. The experiments are limited to only small models, which may limit the generalizability of the results to other models or tasks. The authors have only conducted experiments on small models for image classification and language tasks, so additional experiments on larger and more powerful models may be necessary.

(1) Given our constraints on available hardware, we encounter difficulties in conducting empirical experiments on large-scale, industrial-level MoE models. Nevertheless, we perform additional experiments on the OLMoE model [1] and conduct some different evaluation tasks. The detailed experimental configurations are introduced in our response to Reviewer GE4M. As shown in the following table, our proposed ScMoE architecture not only accelerates the training by 1.13x and inference by 1.19x, but also achieves a marginally higher average accuracy (43.77% vs 43.25%).

(2) Beyond the focus on accuracy, we also simulations to validate the effectiveness of speedup in larger-scale scenarios, as shown in the table "Simulation of larger-scale model training" in our response to Reviewer GE4M. Despite differences in communication and computation durations across configurations, our proposed ScMoE consistently proves to be effective.

2. It relies heavily on the assumption that features in neighboring layers are similar, which is often true. Yet, the supplementary materials show that this similarity is not clear in the OLMoE model. In addition, the authors mention that ScMoE-Pos2 does not perform well, which makes me worry about how well the method can work in different cases. Does the proposed method fail if the features change dramatically between neighboring layers?

(1) Recent studies [2,3,4,5] have commonly observed the high similarity between features in neighboring layers as a characteristic of Transformer-based LLMs, which can be utilized to improve the efficiency of LLM operations.

(2) Our supplementary materials on the OLMoE model further corroborate that the observed feature similarity is within an expected range and aligns with prevalent findings in the field. As illustrated in Figure 15, most data points near the diagonal line are yellow, indicating nearly 100% similarity. Only the output features of 1A and 1M exhibit comparatively lower similarity, yet maintain a significant 50%.

(3) We would like to clarify that the statement "ScMoE-Pos2 does not perform well" was made in a comparative context. In fact, the loss differences among ScMoE-Pos1 (3.2615), ScMoE-Pos2 (3.2818), and ScMoE-Pos2-L0Pos1 (3.2629) are relatively minor. Moreover, the loss values for the standard top-2 MoE (3.3157) and shared-expert MoE (3.2974) are higher than those of the three ScMoE configurations. This indicates that despite variations in feature similarity across certain layers, ScMoE remains effective. Additionally, we will include a detailed comparison of these loss curves in the appendix for further clarity.

(4) Our strategy for selecting shortcut-connected positions, detailed in Section 5.1.3, is designed to achieve optimal accuracy. Additionally, based on these observations of similarity, we can adapt this strategy for subsequent MoE model design by opting not to use ScMoE in the first two Transformer layers. Similarly, DeepSeek-MoE [6] chooses not to apply the MoE module in the first Transformer layer.

Reference

[1] Muennighoff, Niklas, et al. "Olmoe: Open mixture-of-experts language models." arXiv preprint arXiv:2409.02060 (2024).

[2] Sun, Qi, et al. "Transformer layers as painters." arXiv preprint arXiv:2407.09298 (2024).

[3] He, Shwai, et al. "What matters in transformers? not all attention is needed." arXiv preprint arXiv:2406.15786 (2024).

[4] Men, Xin, et al. "Shortgpt: Layers in large language models are more redundant than you expect." arXiv preprint arXiv:2403.03853 (2024).

[5] Gromov, Andrey, et al. "The unreasonable ineffectiveness of the deeper layers." arXiv preprint arXiv:2403.17887 (2024).

[6] Dai, Damai, et al. "Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts language models." arXiv preprint arXiv:2401.06066 (2024).

1. I think the analysis can be stronger if all of the the results are consolidated across different model sizes.

We have summarized and supplemented evaluations of GPT models across various sizes to better demonstrate the impact of model size as a standalone variable. The following table presents the speedup, zero-shot perplexity, peak GPU memory reduction, and MoE block latency reduction of ScMoE in comparison to the standard top-2 MoE. The results demonstrate that our proposed ScMoE consistently achieves improvements across various model sizes.

2. How do the validation loss results in Figure 10 compare with a regular MoE as the baseline?

We observed that the loss curve of ScMoE (3.2629 at the final point) demonstrates faster convergence and a lower final loss compared to the standard top-2 MoE (3.3157 at the final point) and shared-expert MoE (3.2974 at last).

3. How does the overall speedup and benchmark performance scale with model size? Do you expect that the ScMoE approach becomes more viable for larger scale pretraining?

(1) To illustrate the effectiveness of speedup in larger-scale scenarios, we conducted additional simulations, the results of which are detailed in the table "Simulation of larger-scale model training" in our response to Reviewer GE4M. The findings suggest that the relative ratio of communication to computation duration, influenced by various training and model configurations, is crucial for impacting speedup. Optimal speedup is achieved when these durations are balanced. Despite these variations, ScMoE consistently achieves acceleration.

(2) Due to our limited devices, we face challenges in conducting experiments on large-scale MoE models. Nevertheless, we offer theoretical analysis to support the performance of our methods, as discussed in Section 5 and Appendix A.1. Additionally, based on the findings presented in Question 1, we anticipate that the disparity in accuracy between ScMoE and existing MoE architectures will diminish as both model size and the volume of training data increase.

4. Can this shortcut connection approach be applied to already pretrained models? What does the performance look like in such cases? Does your approach necessarily require pretraining a model from scratch to use the shortcut connections.

(1) We propose the ScMoE architecture as a viable option for the backbone of MoE models. In line with the initial experiments conducted on previously proposed MoE architectures, such as standard MoE and shared-expert MoE, we train the ScMoE model from scratch to ensure a fair comparison.

(2) The ScMoE models can be constructed using pre-trained shared-expert MoE models by modifying the inputs of the MoE module. Our experiments with this approach showed a decrease in accuracy, indicating a need for further refinements in the fine-tuning process to preserve accuracy.

(3) Additionally, it is possible to construct and train a ScMoE model based on a pretrained dense model, employing techniques similar to sparse upcycling, which have been applied to standard MoE models.

5. The method seems limited to architectures that use shared experts and alternating MoE layers with typical MLPs. How do you think your results could generalize to other MoE models that either do not use shared experts, or have an MoE for every transformer layer?

(1) We would like to clarify that we have already conducted experiments on the current mainstream MoE structure of placing MoE module in every Transformer block, as presented by the LLaMA2-MoE experiments in Section 4.2.2, Table 2 and Table 7. Specifically, our LLaMA2-MoE experiments utilize ScMoE (Pos-1) and incorporates the MoE module in every Transformer block, which overlaps the communication and computation within the each Transformer block.

(2) Our work focuses on integrating shortcut connections into the MoE architecture to enhance efficiency, which does not rely on the shared-expert MoE. As discussed in Appendix A.2, our DGMoE, which incorporates shortcut connection without shared expert, also offers improvements over the top-2 MoE.

Furthermore, ScMoE functions alongside both the standard MoE and shared-expert MoE architectures. When designing models, practitioners have the option to choose from these three architectures.

(3) We conduct experiments on the OLMoE model, which originally employs standard MoE without shared experts at each layer. As shown in the table "Experimental results on integrating ScMoE architecture into the OLMoE model" in response to Reviewer GE4M, our proposed ScMoE architecture not only achieves acceleration, but also achieves a higher average accuracy.

1. Concerns regarding the model structure may be outdated, particularly in relation to the compatibility of ScMoE with the mainstream MoE architecture, where MoE layers are arranged sequentially rather than interspersed with MLP layers.

We would like to clarify that we have already conducted experiments on the current mainstream MoE structure of placing MoE module in every Transformer block, as presented by the LLaMA2-MoE experiments in Section 4.2.2, Table 2 and Table 7. Specifically, our LLaMA2-MoE experiments utilize ScMoE (Pos-1) and incorporates the MoE module in every Transformer block, which overlaps the communication and computation within the each Transformer block.

Additionally, we conduct experiments on OLMoE model, an advanced MoE model that integrates an MoE module at each layer. As shown in the table "Experimental results on integrating ScMoE architecture into the OLMoE Model" in our response to Reviewer GE4M, ScMoE architecture not only achieves acceleration, but also achieves a marginally higher average accuracy.

2. Concerns regarding stronger baselines.

We select a recently proposed All-to-All optimization technique, "MoE Shared Expert Overlap," which overlaps All-to-All communication with the computation of shared experts, as a stronger baseline for comparison.

As shown in the following table, ScMoE achieves acceleration compared to "MoE Shared Expert Overlap," facilitated by breaking dependencies for improved overlap.

3. Concerns regarding larger multi-node cluster analysis. Does a larger number of experts negatively impact ScMoE’s effectiveness?

To demonstrate scalability in large-scale clusters, we conducted simulations across different number of GPUs, from 16 to 128, as shown in the table "Simulation of larger-scale model training" in our response to Reviewer GE4M. Despite differences in communication and computation durations across configurations, ScMoE consistently demonstrates effectiveness. Notably, in entries No.5 to No.8, an increased number of experts does not negatively affect the effectiveness of ScMoE.

4. Incremental improvement: ScMoE mainly extends shared-expert MoE with shortcut connections.

Our work focuses on integrating shortcut connections into the MoE architecture to enhance efficiency, which does not rely on the shared-expert MoE. As discussed in Appendix A.2, our DGMoE, which incorporates shortcut connection without shared expert, also offers improvements over the top-2 MoE. ScMoE is presented as the most efficient architecture to facilitate our proposed shortcut connections.

Furthermore, ScMoE functions alongside both the standard MoE and shared-expert MoE architectures. When designing models, practitioners have the option to choose from these three architectures.

5. Are GPT-2/LLaMA2 MoE models trained from scratch or fine-tuned from a dense model?

The models are trained from scratch. We introduce the ScMoE architecture as a viable option for pre-training a MoE model from scratch, enhancing the efficiency of its pre-training, fine-tuning, and inference.

6. Are the experiments conducted on FP32 precision of BF16 precision?

BF16 precision.

7. Would it be possible to quantify and report exact percentages of communication hidden by computation in Table 1 or 2?

The exact percentages for Table 1 are directly reflected in Figure 7 (1), which hides 70% communication.

8. What happens when communication is much slower than computation? Can ScMoE still provide significant speedup?

As shown by the simulation results of No. 1 in Question 3, ScMoE achieves a significant speedup of 1.31x, even when the communication time is three times longer than the computation duration.

9. Do you observe any convergence differences between ScMoE and top-2 MoE during training? If so, does ScMoE require different hyperparameters (learning rate, warm-up schedule, etc.)?

We observed that the loss curve of ScMoE (3.2629 at the final point) demonstrates faster convergence and a lower final loss compared to the standard top-2 MoE (3.3157 at the final point). To ensure a fair comparison in our experiments, ScMoE demonstrates comparable accuracy under identical hyperparameters.

10. Have you tested transfer learning performance?

We conduct additional fine-tuning experiments on the trained OLMoE models discussed in Question 1, using 5,000 samples from Alpaca dataset. The results shows that ScMoE still achieves a higher average accuracy (44.88% vs 44.07%).

11. Is ScMoE scalable to larger MoE models (e.g., 13B, 30B, 70B, and even DeepSeek-V3 671B)?

Due to our limited devices, we face challenges in conducting experiments on large-scale MoE models. Nevertheless, we offer theoretical analysis to support the effectiveness of our methods, as discussed in Section 5 and Appendix A.1.

1. I think the accuracy on benchmark datasets are a bit low to provide meaningful comparisons. It'd be better to train bigger models or longer. It would be great if the idea can be validated at larger model scale and longer training runs. Currently the accuracy on benchmarks are a bit low.

Given our constraints on available hardware, we encounter difficulties in conducting empirical experiments on large-scale, industrial-level MoE models. Nevertheless, we perform additional experiments on the OLMoE model [1], an advanced MoE model that activates 1B of its total 7B parameters. We initiate training from scratch for both the original OLMoE and the OLMoE+ScMoE models using 40,000 samples from the TuluV3 dataset [2]. As illustrated in the following table, our proposed ScMoE architecture not only accelerates the training by 1.13x and inference by 1.19x, but also achieves a marginally higher average accuracy (43.77% vs 43.25%).

Experimental results on integrating ScMoE architecture into OLMoE model [1].

Moreover, we offer theoretical analysis upon our empirical findings to support the effectiveness of our methods. Specifically, we observe a correlation between the efficacy of ScMoE and the high similarity of intermediate representations. Given that such high similarity is commonly found in Transformer-based LLMs [3,4,5,6], we believe our ScMoE approach can be effective in larger-scale LLMs. Furthermore, our theoretical analysis of gradient propagation also demonstrates its ability to achieve consistent training while preserving model quality.

Additionally, we conduct simulations to validate the effectiveness of speedup in larger-scale scenarios (larger model and larger number of GPUs). Our simulations employ the ASTRA-SIM simulator [7], which models an A800 cluster in alignment with the real-world HPC cluster utilized for evaluations in Section 4. Specifically, the FP16 computational capability is configured to 312 TFLOPS, the intra-node NVLink bandwidth is set at 400 GB/s, and the inter-node network employs a 200Gb/s (25GB/s Bandwidth) InfiniBand connection.

We conduct simulations of the forward pass for LLaMA2-MoE, maintaining similar attention configurations to experiments in Table 2. Since the load balance of MoE significantly influences its efficiency and is determined by various inputs, we simulate a fully balanced distribution of input tokens to eliminate any random interference in this aspect.

In the simulation, we vary the configurations for the number of GPUs, the number of active and total experts, and expert intermediate size. As illustrated in the following table, while different configurations lead to varying communication and computation durations, our overlapping strategies consistently prove to be effective.

Simulation of larger-scale model training.

Reference

[1] Muennighoff, Niklas, et al. "Olmoe: Open mixture-of-experts language models." arXiv preprint arXiv:2409.02060 (2024).

[2] Lambert, Nathan, et al. "Tulu 3: Pushing frontiers in open language model post-training." arXiv preprint arXiv:2411.15124 (2024).

[3] Sun, Qi, et al. "Transformer layers as painters." arXiv preprint arXiv:2407.09298 (2024).

[4] He, Shwai, et al. "What matters in transformers? not all attention is needed." arXiv preprint arXiv:2406.15786 (2024).

[5] Men, Xin, et al. "Shortgpt: Layers in large language models are more redundant than you expect." arXiv preprint arXiv:2403.03853 (2024).

[6] Gromov, Andrey, et al. "The unreasonable ineffectiveness of the deeper layers." arXiv preprint arXiv:2403.17887 (2024).

[7] Rashidi, Saeed, et al. "Astra-sim: Enabling sw/hw co-design exploration for distributed dl training platforms." ISPASS. IEEE (2020).