Occult: Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference

We thank reviewer NL62 for the dedicated and professional comments. To address your concerns, we provide detailed pointwise responses below:

[Claims and Evidence]

We provide code at https://anonymous.4open.science/r/Occult-D802.

[Theoretical claim 1: Communication complexity in Fig. 2]

Thank you for the careful review. We admit there's a typo in Fig. 2(c): the blue token routed to Expert 1 and the yellow one routed to Expert 2 should be exchanged due to our intentional modification of routing choices to ensure tokens only activate experts on the same device, thereby reducing all-to-all communication volume.

To clarify the communication complexity in Fig. 2:

• Fig. 2(a): Each token is repeated twice, yielding $C_{\mathcal{T}}$ = 2
• Fig. 2(b): Red and green tokens are repeated once while yellow and blue tokens are twice, yielding $C_{\mathcal{T}}$ = 1.5
• Fig. 2(c): Each token is repeated once, yielding $C_{\mathcal{T}}$ = 1

The confusion may stem from insufficient emphasis on top-k value. This example uses 2 devices, 4 tokens, 4 experts, and top-2 routing. According to our derivation in Section 3.2, communication complexity bounds are $1\leq C_{\mathcal{T}}\leq 2$, with all 3 cases in Fig. 2 falling in this interval. We will revise Fig. 2 and clearly explain this modification in Fig. 2(c) to make it more comprehensible.

[Experiments & Analysis 1: Multi-node Training]

Our research is targeted at communication efficiency, so it inherently performs better than conventional expert parallelism on multi-node training, where inter-node communication is the main bottleneck. We further examine 8-way (1 node ) and 16-way ( 2 nodes ) expert-parallelized training for DeepSeek-MoE, where latency and evaluation comparison are also visualized at https://anonymous.4open.science/r/Occult-D802.

[Experiments & Analysis 2: Fine-tuning baseline for fair comparison]

In Sec. 5.2 and Fig. 7, we provide the evaluation results for

• Original model (no tuning), shown in yellow dashed line
• Tuned model with pruning, shown in brown dots ( similarity-based pruning ) and green stars ( router-based pruning ).
• Standard tuning, shown in pink diamond.

In the caption text of Fig. 7, we indicated that pruning within 4 devices is equivalent to standard tuning since we use 4 devices for distributed training. Therefore, the effect of pruning can be derived by comparing the brown dots & green stars with pink diamonds. The yellow line serves as a reference.

[Strengths and weaknesses 1: Fig. 2]

Thank you for the feedback. We will improve the caption to clarify the meaning of $D_i$, $D_i^j$, and $E_i$, and add explicit in-text references to better guide readers through this illustration of different communication strategies.

[Strengths and weaknesses 2: Communication complexity]

We use communication complexity to indicate the ratio of all-to-all communication volume to the number of tokens, approximated by the average token replication count. This approach addresses two key challenges:

• During all-to-all communication, a token replica may either be retained locally or transmitted to other devices based on its routing choice, posing uncertainty in precisely measuring inter-device communication volume across different tasks.
• To establish a more stable metric, we use average token replication times, which is also the least upper bound for the ratio of inter-device communication volume to token count, minimizing the impact of dynamic routing uncertainties.

We appreciate you highlighting this concern and will expand these explanations in Section 3.2 to provide better clarity.

[Strengths and weaknesses 3: Figure 4]

We apologize for the missing in-text reference to Fig. 4. We will add detailed descriptions, including:

• $`SFD`$ tokens serve as the all-to-all content, constructed from $`ORI`$ tokens based on $BRIM_0$
• $`FFN`$ utilizes $BRIM_1$ as an auxiliary input to guide the token layout of the output tensor
• Intermediate tokens are organized densely in the $`EPD`$ state

I hope they can help readers better understand the token state transitions and the role of $BRIM$s in our framework.

[Other comments]

Thanks for pointing out. We will fix these typos carefully in our revised draft.

[Question: Huggingface Latency]

Prefilling latency is the primary bottleneck when generating a small number of tokens. Hugging Face’s native API employs pipeline parallelism (PP), which outperforms expert parallelism (EP) under the following conditions:

• Limited GPU amount ( our experiments use only 4 GPUs )
• Deep model architectures ( the three evaluated models contain 16, 24, and 28 layers, respectively )

Thereby, the communication volume of EP in the prefilling stage is much larger than PP. However, the decoding speed of EP is much faster, profit by the efficient MoE computation.

Thank you again for your comments. We will incorporate these refinements in our revised draft.

We thank Reviewer 1GwW for recognizing that "the paper addresses a key limitation in MoE scalability", “experiments generally robust”, and that "Occult provides a practical, well-validated optimization." To address your questions, we provide pointwise responses below.

[Potential concerns 1: Different task types]

We've evaluated the performance on 23 benchmarks with different tuning strategies in Tab. 3,4,5 on pages 16, 17, and 18, covering natural language understanding, commonsense reasoning, and math. Besides, we conduct additional experiments on HumanEval (coding) and GSM8K (math) datasets for DeepSeek-MoE. The results below consistently demonstrate the effectiveness of our efficient pruning methods:

HumanEval:

GSM8K ( flexible-extract ):

[Potential concerns 2: Scalability]

Our occult is more feasible for modern fine-grained MoE-LLMs such as DeepSeek-MoE. Sadly, models with a scale larger than 16B are out of our capability to train ( such as DeepSeek-V2 ). We conduct additional experiments on 8 GPUs and 2 x 8 GPUs (two nodes) to demonstrate the scalability of our method, using DeepSeek-MoE with 8- and 16-way expert parallelism, taking batch size 32 per GPU:

Avg training latency per step (s):

Occult acceleration can be more apparent on a better-grouped expert placement, thereby Occult can better improve training efficiency on 8-way expert parallelism.

We also evaluate 8- and 16-way EP on MMLU and MathQA:

8-way EP:

16-way EP:

[Strengths and Weaknesses 2: Collaboration pruning trade-offs]

We've analyzed the impact of different pruning settings on both performance and efficiency in Fig. 7, 8, 10, and 11. A more comprehensive accuracy analysis is provided in Tab. 3,4,5 on page 16.17.18. Additionally, we conduct an additional ablation study to analyze the pruning impact on efficiency with DeepSeek, using 8-way EP:

[Comments 1: To multi-node setting]

Thanks for the practical question. To scale up our method to the multi-node settings, we further conduct extra experiments on multi-node settings, as shown in [Potential concerns 2: Scalability], and demonstrate the scalability of our method. To scale it to a very large cluster, a combined parallel strategy for the MoE layer is required to adapt to the hardware resources, including:

• Data parallelism is required to repeat the basic unit of expert parallelism (e.g., 16 GPUs in 16-way EP ), aiming at avoiding heavy communication overhead across nodes ( e.g., repeating the parameter of each MoE layer for 4 times in a 64-GPU cluster )
• Expert parallelism can be organized across different nodes, placing the grouped experts in Occult on the same node
• Tensor parallelism can be adopted in intra-node GPUs for very large expert since the intra-node bandwidth is usually abundant

[Comments 2: RL-based expert assignment]

End-to-end rescheduling may be impracticable, as dynamically tuning expert placement requires huge additional bandwidth, although it can be asynchronous.

While our current heuristic approach yields strong results, we believe our expert assignment algorithms can be further improved for enhanced performance. RL-based methods can learn the collaborative communication patterns from the profiling dataset, which can benefit the expert assignment task and improve generalization. We have not tried similar methods yet, but it is promising for our future research.

Thank you again for your questions. We will incorporate the refinements in our revised draft.

We thank tR7y for recognizing that "the motivation of the paper is clear" and that "experiments show that the proposed method achieves faster speed and higher performance compared to baseline methods." To address your questions, we provide pointwise responses below.

[Experiments & Analysis 1: Larger MoE and fewer-expert models]

Occult is more feasible for fine-grained MoE-LLMs such as DeepSeek-MoE. 16B is a common choice, and larger models usually contain hundreds of billion parameters ( such as DeepSeek-V2 ), which is beyond our capability.

We've conducted additional experiments on Mixtral (8 experts) to demonstrate generalizability to models with fewer experts. Limited to GPU memory, we only tune the last 2 layers with 80 ( bs ) $\times$ 128 ( seq length ) tokens with 4 GPUs ( 2 experts per GPU ). To illustrate the scalability, we also performed training with 8 GPUs on DeepSeek-MoE, showing that the communication savings can scale with device count for expert parallelism ( EP ). We tune all the MoE layers with 32 ( bs ) $\times$ 128 ( seq length ) tokens.

We also provide the evaluation results for 8-way expert parallelized DeepSeek-MoE with Occult to demonstrate its effectiveness:

[Experiments & Analysis 2: Performance improvement explanation]

The motivation for our proposed pruning methods is to replace some of the original routed experts with carefully assigned alternatives based on certain rules, so that the modified routing choices can be more grouped within some GPUs rather than dispersed, thereby reducing the all-to-all communication overhead. Visualizations in Fig. 8 demonstrate that two-device pruning preserves more essential collaboration patterns found in the original model, while single-device pruning might lose important expert correlations. Our router- and similarity-based pruning approaches are capable of maintaining the most critical expert collaborations, helping the model retain or even enhance its capabilities as a kind of regularization [1, 2]. This targeted preservation of collaboration patterns explains the performance improvements shown in Tab. 3 and Fig. 7.

[Question: Difference with existing method for MoE]

Our framework differs from existing MoE libraries in several fundamental ways, making it more communication-efficient:

• Novel data structure for token management: As described in Sec. 4.1, we introduce "Bidirectional Re-Index Matrix ($BRIM$), a novel data structure for unified data management" which efficiently tracks token states across different processing stages. Unlike existing methods that use general-purpose data structures, $BRIM$ is specially designed for MoE workflow with optimized memory access patterns.
• State-based token representation: As stated in lines 206-216, "we outline them as 3 states: Original ($`ORI`$)... Simplified ($`SFD`$)... Expanded ($`EPD`$)," allowing us to maintain tokens in "the token counts across states follow $`ORI`<`SFD`<`EPD`$." This contrasts with existing methods that simply replicate $k$ times for each token, regardless of collaboration patterns.
• Two-stage aggregation: As described in lines 199-203, we implement "summing the intra-collaboration results before all-to-all, and summing the inter-collaboration results after all-to-all, for each token $x$." This two-stage token aggregation enables symmetrical efficient all-to-all communication for both dispatch and combine operations in MoE pipeline.

Thank you again for your questions. We will emphasize these distinctions in our revised draft.

[1] Giles, C.L. and Omlin, C.W., 1994. Pruning recurrent neural networks for improved generalization performance. IEEE Transactions on neural networks, 5(5), pp.848-851.

[2] Han, S., Mao, H. and Dally, W.J., 2015. Deep compression: Compressing deep neural networks with pruning, trained quantization and Huffman coding. arXiv preprint arXiv:1510.00149.

We sincerely thank reviewer geV9 for recognizing that our approach "enhances the training and inference efficiency of MoE models, which is crucial for future LLM deployment." To address your questions, we provide pointwise responses below.

[Experiments & Analysis 1: Model configuration]

Thanks for the constructive feedback. We've added a supplementary table detailing model configurations:

[Experiments & Analysis 2: Pruning strategies]

The performance differences between router- and similarity-based pruning arise from how they handle expert replacement:

• Router-based pruning uses only routing scores, which may not involve expert attribute
• Similarity-based pruning considers inter-expert similarity, obtained from a profiling dataset

As shown in Figure 7, similarity-based pruning achieves comparable or better performance than router-based approaches in most cases of models and tasks. For example, in the OLMoE results on RTE, similarity-based pruning (shown as "Pruning (Similarity-based)") achieves approximately 3% higher accuracy than router-based pruning when restricting collaboration to 1 GPU.

Similarity-based pruning can utilize prior knowledge from profiling datasets, which may advance the downstream tasks under similar domain distribution. However, it would expire on out-of-distribution tasks, where router-based pruning may perform better.

[Experiments & Analysis 3: throughput, TTFT, and TPOT]

Following your valuable suggestion, we will adopt “throughput (tokens/s), Time To First Token (TTFT), and Time Per Output Token (TPOT)” in Fig. 9 & 10 under varying batch sizes and sequence lengths in the formats below:

Prefilling:

Decoding:

[Strengths And Weaknesses 1: Cost of pruning]

Pruning does not require more time to converge because

• Results in Fig. 7 experiments are all reported with 1 epoch training on the Alpaca dataset, on some tasks it can even perform better than standard SFT.
• The average latency per training step can be greatly reduced with Occult.

We report the latency and memory cost of pruning for DeepSeek-MoE here, with batch size 16 and 8 GPUs:

Our prune algorithms can accelerate clock-time latency and reduce memory costs for training.

[Strengths And Weaknesses 2: Bandwidth impact]

We conducted additional experiments with varying interconnect bandwidths on different machines. The acceleration of Occult is more significant with lower bandwidth. We use DeepSeek-MoE with 8-way EP and batch size 32:

[Strengths And Weaknesses 3: Scalability]

Our approach can remain effective at a very large scale with a hierarchical strategy:

• Intra-node optimization: Within each node (typically 8 GPUs connected via NVLink), we can place multiple experts on the same node to maintain high intra-collaboration within the node boundary, where communication is fast.
• Inter-node optimization: Cross nodes, expert placement rescheduling & collaboration pruning can be applied to minimize cross-node communication, which is typically the latency bottleneck.

This scaling profile can be beneficial since inter-node communication is often 2-10× slower than intra-node communication, making the reduction of cross-node traffic especially valuable. In large-scale deployments, Occult's benefits should increase rather than diminish as communication overhead becomes more dominant.

[Comments: Notation issue]

Thanks for the constructive suggestion on “faster” and “speedup”, we will fix it in the revised paper.

[Question 3: Integration with DeepEP]

DeepEP mainly contains 2 stages for hardware-tailored MoE communication:

• Inter-node all-to-all: sending tokens to the corresponding GPU on the target node
• Intra-node all-to-all: sending tokens to the target GPUs containing the target expert

In our understanding, Occult is orthogonal to DeepEP since it optimizes all-to-all communication volume along the dimension of top-k aggregation, which can be combined with DeepEP for enhanced efficiency.

Thank you again for your questions. We will include these additional experiments and analyses in our revised draft.