Occult: Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
Authors: Shuqing Luo
Pingzhi Li
Jie Peng
Yang Zhao
Yu Cao
Yu Cheng
Tianlong Chen
Venue: ICML 2025 poster (2025)
Keywords: Mixture-of-Experts
Communication Efficiency
Expert Parallelism
Primary area: deep_learning->large_language_models
TL;DR: Optimizing all-to-all communication in expert parallelism using algorithm-system co-design

Abstract: Mixture-of-experts (MoE) architectures could achieve impressive computational efficiency with expert parallelism, which relies heavily on all-to-all communication across devices. Unfortunately, such communication overhead typically constitutes a significant portion of the total runtime, hampering the scalability of distributed training and inference for modern MoE models (consuming over 40% runtime in large-scale training). In this paper, we first define $\textit{collaborative communication}$ to illustrate this intrinsic limitation, and then propose system- and algorithm-level innovations to reduce communication costs. Specifically, given a pair of experts co-activated by one token, we call them as $\textit{collaborated}$, which comprises $2$ cases as $\textit{intra-}$ and $\textit{inter-collaboration}$, depending on whether they are kept on the same device. Our pilot investigations reveal that augmenting the proportion of intra-collaboration can accelerate expert parallel at scale. It motivates us to strategically $\underline{\texttt{o}}$ptimize $\underline{\texttt{c}}$ollaborative $\underline{\texttt{c}}$omm$\underline{\texttt{u}}$nication for acce$\underline{\texttt{l}}$era$\underline{\texttt{t}}$ed MoE training and inference, dubbed $\textbf{\texttt{Occult}}$. Our designs are capable of $\underline{either}$ delivering exact results with reduced communication cost, $\underline{or}$ controllably minimizing the cost with collaboration pruning, materialized by modified fine-tuning. Comprehensive experiments on various MoE-LLMs demonstrate that $\texttt{Occult}$ can be faster than popular state-of-the-art inference or training frameworks (over 50% speed up across multiple tasks and models) with comparable or superior quality compared to the standard fine-tuning. Codes will be available upon acceptance.

Occult: Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
Shuqing Luo
Pingzhi Li
Jie Peng
Yang (Katie) Zhao? 2 Yu (Kevin) Cao Yu Cheng Tianlong Chen
Abstract
Prefiling Latency Comparison (16384 Tokens)
GLM
Owan-Nce
DespSrelette
Lalenoy
Mixture-of-cxperts (MoE) architectures could achieve impressive computational efficiency with expert parallelism, which relies heavily on all-to- all communication across devices. Unfortunately, such communication overhead typically consti- tutes a significant portion of the total runtime, hampering the scalability distributed training and inference for modern MoE models (consum- ing over 40% runtime in large-scale training). In this paper, we first define collaborative commu- nication to illustrate this intrinsic limitation, and then propose system- and algorithm-level innova- tions to reduce communication costs. Specifically, given pair of experts co-activated by one token, call them collaborated, which comprises 2 cases intra- and inter-collaboration depending whether they are kept the same device. Our pilot investigations reveal that augmenting the proportion of cntaacoollaooration can accelerate expert parallel at scale. motivates us to strate- gically optimize collaborative communication for accelerated MoE training and inference, dubbed Occult. Our designs capable of either delivering exact results with reduced com- munication cost, or controllably minimizing the cost with collaboration pruning, materialized by modificd fine-tuning. Comprehensive experi- ments on various MoE-LLMs demonstrate that Occult be faster than popular state-of-the- art inference training frameworks (more than 1.5x speed up across multiple tasks and models) with comparable or superior quality compared to the standard fine-tuning. Code is available at http:///ttuub.com/UUTEEEE---LLs
TUM
Desaat
Rlgedd SCO
(No Frune) (Prun GPUssPrane Gru
Deooding Lalancy Comparison (12800 Frompt Tokans, Tokeens, 512)
Laleroy (6)
111
curSC
DeepSnel-ee
Hugorg
(Pack Srarso Outtrr
Regedi
[No Pruns)
OPUsgPrank OPUN
Latancy Comparison (Batch Seq. Length
per
TSLP
DeepSpeed
Regablede Ragasledoo
Orted CPU
(CPU Orioad)
(S7)
Sparna) (Grouped)
(Ns Prune)
Figure Latency Comparison with Multiple Models Tasks. Occult can accelerate training & inference for modern MoE. LLMS on communication-imtnnn tasks.
1. Introduction
Transformer-based (Vaswani, 2017) large language mod- els (LLMs) (Achiam et al., 2023; Touvron et al., 2023) have demonstrated tremendous success in broad spectrum of downstream tasks, and scaling up model parameters the de facto approach to train more powerful LLMs (Jiang et al., 2024; Liu al., 2024b). Mixture-of of-uue----- (MoE) empow- L.MMs with efficient parameter scaling via sparsity. replaces the standard feed-forward network with group of experts, and each token only activates subset, which significantly improves computation efficiency.
However, the giant scale of MoE parameters and the limited memory a single GPU make it necessary to shard an MoE layer across multiple GPUs for parallelized training and in- ference. Among the common techniques, expert parallelism is scalable and efficient approach (Lepikhin et al., 2021; Fe- dus et al., 2022), which distributes experts multiple GPUs and dispatches tokens to the corresponding devices via all- to-all communication. As illustrated Figure 2(a), we can outline a typical forward pass as Dispatch + All-to-All + Computing All-to-All - Combine.
'The University of North Carolina Chapel Hill PUniversity of Minnesota Twin Cities 'The Chinese University Hong Kong. Correspondence to: Tianlong Chen <tianlong @ es edu.e
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
Previous work has demonstrated that all-to-all communica- tion is the main bottleneck (Li et al., 2023a) for both training and inference, especially under heavy workloads (Hwang et al., 2023), where it can count for over 40% runtime. The reason is that communication across GPUs is much slower than computation on a single one. GPUs on the same ma- chine are typically connected via PCIe or NVLink, where the former provides low bandwidth, and the latter is higher but only available on some cutting-edge devices. This makes communication efficiency pivotal subject for scalable MoE deployment. For modern L.Mss, heavy workload is com- mon teaser, usually emerging in the following tasks:
Table I. Approximated Linear Correlation between Cj and runtime. 2 prompt tokens the prefilling stage are examined.
OLMel
MegaBlocks
iccuit Trivial
Geouit Rescheduled
Cocult Pruns, GPUs
Pranc. GPUs
Occult
Cy
E(Cr)
3,88
8.022
1,88
Intra/ / Inter
0.33/0.67
0.33/0.67
0.46/0.54
0/60/0.00
1/0
Laiesey (4)
24.50
15.33
12.63
9:22
5.82
For intra collaboration, only one replica of x is re- quired in the Dispatch stage. For Combine stage the output of E, and E3 can be pre-aggregated locally.
* For inter-collaboration, typically the token should be duplicated twice. But it can be transformed into intra- collaboration, only one replica would be required.
Training: Large batch can better utilize hardware re- sources, accelerating both pre-training and fine-tuning.
Long Sequence Prompting (Prefilling Stage of Infer- ence): A single forward pass can batch massive tokens.
Therefore, the communication cost for expert parallelism can be optimized via maximizing intra-collaboration and minimizing inter-collaboration To formulate it, we first propose a simple and practical criterion to measure the com- munication budget, i.e., the average number of replicas for each token, denoted as CT. In Table 1, we present the strong linear correlation between Cy and runtime. To push the effi- ciency of expert parallel to the limit, Occult t first resched- ules expert placement leveraging profiled collaboration. We utilize the statistics of intra- and inter-collaboration from a profiling dataset to derive a rescheduled placement, which delivers around 20% reduction on CT and runtime, as shown in Table 1. Furthermore, we propose collabora- tion pruning for controllable and minimized Cy. which pushes communication efficiency to the limit controllably.
Concurrent Generation Requests (Decoding Stage of Inference): Requests from different users are batched together, and decoding is implemented step by step.
The above motivates us to explore how to optimize commu- nication across devices to achieve lower latency and higher throughput. Previous MoE libraries (Hwang et al., 2023; Gale et al., 2023) typically make & replicas for each token in expert parallelism, and dispatch them to different devices, constituting the content of all-to-all communication, In fact, not all content is necessary: if subset of activated experts for one token is kept on the same device, then only one replica would be required for this device. In the ideal case, if the activated experts are all kept on the same device, then the allll cost would turn out to be almost negligible.
As algorithm-system co-design scheme, Occult can ac- celerate MoF. training and inference with expert parallelism on communication-intent tasks, Our contributions are:
To implement this efficient communication, we develop a novel sparse matrix multiplication kernel tailored for ef- ficient all-to-all. Note that if only reusing previous libraries, the streamlined token replicas for efficient communication cannot directly serve as the input/output of expert comput- ing, otherwise extra memory footprint would be required. Therefore, we build a new sparse matrix multiplication ker- nel for both forward and backward passes in Section 3. aiming at reducing unnecessary memory allocation or ac- be critical
* Novel Perspective for Efficient Expert Parallel. We propose to view the all-to-all communication in ex- pert parallelized MoE workflow from novel collab- orative perspective, and propose an algorithm-system co-design scheme dubbed Occult to accelcrate both training and inference for MoE-LLMs.
cess. Based this, expert placement turns for communication efficiency: an ideal placement should make the routed experts for each token fall into as few de- vices as possible.
* Higher Quality than Top-k Routing. We evaluate the performance of MoE-LLMs tuned with collaboration pruning extensive benchmarks, where Occu 1 t can achieve comparable better performance than top-k routing with improved computational efficiency.
To this end, we propose to reformulate all-to-all in MoE workflow as Collaborative Communication. Given expert E, and F, co-activated by token ær, we denote them col- laborated. Furthermore, expert collaboration can be divided into two categories: (1) Intra Collaboration; if E; and Ej are kept on the same device, and (2) Inter Collaboration: if E, and E.j are kept on different devices.
* Faster Training and Inference. Occult t achieves wall-clock speedup for training and inference with MoE-LLMs on communication-intensii tasks. 3 fron- tier models and 3 tasks, including prefilling, decoding, and training, are examined, with outlined compari- son in Figure 1. Occult consistently surpasses other libraries and popular frameworks to varying degrees.
This collaborative perspective inspires us with innovative solutions to advance MoE communication efficiency:
2

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
GPU-6
Dg
Dj
Eg
DE Al
Combne
Un
2.2. Measuring All-to-All Communication Complexity
DF
E2
We propose to measure the all-to-all communication com- plexity with CT, motivated by 2 aspects:
GPU-
Combino
D:
In the Dispatch stage, typically replicas are prepared for each token, serving as the all-to-all content. An in- dividual replica may be reserved locally or transmitted to another GPU, based on the routing choices.
(a) Classical MoE workflow
(CT 2).
pg AV ng
bg Al- Ao og
Caneiro
Do
2. Exactly measuring the communication volume Icads to varied evaluations on different datasets. CT can tackle this issue, since it is the supremum of inter-device all- to-all communication volume, making it task aask-agnostic.
Dj
DP
Combioo
D2
D
As shown in Table 1, CT exhibits a strong linear correla- tion with wall-clock runtime, since highly related to communication overhead and memory footprint. Based on this empirical relationship, we formulate the boundaries of the all-to-all communication budget: Consider distributed system with N. experts and Nu devices implementing top-k routing, where Ne is divisible by Na. We define CT and CT as the lower and upper bounds of CT respectively, where:
(b) Occult workflow w/o collaborative pruning (Cr 1.5).
GPt-
/AHr-i ng
ANtove dg
GPU-1
Combre
k Na I S CT S min(k: Na).
(c) Occult workflow w. collaborative pruning (C7 = 1).
1 1
N.
(5)
Figure 2. MoF workflows with 3 all-to-all communication strategies. We take devices (Dp and D1) for expert parallel, and D/ denotes tokens from device i dispatched to device j.
Therefore the theoretically optimal budget for all-to-all com- munication be formulated Cy
N
2.3. Optimizing All-to-All Communication Complexity In Figure 2, we progressively optimize CT to improve com- munication efficiency of expert parallelism. Figure 2(a) shows the conventional expert parallelism. A key obser- vation is that not all tokens require k replications: the red and green tokens necessitate only single copy, since both replicas are transmitted to the same device. Figure 2(b) de- picts how leveraging this insight reduces the communication cost Cf from 2 to 1.5. Nonetheless, the yellow and blue tokens still incur duplicate transmissions. Through strategic optimization of the routing policy, as shown in Figure 2(C), we can achieve further reduction in CT to 1, i.e. reaching the theoretically minimum communication overhead.
2. Preliminary
2.1. Mixture-of-Experte and Expert Collaboration
Given an input token embedding x, the output of an MoE layer is the weighted sum of outputs from the No experts
N,-1 Σ R(x).
MoF(x:)
: Fx()),
(1)
1=0
where R(x); is the output of router network R(-) for the i-th expert. For each token, the MoE layer aggregates the output of & experts based on the top-k highest scores obtained from the Softmax value of a gating function g(-), which is usually a single linear projection layer:
2.4. Why Collaboration Pruning?
R(x)) Top-K K (Softma(((()))) k),
To minimize the communication complexity Cr, two poten- tial approaches emerge:
is in the top k,
Top-K(v,k) 0, otherwise.
Pure system design: Dynamically optimize expert
placement for each mini-batch reduce Cy.
To quantify the interactions among experts, we construct graph C based on the expert co-activation patterns within token batch B. Each edge represents the co-activation times between 2 co-activated experts:
Algorithm-syste system co-design: Adapt the routing policy to align with pre-defined expert placement, thereby reducing C7 while maintaining model performance. Our empirical analysis demonstrates the unfcasibility of the pure system design approach. For a collaboration graph constructed from a mini-batch B, we analyze the number of experts in the maximal connected subgraph the number of tokens in B increases. Figure 3 shows this analysis across
C,, = Σ 1(R(x); + (OAR()), 0}.
(3)
FES
We further normalize all edge values in C by dividing by the maximum edge value, yielding matrix P E N.N. with values in the interval [0. 1): Pij = C2.z/
'The lower bound of CT is the numher of devices that A cover, which equals to [k/(NN/Na) = K..8
max Ci.j- 05i45N.-1
(4)
N.
3

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
0 Madrnal Coroected Suegene Buegrphh
Top-2 rouling. 16 experts
Top-4 rouling, 8D
Top-8 muding, 64 experti
Top-8 muding, 84
100%
75%
75%
75%
-
59%
Total Experts Phi3..5MMoE-nnttue
Total Experis
25%
Total Esperts: DeegSeek-MoE
Total Experte
5-MoE-A2
in Min datch
Tokens Batch
Tokers Nin Batthh
tokene Mini Batch
Figure 3. Empirical correlation between the expert amount in the maximal connected sub-graph and token amount in a mini-batch for top-2, 4, 6, and 8 routing. We examine the last MoR layer for each model with the same inputs.
four MoE-LLMs with top-2, 4, 6, and 8 routing. The results show that the maximal connected subgraph converges to the complete graph as the mini-batch token count increases, implying that optimal expert partitioning across distributed devices for minimal CT is practically unattainable. There- fore, we adopt algorithm-systen co-design approach to
* Simplified (SFD): The communication tensors for all- to-all operations, serving as the input of the first linear layer and the output of the second linear layer. Each token is replicated fewer than & times.
* Expanded (EPD): The intermediate token tensors be- tween the first and second linear projection layer, where each is replicated exactly &o times.
minimize the communication budget.
The token counts across states follow ORI < SFD < EPD. All token tensors are maintained continuously in GPU HBM. Each expert computing operator bridges two states for mem- ory efficiency, taking one as input and producing another as output. However, direct implementation across these states presents challenges. State transitions require either token replication or reduction during computation on SRAM, ne- cessitating token indexing. Moreover, this parallel token indexing process is still complex, due to the simultaneous consideration of devices, experts, and tokens.
3. Method
In this section, we introduce Occu 1 t, an algorithm-system co-design approach that accelerates the MoE pipeline with expert parallelism by optimizing collaborative communi- cation. Section 3.1 introduces our tailored sparse matrix multiplication (SMM) implementation for efficient all-to- all communication. Section 3.2 describes the algorithmic design for expert placement rescheduling. which critically impacts the communication budget. Section 3.3 presents our collaboration pruning schemes that push CT to the limit.
To tackle this teaser, we introduce Bidirectional Re- Index Matrix (BRIM), a novel data structure for unified data management. Each MoE layer requires 2 BRIMs, as pre- sented in Figure 6. A BRIM provides two functions, namely Scattering and Merging, which support basic MoE opera- tions including Dispatch, Combine, and expert computing. Each function implements a two-stage process, aligned with our two-stage aggregation for Equation 1:
3.1. Sparse Matrix Multiplication for Efficient All-to-All
Our system design originates from the key insights for effi- cient all-to-all communication. Consider 2 experts E, and E, that are co-activated by token æ; where the expert com- puting results of them must be aggregated along with hs 2 other experts for x. When F, and F. are intra-collaborated, the forward MoE pass exhibits the following properties:
Merging: Stage one (EPD + SFD) integrates with SMM, while stage two (SFD + ORI) integrates with Dispatch & Combine. As shown in the left part of Figure 5, merging operates along the 1g dim of BRIM.
The Dispatch stage requires only single replica of æ for all-to-all communication.
* Scattering: Stage one (ORI + SFD) integrates with Dispatch & Combine. while stage two (SFD + EPD) integrates with SMM. As shown in the right part of Fig- ure 5, scattering operates along the 2nd dim of BRIM. We implement the basic MoE operations in Occult with Triton (Tillet et al., 2019). For BRIM-based SMM operators, each threud-block is assigned sub-vector with a fixed num- ber of elements in a single BRIM, row to enable tiled matrix multiplication." 2 The expert can be determined by the sub- vector's 's x z-coordinate, ensuring that each thread- -lock loads
The Combine stage allows pre-aggregation of expert computing results from E, and E,, maintaining consis- tency with Dispatch.
To implement this symmetric and efficient communica- tion, the aggregation in Equation 1 should be split into two stages, i.e., summing the intra-collaboration results before all-to-all, and summing the collabrollaoo results after all-to-all, for each token æ.
To implement these ideas and enhance efficiency, we need to deal with multiple tensor states, since tokens are selectively repeated. In this paper, we outline them as states:
Siine BRIM is very sparse, directly splitting it into vectors with fixed length would cause decrease in GPU utilization. There- fore, we further condense the obtained vectors from BRIM re- duce the total number of thread-blocks during runtime.
Original (ORI): The raw input and output tensors for an MoE layer, where each token appears exactly once.

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
GPU-0
Dispatch
All-to-All
Intermediate
All-to-All
Combine
0
Input
0
1
0
0
Output
0
1
0 2
MLP-0
2
0 X
MLP-1
1
0 1 X
2
2
2
X
3
2
X
2
2
2
0
0
0 1
0
X
MLP-0
X
MLP-1
2
2
2
3
ORI
BRIM,
SFD
SFD
BRIM;
EPD
BRIM;
SFD
SFD
BRIM,
ORI
Figure 4. MoE workflow with proposed bidirectional re-index-guided all-to-all communication and expert computing, We take the forward process for illustration, where token tensors stated ORI, SFD and EPD, compactly stored HBM and guided by BRIMs.
ORI
BRIM
Device-0
SFD
Algorithm Building expert placement.
Require: Collaboration graph E RRNxx,,, number of devices Na for expert parallelism.
2
2
Initialize expert placement & with N₂ empty lists.
Merging-1
Scattering-O
for d <- 0, Na- do
Merging (2nd stage) and Scattering (I, stage) with BRIMo
if
then Choose the most collaborative experts.
i argmax(P) and push i, to Lid-
SFD
BRIM,
Expert-0
EPD
else b Expert with the least collaboration with used ones.
Initialize Cineer as empty dict.
0
X
1
for e- 0, Ne -1do do
2
2
if not in C then Traverse the unused experts.
3
Set Cimer as 2] Eiec P(re)-
Merging-O
Scattering-1 1
end if
(b) Merging (Iu stage) and Scattering (2nu stage) with BRIM,
end for Expert with least collaboration with used ones. Index the minimal value from Caser, and push it to Li4- end if Initialize Cia] with 1 or 2 experts.
Figure 5. Functions of 2 BRIMs. Both BRIMa and BRIM; are 2D matrices. BRIM,'s 's dimensions defined by the number of devices for expert parallelism (first dimension) and the token count the CRI tensor (sccond dimension). BRIM, dimensions are defined by the number of local experts per device (first dimension) the token count in the SFD tensor (second dimension).
while len (en() S Ne/Na do Progressive expert selection.
Initialize Cauaa as an empty dict.
for + 0, N₁
do
if e not in C then Traverse the unused experts.
Set Canaa [e] as Zrjj
1 >>Ec(a) Plo.e)
end if Collaboration with experts on device end for Expert with the most collaboration with Diaj- Index the maximal value from Cisra,, and push it to Li4-
weights from one expert. Merging and scattering functions primarily differ in their I/O patterns:
end while end for return L.
For merging (first stage): Each thread-block reads from the EPD tensor based the BRIM, sub-vector values, and writes the cumulative dot-product results to the SFD tensor based on the BRIM; vector's y-coordinates. Multiple thread-blocks may write to the same HBM region since results from different BRIM, rows in the same column require aggregation (Figure 5(b)). This can be handled either through at omi c c_add or by se- rially aggregating results across local experts, yielding identical results with similar runtime.
together with a PyTorch-style pseudo-code for the whole pipeline in Appendix A.I.
3.2. Expert Placement Rescheduling
Since Occult's 's communication efficiency CT improves with increased intra collabbollboottoonn and decreased inter- collaboration, expert placement becomes critical for op- timizing all-to-all communication. However, determining optimal placement challenging due the high variability in routing choice distributions across different mini-batches. As discussed in Section 2.4, we propose determining fixed near-optimal expert placement for off -the-shelf MoB-LLMs using profiling dataset.
For scattering (second stage), each thread-bloce reads from the SFD tensor based on the BRIM, sub- vectoretor's 's y-coordinates and writes the cumulative dot-product results to the EPD tensor based on the values in the BRIM, sub-vector.
We provide the implementation details for each basic op- eration in our refactored MoE workflow in Appendix A.2,
We build a collaboration frequency graph P and formulate the rescheduling objective as evenly partitioning P into No sub-graphs, maximizing intracoollaboration while minimiz-
Some data types may not be supported for atromic c_add NVIDIA GPUs. We initialize the output tensor with Float32, and transform to target type after computing.
5

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
Pipeline:
ORI
Dispatch
SFD
All-to-All
SFD
MLP-0
EPD SFD
MLP-1
All-to-All
SFD
Combine
ORI
Forward Pass
BRIM
Scattering-0
BRIM,
Scattering-1
Merging-0
BRIM,
Merging-1
Functions:
Merging-1
Merging-O 0
Scattering-1
Scattering-O
Backward Pass
Figure Correlation between functions BRIM and the basic operations MoE workflow. Each BRIM provides 2 functions. BRIM, is employed for Dispatch & Combine, while BRIM; is adopted for the sparse matrix multiplication.
ing inter-collaboration. Each sub-graph represents device- specific expert assignments. The assigned expert placement L comprises Na lists. We quantify collaboration metrics:
Soppert-Siillarit-bbaaseee Pruning. We prepare an expert similarity table for each MoE layer, named as T E ZN,N,, and then traverse the k routed experts in descending order of routing score. Experts outside the Na devices range replaced with the most similar available alternative from T that satisfies two conditions:
colab-oollaooratioon for device d is measured from the average collaboration frequency among all expert pairs device d. With expert pair set Sa- j(i,j)li e La AJj E La Aat
(1) falls within the range of Na devices, and (2) has not been selected before.
cd Cnntr
1
Eusness Powl
Following the empirical insights from MC-MoE (Li et al., 2023b), we measure the expert similarity using router logits obtained from the inference process on profiling dataset. Implementation details are provided in Appendix A.3.
[Sal
Inter-collaborationn between devices d1 and d2 mea- sures the average collaboration frequency across all feasible cross-device expert pairs:
4. Experiments
Ch
Ld
Ldal
4.1. Experimental Setup
To derive near-optimal expert allocation scheme from P, we develop clustering algorithm in Algorithm 1, inspired by farthest point sampling in point cloud learning (Qi et al., 2017). Compared to trivial expert placement, our reschedul- ing algorithm reduces Cy by ~ 20% while maintaining exact computation results, as shown in Table 1.
Models, Tasks, and Datasets. Models: We examine the effectiveness of Occult on three frontier MoF-LLMs: OLMoE-IB-7B (Muennighoff et al., 2024), Qwenl.5-Mon
A2.7B (Team, 2024), and DeepSeek- (Dai et al., 2024).
The configurations are provided in Table 2. Tasks and datasets: To validate the effectiveness of collaboration prun- ing, we tune the models with Alpaca (Taori et al., 2023), and evaluate on six tasks: Winogrande (Sakaguchi et al., 2021), WSC (Levesque et al., 2012), PIQA (Bisk et al., 2020), RACE (Lai et al., 2017), MathQA (Amini et al., 2019) and RTE (Bentivogli et al., 2009). more compre-
3.3. Expert Collaboration Pruning
We further optimize collaborative communication through algorithm-level collaboration pruning. This approach re- stricts the routing choice of each token to a pre-defined range of N₂ devices. Collaborations outside this scope are pruned and replaced by a legal alternative. The Na devices determined by traversing the tokens' ls routed experts in descending order of routing score until Na devices are reached. 4, Expert replacement follows two proposed crite- ria for pruming implementation.
hensive evaluation with extensive benchmarks is provided in Appendix B.
Baselines and Evaluation Metrics. Baselines: For infer- ence, we compare against vllm (Kwon et al., 2023) and HuggingFace. For training, we compare against Py Torch fully-sharded data parallel (FSDP) and DeepSpeed (Rajb- handari et al., 2022). We evaluate two popular MoE braries: Tutel (Hwang et al., 2023) and MegaBlocks (Gale et al., 2023) (For MegaBlocks, both block sparse matrix mul- tiplication and grouped GeMM for dMoE examined). Since Occult is orthogonal to these general frameworks, vllm and HuggingFace serve as references only. Evaluation metrics: We report accuracy for NLP tasks involved in this section, with Fl score and exact match score results pre- sented in Appendix B.
Pruning. We start by getting the gating scores for all experts as the router network output, denoted as {r(x)) -.N.- While conventional top-k routing selects the largest l values for forward pass and weighted aggregation (Fquations and 2), we modify it with 2 steps: Gating score sorting: Sort {r(*),) N, [r( in descending order.
Beetiin-Scooe-baae
Expert selection: Select top-k experts within the range of the Na devices based on Nr
*Pruning would not he required if the routed experts cover less than N₂ devices
Hardware and Software All the experiments are con- ducted on a single node using PCle-connected NVIDIA
6

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
Turiig (No Phinhal
Phunin (iiiasse)
Phiring (huleshasdd
MM Performanoe
WinoGrande
WSC
PIQA
RACE
MathQA
RTE
Viec Reppiition Times of Tibee
Mdec Repittion Timas of Toden
Moex Reppitton
Teman of Token
Mbcc flapiition Times Tokeen
Mac Rappition Tirea ol Tokenn
Mac Raptition Tires of Tolan
Cover-MoE Performance
WinoGrande
WSC
PIQA
RACE
MathOA
RTE
Vax Repiltion Times of
Max Repiltion Times of Taloo
M0x Repeitoo Times of Takon
MMc Repiition Timee Taken
Rephition Timee at
MM Repittin Tinee of Tokon
Deep Purformance % 3
WinoGrande
WSC
PIQA
RACE
MathOA
RTE
Elepitiion
Figure 7. Performance Comparison for Collahoration Pruning. Comprehensive evaluation across three MoF architectures shows performance trends under different pruning strategies. Note that d-vevece collaboration pruning is equivalent to standard training with original top-k routing.
Coleeiivll Freuuncy
Colleboration the 1st Expert in the 1st MoE layer of OLMoE
Although router-based pruning is easier to implement, similarity-based pruning consistently outperforms it.
No Collaboration Pruning
Explanation for Performance Enhancement. To inter- pret the collaboration mechanism, we visualize an expert from OLMoE's s first MoE layer in Figure 8. Single-device pruning only preserves certain correlations while others are wiped out. In contrast, two-device pruning can ensure the potential collaboration between any pair of experts, main- taining correlation patterns that more closely align with the original model's structure.
Expert Indec in the Tall MoE Layer
Collaboraton Expert the MoE OLMOE (Prune within GPUs)
Coukctrate Fiqquene
Pruning GPUs
Expert Indox MoE Layer Collaboration Expert Moll ayer OLMell (Prune GPU)
CoAodvall Frequecy
Pruning within GPP
4.3. Accelerated Expert Parallelism
Prefilling. Figure 9 demonstrates Occult's t's superior la- tency performance compared to existing frameworks on communication-inteniii tasks. Using 244 prompt tokens with Occult, pruning on two devices, as our benchmark:
MoL
Figure 8. Collaboration Analysis for OLMoE using 4-way ex- pert parallelism. Examining expert 0 in the first MoE layer shows stable collaboration between no-pruning and 2-GPU scenarios, while 1-GPU pruning leads to some patterns' vanishing.
* For OLMOE, it is 8.66x faster than Tutel, 1.70x faster than MegaBlocks and 1.86x faster than vllm.
A6000 GPUs, each with 48GB HBM memory. We take BFloat16 for MoE parameters and activations, and Occult implements atomi c. add operations in Float32, Batch size and token amounts are recorded for individual devices.
* For Qwen-MoF, is 2.74x faster than Tutel, 0.71x faster than MegaBlocks and 1.60 faster than vllm.
* For DeepSeck-MoE, it is 5.64x faster than Tutel, 0.66 faster than MegaBlocks and 1.51x than vllm.
4.2. Performance of Collaboration Pruning Effectiveness of Collaboration Pruning. Figure 7 presents performance comparisons for both router- and similarity- pruning strategics with various device amounts. Our results demonstrate that:
Pruning within 1 device can achieve the highest efficiency, but with slightly inferior quality, while pruning within 2 devices can achieve a superior trade-off.
Decoding. We visualize the latency comparison for full generation in Figure 10. While Occult excels with large batch generation and demonstrates the most consistent scal- ing across MoE-LLMs, indicating superior throughput, its acceleration is limited for small batch sizes where commu- nication is not necessarily a bottleneck.
While restricting collaboration within 1 device may hamper the model performance, expanding the con- straint to 2 devices can achieve comparable or superior quality than standard training.

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
8
Tutel
Owen Prefiling Latero
Tutel
DeepSrekk MGE Prefiinng Larency
Tukel
MegaBlocke (Block Sparss MM)
MegaBlocke (Block Sparse MM))
100
MegaBlodk (Block Sparee NM)
MegeBlocks (Grouped GeMM))
MegaBlocks (Grouped GeNM:
MegaBlocks [Grouped GeMM)
60
Ocoult (No Pruning)
Ocoult (No Pruning)
Occull (No Frunnn)
OLMoE Pretting
Dosult (Pruning Devices)
Occult (Pruning Devicee)
Occult (Pruning, Devices)
Dooult (Pruning. Devices)
Ocoult (Pruning
60
Occult (Fruning, Devices)
8
Mumber of Tokens
Number of Tokens
Number of Tokens
Figure 9. Prefilling Latency Comparison with 4 GPUs. Occult achieves superior efficiency under scaled workload.
viln
HuppingFaou
HugaBiccka (Bick Sparse 186)
VageBBaccs (Orouped GaNMM)
Ooari (ho Pruring
*
Qoadit (Puring, Devian)
Occul (Prunig Devicea)
Deeep Seetic--oo Diecoding Leberry (0)
QUICCE Decadre Labenay (3)
15
Latenry
150
3330233
990
800
133
133
Cowern
155
353
266
408
223
300
206
Culpui Sequence Length
Quiput Sequencee Length
Outpul Sequerce Length
Figure Decoding Latency Comparison with GPUs. Analysis with fixed prompt tokens (12800) and batch size (512) demonstrates
Occult's 's consistent latency advantages on communication intensivv decoding tasks.
Training Latency Analysis for OLMoE (Bandwicth=46GBBs)
9x, and 20x speedup compared to MegaBlocks (Gale et al., 2023), FSDP (PyTorch Fully Sharded Data Par- allelism (Zhao et al., 2023)), DeepSpeed (Rajbhandari et al., 2022), and Tutel (Hwang al., 2023), respectively, Occult is scalable for increasing GPUs involved in ex- pert parallelism, while exhibiting the most stable latency increasing trend with batch size growth compared to other expert-parallel frameworks.
Py fonch SSP
MegetBludo (Beck
Porch FSDP offoad DeepSpeed Zero Stage
MegaB
Average Latency per Training Step (s)
20
Deeplipred Zero Stage 2
(No Fruning) Coout (Phuring Devi
Deepdipeed Zere Stage 3 Tutal
Ocout (fturina Devca)
5. Related Works
Open-Source MoE-based LLMs. Most of the modern MoE-LLMs adopt top-k routing with k > 2, implying that collaborative communication is universal. This enables Occult as a general approach to advance efficiency for scalable expert parallelism. We investigate some frontier open-source MoE-LLMs in Table 2. The attempts in the early stage for sparsely scaling language models include Nllb (Costa-jussà al., 2022) and dus al., 2022). In recent years, have witnessed the vibrant development of modern MoE-LLMs. For instance, Mixtral (Jiang et al., 2024) outperforms Llama2-70B (Tou- vron et al., 2023) with 8 x 7B parameters. Phi-3.5-MoE: (Ab- din et al., 2024) contains 16 x 3.8B parameters, which is on par with Gemini-1.5-Flast and GPT-4o-mini. Qwenl.5- MoE (Team, 2024) activates 2.7B parameters and matches the performance of 7B models. DeepSeek-MoF (Dai et al., 2024) introduces shared experts learn common knowl- edge. OLMoE (Muennighoff et al., 2024) open-sources all data related to model training. PowerMoE (Shen et al., 2024b) trains the MoE model with a new learning rate sched- uler. JetMoE (Shen et al., 2024a) surpasses Llama2-7B (Tou- vron al., 2023) with low training cost.
12
16
20
24
Batch Size
Figure 11. Training Latency Comparison with GPUs. With fixed sequence length of 512 tokens, Occult demonstrates supe- rior performance across varying batch sizes.
(Fe-
Training. We evaluate the average latency per training step over 1k iterations with sequence length 512. We evalu- ate on OLMoE (Muennighoff et al., 2024) and DeepSeek- MoE (Dai et al., 2024) through tuning all the MoE mod- ules and freeziig other parameters. We take a compre- hensive comparison among popular training frameworks with 4 GPUs in Figure 11, and examine the expert paral- lelism training frameworks with 8 and 16 GPUs in Figure 12. Experimental results demonstrate that Occult possesses superior memory efficiency, supporting the largest batch size under the same memory budget. Occult can also speed up ane-ttunigg for MoE-based LLMs. For batch size 8, Occult (two-device pruning) achieves 1.54x, 2.65x,
8

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
DeepSook.Me MoE Training (B-way Expert Paralkllam)
DeepSosk MrE Trainin (16-way Expert Paralialim))
OLMoE Trainina (6-way Expert Paralielsm)
16
Aveeagee Lannnry per Treining Step
Apeal Agg Latentc UNI Treniing SiAp
17.6
Traniins
15.0
12.5
12.0
Blick Anc
74
26 hu
Hogatos 20 eou Coord)
ing. (Duvee))
Panre 4Durmm
(Qvier))
Crrah (Overrn)
32
24
Batch Size
Batch Size
Batch Size
Figure 12. More training latency comparison for expert parallelism frameworks. Owning to the communication- and memory-efficient design, Occult achieves superior training efficiency under both 8- and 16-way expert parallelism configurations.
Table Frontier Open-Souree MoE-LLMs.
The models we used for experiments are highlighted in bold.
efficiency. MC-MoE (Li et al., 2023b) merges the experts into low-rank and structurally sparse alternatives for better efficiency of inference.
Model
A * Exports
# MoE Layers
Params
Principles of GPU Acceleration. Modern GPUs provide massive threads for parallel execution. Threads are grouped into thread- thread-blocks, which are executed on streaming multi- processors (SMs). GPUs have a memory hierarchy, outlined as large but slow-accessed high bandwidth memory (HBM) and small but faster-accessed shared memory (SRAM). Ma- trix multiplication is optimized on GPU using tiling, i.e. partitioning the output matrix into small 2D blocks, where each block is computed using a thread-block in parallel. The size of an individual block can be manually adjusted to improve runtime performance.
Mixtral-sx78 (Jinng et al., 2024)
2
32
46.7 B
Mixtral Sx22B (Jingg et al., 2024)
2
56
141 B
Phi- (Abdin al., 2024)
2
41.9
Minina--D1 (Li et al., 2025)
2
32
80
456 8
O5 (Team, 2024)
4
24
14.3 8
DeepSeck-MoE (Dai et al., 2024)
6
64
27
16.4
DeepSeck V2 (Liu al. 2024a)
6
160
59
236 B
OL.Me. (Mucnnighoff'e et al., 2024)
64
16
6.02 8
Qwen3-30H-A38 (Team, 2025)
8
128
48
30.5 8
Owen3-235B-A22 (Team, 2025)
B
DeepSeck V3 (Liu al., 202
8
256
58 065 B
System Designs for Efficient MoE. DeepSpeed- MoE (Rajbhandari et al., 2022) designs a new MoE architecture and an optimized inference system for efficient and scalable serving. Tutel (Hwang et al., 2023) provides a scalable MoE framework with adaptive parallelism/pipelin- ing optimization at runtime. MegaBlocks (Gale et al., 2023) proposes block-sparse matrix multiplication to enable no discard of tokens. Hexa-MoE (Luo et al., 2024) proposes expert-specific operators as an alternative to GeMM or grouped GeMM to make MoE-training heterogeneous-awaree Janus (Liu et al., 2023) proposes a data-centric strategy to eliminate all-to-all communication overhead under a large workload. APTMeE (Wei et al., 2024) proposes affinity-aware offloading technique to enhance pipeline parallelism for fine-tuning MoE-LLMs.
6. Conclusion
In this paper, we introduce Occult, an algorithm-systee system co-design approach to optimize all-to-all communication in expert parallelism for accelerated MoE training and in- ference. Comprehensive experiments show that it can con- sistently speed up communication-inien LLMs. Occult is orthogonal to popular frameworks, there- fore, it can be further integrated with them for enhanced efficiency. As preliminary research to advance efficiency in expert parallelism, some problems remain unsolved, such as workload balancing. We will continue to explore this topic and try to provide improved solutions.
tasks for MoE-
A ckdnwweegements
Pingzhi Li and Tianlong Chen are partially supported by Amazon Research Award, Cisco Faculty Award, UNC Ac- celerating AI Awards, NAIRR Pilot Award, OpenAI Re- searcher Access Award, and Gemma Academic Program GCP Credit Award.
Algorithm Designs for Efficient MoE. Hash- Layer (Roller et al., 2021) replaces the gating layer in the common MoE model with a precomputed hash function to reduce computation cost. MoE-I2 (Yang et al., 2024) proposes a two-stage compression scheme, including inter-expert pruning and intra-expert low-rank decomposition. Pre-gated MoE (Hwang et al., 2024) proposes a pre-gating function to enable the pre-fetching of MoE parameters so that it can be efficiently served on memmoy--onsttriieed devices. Expert Pruning (Lu et al., 2024) proposes post-training approaches for task-agnostic and task-specific expert pruning and skipping with MoE-LLL to reduce the model size and improve inference
Impact Statement
This paper aims to improve the communication efficiency for training and inference with modern MoE-based LLMs. The efficiency advantage of such models might help de- mocratize access of language models. On the other hand, whether such new algorithm would affect known issues such as biased and harmful outputs of language models remains an unexplored research question.

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
References
Fedus, W. Zoph, B., and Shazeer, N. Switch transformers: Scaling to trillion parameter models with simple and ef- ficient sparsity. Journal of Machine Learning Research, 23(120):1-39, 2022.
Abdin, M., Aneja, J., Awadalla, H., Awadallah, A., Awan, A. A., Bach, N., Bahree, A., Bakhtiari, A.. Bao, J., Behl, H., et al. Phi-3 technical report: A highly capable lan- guage model locally on your phone. arXiv preprint arXiv:2404. 14219, 2024.
Gale, T., Narayanan, D., Young, C.. and Zaharia, M. Megablocks: Efficient sparse training with mixture-of- experts. Proceedings of Machine Learning and Systems, 5:288-304, 2023.
Achiam, J., Adler, S., Agarwal, S., Ahmad, L., Akkaya, I., Aleman, F. L., Almeida, D., Altenschmidt, J., Altman, S., Anadkat, S., et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774, 2023.
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M., Song, D., and Steinhardt, Measuring massive multitask language understanding. In International Conference on Learning Representations, 2021.
Amini, A., Gabriel, S., Lin, P., Koncel-Kedzione Kedziorski, R., Choi, Y., and Hajishirzi, H. Mathqa: Towards interpretable math word problem solving with operation- based formalisms. arXiv preprini Xrxiv: 13319, 2019.
Hwang. C., Cui, W., Xiong, Y., Yang, Z., Liu, Z.,, Hu, H., Wang, Z., Salas, R., Jose, J., Ram, P., et al. Tutel: Adap- tive mixture-of-experrs at scale. Proceedings of Machine Learning and Systems, 5:269-287, 2023.
Aroca-Ouellette, S., Paik, C., Roncone, A., and Kann, K. Prost: Physical reasoning about objects through space and time. In Findings of the Association for Computational Linguistics: ACL-IJCNLP 2021, pp. 5597-4608, 2021.
Hwang, R., Wei, J., Cao, S., Hwang, C., Tang, X., Cao, T., and Yang. M. Pre-gated moe: An algorithm-sysce systtmm co- design for fast and scalable mixture-of-expper inference. In 2024 ACM/IEEE 51 Annual International Symposium Computer Architecture (ISCA), pp. 2024.
Bentivogli, L.. Clark, P., Dagan, I., and Giampiccolo, D. The fifth pascal recognizing textual entailment challenge. TAC, 7(8):1 2009.
IEEE,
Bisk, Y., Zellers, R., Gao, J., Choi, Y., et al. Piqa: Reasoning about physical commonsense natural language. In Pro- ceedings of the AAAI conference on artificial intelligence, volume 34, pp. 7432-7439, 2020.
Jiang, A. Q. Sablayrolles, A., Roux, A.,, Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., Casas, D. d. I., Hanna, E. B., Bressand, F., et al. Mixtral of experts. arXiv preprint arXiv:2401.04005 2024.
Clark, C., Lee, K., Chang, M.-W., Kwiatkowski, T., Collins, M., and Toutanova, K. Boolq: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Hu- man Language Technologies, Volume (Long and Short Pupers), pp. 2924-2936, 2019.
Khashabi, D., Chaturvedi, S., Roth, M., Upadhyay, S., and Roth, D. Looking beyond the surface: A challenge set for reading comprehension over multiple sentences. In Pro- ceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long Pupers), pp. 252-262, 2018
Costa-jussà, M. R., Cross, J., Çelebi, O., Elbayad, M., Heafield, K., Heffernan, K., Kalbassi, E., Lam, J., Licht, D., Maillard, J., et al. No language left behind: Seal- ing human-centered machine translation. arXiv preprint arXiv:2207.04672.1 2022.
Kwon, W., Li, Z., Zhuang. S., Sheng, Y., Zheng, L., Yu, C. H., Gonzalez, J., Zhang, H., and Stoica, I. Efficient memory management for large language model serving with pagedattention. In Proceedings of the 29th Sym- posium on Operating Systems Principles, pp. 611-626, 2023.
Dai, D., Deng, C., Zhao, C., Xu, R., Gao, H., Chen, D., Li, J., Zeng, W., Yu, X., Wu, Y., et al. Deepseekmoe: Towards ultimate expert specialization in mixture-of-experts lan- guage models. arXiv preprint arXiv:2401.06066 2024.
Lai, G.,, Xie, Q., Liu, H., Yang, Y., and Hovy, E. Race: Large-scale reading comprehension dataset from exami- nations. arXiv preprini arXiv: 1704.04683. 2017.
Dasigi, P., Lo, K., Beltagy, I., Cohan, A., Smith, N. A., and Gardner, M. A dataset of information-sceking questions and answers anchored in research papers. arXiv preprint arXiv:2105.03011, 2021.
Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., and Chen, Z. Gshard: Scaling giant models with conditional computation and automatic sharding. In International Conference on Learning Rep- resentations, 2021.
Dolan, B. and Brockett, C. Automatically constructing a corpus of sentential paraphrases. In Thind international workshop on paraphrasing (IWP2005), 2005.
10

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
Levesque, H., Davis, E.,, and Morgenstern, L.. The winograd schema challenge. In Thirteenth international confer- ence on the principles of knowledge representation and reasoning, 2012.
Mihaylov, T., Clark, P., Khot, T., and Sabharwal, A. Can suit of armor conduct electricity? a new dataset for open book question answering. arXiv preprint arXiv:8809 180p:02799,,7 2018.
Li, A., Gong, B., Yang. B., Shan, B., Liu, C., Zhu, C.,, Zhang, C., Guo, C., Chen, D., Li, D., et al. Minimax-01: Scaling foundation models with lightning attention. arXiv preprint S.Ii::2501.083/33 2025.
Muennighoff, N., Soldaini, L., Groeneveld, D., Lo, K., Mor- rison, J., Min, S., Shi, W., Walsh, P., Tafjord, O., Lambert, N., et al. Olmoe: Open mixture-of-experts language models. arXiv preprint arXiv:2409.02060 2024.
Li, J., Jiang, Y., Zhu, Y., Wang, C.,, and Xu, H. Accelerating distributed [MoE] training and inference with lina. In 2023 USENIX Annual Technical Conference (USENIX ATC 23), pp. 2023a.
Qi, C. R., Yi, L.,, Su, H., and Guibas, L. J. Pointnet++: Deep hierarchical feature learning on point sets in met- ric space. Advances in neural information processing systems, 30, 2017.
Li, P., Zhang, Z., Yadav, P., Sung, Y-I.., Cheng, Y., Bansal, M., and Chen, T. Merge, then compress: Demystify efficient smoe with hints from its routing policy. arXiv preprint 3r1v:2310.00134, 2023b.
Rajbhandari, S., Li, C., Yao, Z., Zhang, M., Aminabadi, R. Y., Awan, A. A., Rasley, J., and He, Y. Deepspeed-moe: Advancing mixture-of-expert. inference and training to power next-generation ai scale. In International con- ference on machine learning, pp. 18332-18346. PMLR, 2022.
Liu, A., Feng, B., Wang. B., Wang, B., Liu, B., Zhao, C., Dengr, C., Ruan, C., Dai, D., Guo, D., et al. Deepseek-v2: A strong, economical, and efficient mixture-of-expers language model. arXiv preprint arXiv:2405.0444 2024a.
Reddy, S., Chen, D., and Manning, C. D. Coqa: A conversa- tional question answering challenge. Transactions of the Association for Computational Linguistics, 7:249-266,
Liu, A., Feng, B., Xue, B., Wang, B., Wu, B., Lu, C., Zhao, C., Deng, C., Zhang, C., Ruan, C., et al. Deepseck- v3 technical report. arXiv preprint 1rXiv:2412.19437, 2024b.
2019.
Roemmele, M., Bejan, C. A., and Gordon, A. S. Choice of plausible alternatives: An evaluation of commonsense causal reasoning. In 2011 AAAI spring symposium series, 2011.
Liu, J., Cui, L., Liu, H., Huang, D., Wang, Y., and Zhang, Y. Logiqa: a challenge dataset for machine reading com- prehension with logical reasoning. In Proceedings of the Twenty-Ninth International Conference on Interna- tional Joint Conferences on Arlificial Intelligence, pp. 3622-3628, 2021.
Roller, S., Sukhbaatar, S., Weston, J., al. Hash layers for large sparse models. Advances in Neural Information Processing Systems, 34:17555-17566 2021.
Sakaguchi, K., Bras, R. L., Bhagavatula, C., and Choi, Y. Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 64(9):99-106, 2021.
Liu, J., Wang, H., and Jiang, Y. Janus: A unified dis- tributed training framework for sparse mixture-of-expere models. In Proceedings of the ACM SIGCOMM 2023 Conference, pp. 486 498, 2023.
Shen, Y., Guo, Z., Cai, T., and Qin, Z. Jetmoe: Reaching llama2 performance with 0.1 m dollars. arXiv preprint arX1v:2404.07413, 2024a.
Lu, X., Liu, Q., Xu, Y., Zhou, A., Huang, S., Zhang, B., Yan, J., and Li, H. Not all experts are equal: Efficient expert pruning and skipping for mixture-of-experr large lan- guage models. arXiv preprint 20Xi::2402.114800 2024.
Shen, Y., Stallone, M., Mishra, M., Zhang, G., Tan, S., Prasad, A., Soria, A. M., Cox, D. D., and Panda, R. Power scheduler: A batch size and token number agnostic learning rate scheduler. arXiv preprint arXiv:2408. 13359, 2024b.
Luo. S., Peng, J., Li, P., Wang, H., and Chen, T. Hexa- moe: Efficient and heterogeneous-awant moe accelera- tion with zero computation redundancy, arXiv preprint arXiv:2411. 2,0 2024.
Miao, S.-Y., Y. Liang, C.-C., and Su, K.-Y. A diverse corpus for evaluating and developing english math word problem solvers. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pp. 975-984, 2020.
Socher, R., Perelygin, A., Wu, J., Chuang, J., Manning, D. Ng, A. Y., and Potts, C. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 conference on empirical methods natural language processing, pp. 1631-1642, 2013.

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
Taori, R., Gulrajani, I., Zhang, T., Dubois, Y., Li, X., Guestrin, C., Liang, P., and Hashimoto, T. B. Alpaca: A strong, replicable instruction-followii model. Stanford Center for Research on Foundation Models. https://crfin. stanford. cdu20223/03//1/alpaaae html, 3(6)77, 2023.
Yang, C., Sui, Y., Xiao, J., Huang. L., Gong, Y., Duan, Y., Jia, W., Yin, M., Cheng, Y., and Yuan, B. Moe-12: Compressing mixture of experts models through inter- expert pruning and intra-expert low-rank decomposition. arXiv preprint 24 2024.
Team, Q. Qwen m.eme..-mme:: Matching 7b model per- formance with 1/3 activated parameters, February 2024. URL https: /twenllmmmmmmmmmmmmmm/.mmmmmmmm. weattub.io/bllee qwen-moe/.
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. Hellaswag: Can machine really finish your sentence? arXiv preprint arXiv: 0900999000000,.. 2019.
Zhao, Y., Gu, A., Varma, R., Luo, L., Huang, C.-C., Xu, M.. Wright, L.., Shojanazeri, H., Ott, M., Shleifer, S., al. Pytorch fsdp: Experiences scaling fully sharded data parallel. Proceedings of the VLDB Endowment, 16(12): 3848-3860, 2023.
Team, Q. Qwen3: Think deeper, act faster, April 2025, URL https: /twp///wemmmmmmmmmm//m/ wwt.ll..............m wt/blo///b/oo/ooo /ttpb///////lg//////////////// qwen3/.
Tillet, P., Kung, H.-T., and Cox, D. Triton: an intermediate language and compiler for tiled neural network computa- tions. In Proceedings of the 3rd ACM SIGPLAN Inierna- tional Workshop on Machine Learning and Programming Languages, pp. 10-19, 2019.
Touvron, H., Martin, L., Stone, K., Albert, P., Almahairi, A., Babaei, Y., Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., et al. Llama 2: Open foundation and fine- tuned chat models. arXiv preprint arXivv2307.09288, 2023.
Vaswani, A. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Wang, A., Singh, A., Michael, J., Hill, F., Levy, O., and Bow- man, S. GLUE: A multi-tas' benchmark and analysis plat- form for natural language understanding. In Linzen, T., Chrupata, G., and Alishahi, A. (eds.), Proceedings of the 2018 EMNLP Workshop BlackboxNI.P: Analyzing and In- terpreting Neural Networks for NLP, pp. 353-355, Brus- sels, Belgium, November 2018, Association for Compu- tational Linguistics. doi: 10. 10.18653/v1/11185444 URL https: / wwwpaaaalancoooo..coocaoooooooaa wrg/worg//............. -54454
Wei, Y., Du, J., Jiang, J., Shi, X., Zhang, X., Huang, D., Xiao, N., and Lu, Y. Aptmoe: Affinity- aware pipeline tun- ing for moe models on bandwidth-constraine In SC24: International Conference for High Performance Computing, Networking, Storage and Analysis, pp. 1-14. IEEE, 2024.
gpu nodes.
Welbl, J., Liu, N. F., and Gardner, M. Crowdsourcing multiple choice science questions. In Proceedings of the 3rd Workshop on Noisy User- seerrgenerated Text, pp. 94- 2017.
Williams, A., Nangia, N., and Bowman. S. A broad- coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Confer- ence of the North American Chapter of the Association for Computational Linguistics: Human Language Tech- nologies, Volume (Long Papers), pp. 1112-1122, 2018.
12

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
A. Algorithm Details.
A.1. Refactored MoE Pipeline
We demonstrate the refactored MoE pipeline in Algorithm 2. Notice that the routing weights for the top-k experts are modulated in the intermediate tokens rather than the output ensure symmetric pipeline. We will provide the details for each basic MoE operator in the following.
Algorithm 2 PyTorch-style pseudocode of MoE pipeline.
x.shape: (num pot__ttkks_orr.. hidden_aize)
def MoE_ (sEL,, x):
3
(-) Assign takens to experts.
ids.shape: (num_tokang,
self fttp..)) w.shapo: (nun_takena_orl, self.top_k)
ids, routerix) (x)
Build BRIM-D for dispatch combine.
1orim_0.ahapp (num_devices, tun_tokene_orli
9
brin_0 - build_brin_ (x, ids)
20
Dispatch
x.shapa: tom__tokk__ff__... midden_size) ida.shape: (nur_tokens__fd__)) self.top_k) w.shape: (num_tokens_ _ad___ self.top_5 H, brim_0)
33
2, ids, dispatch
26
(4) 1at all-to-all communication
x.shape: tnum_tokens_ sfd_ hidden_sizes size), shapstd (nin_tokess_sf_)) self.top_k)
w_sEd: (num wwweesookeeaocoooeeeee self. self.top_kk
99
"da_atd, M_ate all all__o_a __o_all_0(M ids, brim_0)
21 (5) Build BRIM-1 for expert computing brin_1-shapez (num loca _ecl__xx,,, num fur_tokens_affl sain_w.ahaape brin_l, brim_w buile_brin_ (x, ida_atd, M_ato)
27
nur_tokens_std_]
25
Computing with local experts shapd.shape: (num Inum_token wtt.ttookkeeppeeeeeeeeee intermellate_siann x.shape: (nun_tokera_afd_) hloden_aizze 28 x_edd - ome_scttteriagle (x, self.wl, brim
x_edd sright_modulateli (x__ep, brim_l, brim_M) arm_mergangix_cee self.w2, brin_1)
31
(T) 2nd all-to-al1 communication
x.shapa: (num_tokenn_sat afd_0, hidden_size)
34
x all_to_all__ brim_l]
35
Combine
x.shape: fnun_tokens_ce eri, hidden_sizei
38
combine (x, brim_31
40
return x
weight.modulate e: assign weights to tokens with EPO state in parallel.
A.2. Algorithms Details for Basic MoE Operations
Build BRIM BRIM, is constructed each device before Dispatch guidance. We provide the details in Algorithm 3.
Algorithm Building BRIM
Require:
Input tokens & shaped as (Nei, D), routing choice r shaped as (Non, No), number of devices N₂ for expert parallel, expert placement L shaped as (Na, [N,/Na]).
Initialize BRIMo with empty tensor shaped (Na, Nun).
Initialize ctr 4-
for + N, - do for - 0, Non - do if token It., activates experts kept device then BRIMo (d, t] ctr. ctr ctr
else
BRIMo (d, 1] &- -
end if
end for end for return BRIMu.
13

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
Dispatch. In Dispatch stage, tokens, routing choices, and routing weights are processed together selectively repeat for each token. We provide the algorithm details in Algorithm where only tokens are processed example.
Algorithm 4 Dispatch.
Require:
Input tokens æ shaped as (Nai, D), BRIMo shaped as (Na, Noi). tile size NM, Initialize Naa with the amount of non non-negen negative elements in BRIMo. Initialize æa5a with empty tensor shaped (Nata, D).
Prepare Na- /Naa/Mi] [D/N]] thread-blocks
parfor d <- 0, Na
do
[M] parfor fm - 0, Naci / - do
parfor tn <- 0, D/N] 1 do
Load BRIMa (d, M: (Im 1) M] from HBM SRAM, denoted as
Load zz (4m M (4m + 1) M,
(4n + 1) -]] from HBM to SRAM, denoted
Write a to ao,af la N: (4n + 1) N]] from SRAM to HBM.
end parfor
end parfor
end parfor
return Vzld.
All-to-All. We take the all.toall.sinnee interface of torch. distributed implement the all-to-all commu- nication in Occult. On each device, the tokens after Dispatch can be further transmitted to device d based on the non-zero elements in BRIMold, [d, *], which serve as the indices.
Build , BRIM| is constructed after the 1st all-to-all communication to guide the sparse matrix multiplication kernels. The details for construction are provided in Algorithm 5.
Algorithm 5 Building BRIM,
Require: Input tokens Tats shaped as (Nas,, D). routing choice Pas shaped as (Nats, NVreesp), local expert Live shaped as (Niocesp ) Initialize BRIM; with an empty tensor shaped as (Nincesps NVat).. Initialize ctr +- O.
for O. Novepp do for 0, Neaa do if token Lmall, activates expert Line [e] then
BRIM [e, 1] <- ctr
ctr clr
else
BRIM, (e, t] f-
end
end for
end for return BRIM
Sparse Matrix Multiplication (Scattering). Sparse matrix multiplication with scattering function takes a token tensor with SFD state and weights as input, and token tensor with EPD state as output, where the weight tensor is dense while the others are sparse. Details are provided in Algorithm 6.
14

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
Algorithm 6 Sparse Matrix Multiplication with Scattering.
Require: Input tokens aata shaped as (Nau, Din), expert weights wo shaped as (Niseesps Din Dowl), BRIM, shaped as (Nucesp Naa., tile size (M, N).
Initialize output tokens Red with an empty tensor shaped as (Nopds Dous).
Prepare -incesp'' [Naa/M] D [ thread-blocks.
parfor <- 0, Nixcopp
do
parfor fm f- 0, Natu/M]
do
parfor fa + 0, Dow// /N] do
On chip: Initialize
E R
for < 0, [Dm/K]]
do
Load 4rsaa M: (4m + 1) M, K: (k+ 1) K] from HBM SRAM a.
Load wle, k K (A: + 1) K, ln N: (ln (1++0) 1) X] from HBM to SRAM as b.
On chip: f c a -.
end for
Load BRIM, (e, fm M: (Im + 1) MM] from HBM SRAM, denoted as v.
Write c to appd U, bn N: (tn + 1) N] from SRAM to HBM.
end parfor end parfor
end parfor
return Zojd-
Weight Modulation. Since the MoE pipeline of Occult is symmetric, we modulate the routing weights on the interme- diate tokens rather than the output, which delivers same results. Details are provided in Algorithm
Algorithm 7 Weight Modulation.
Require: Intermediate tokens Tqu shaped (Neu,, D), routing weights shaped (Nea, No), BRIM, shaped (Ninceyy Nasa local expert list Lise shaped as (Niweeps) ). tile size (N.). Initialize modulated tokens weys with empty tensor shaped as (Nept, D). Prepare Muceup Naa [D/N] thread-blocks.
parfor +0. N
do do parfor tn + 0, [D/N] - do
parfor + 0, Nat
if BRIM [e, 2 0 then
Load Tepp BRIM, [e, (), fn N: (ta + 1). N] from HBM to SRAM, denoted as
Write (t, Liee [e]] to a [BRIM] [e, fa N (fa + 1) -N] from SRAM HBM.
end parfor
end parfor
end parfor
return T'pd
Sparse Matrix Multiplication (Merging). Sparse Matrix Multiplication with merging function takes a token tensor with EPD state and weights as input, and a token tensor with SFD state as output, where the weight tensor is dense while the others are sparse. Details are provided in Algorithm 8.
15

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
Algorithm 8 Sparse Matrix Multiplication with Merging.
Require: Input tokens zpp shaped as (Neps, Din), expert weights to shaped as (Nwseexp Dm, Dow)) BRIM, shaped as (Niscesps Naaa, tile size (M, N).
Initialize output tokens Dat with a zero tensor shaped as (Neta, Dout).
Prepare -incesp'' [NaaMM] [Dea/ [N] thread-blocks.
parfor e <- 0, Nixcopp
do
parfor fm f- 0, Nata/M [M] do
parfor fa + 0, Dow// /N] do
On chip: Initialize
0 E R
Load BRMM [e, 4m M: (4m 1) MM] from HBM SRAM, denoted as v.
for 0- [
do
Load Øed 1. k K (k: + 1) 1)).K K] from HBB to SRAM as a.
Load wle, k- K (k + 1) K, ta N: (ta + 1) NJ from HBM to SRAM as b.
On chip: f c b.
end for
Write c to Irsal [tm M (tm + 1) M, bn N (to + 1) N] from SRAM to HBM via Atomic.Add.
end parfor
end parfor
end parfor return Wald-
Comhine. In Combine stage, different replicas of the token computed different devices are merged into the final output of MoE layer. Details are provided in Algorithm 9.
Algorithm Combine.
Require: Input tokens Data shaped as (Nau, D). BRIM, shaped as (Na, Noa), tile size (M, N). Initialize meii with zero tensor shaped (Non: D).
for 0, Na do
parfor fm f- 0, Naal - do
parfor i, +- 0, D do
Load BRIMa (d, 4m M: (Im + 1) M] from HBM to SRAM, denoted as
Load zar(e (V, N: (l, NJ from HBM to SRAM, denoted
Add e Toa (m (fm 1) f, - (fa + 1) N] from SRAM to HBM.
end parfor
end parfor
end for return Zori-
A.3. Similarity Table Construction
The empirical insights of MC-MoE (Li et al., 2023b) inspire us to measure the expert similarity with router logits, iee, experts with similar router logits tend to be more interchangeable with each other. A profiling dataset is required in this way. We denote the profiled router logits from model inference as H E R.NNN including N..i tokens and N, experts. Similarity table T E N., can be constructed via the cosine similarity between H² and H:
(H H.j)
T | ||H
(6)
To make the similarity table more representative, the profiled token amount should be large enough. However, directly compute cosine similarity for it would cause overflow for floating point number. Therefore change to compute the squared similarity:
(H) Ho, H²
Th
(7)
| ||H_.] 2
We further take an average for both numerator and denominator to normalize their scope:
Non (HT Hm.)² )
T
(8)
Now ||H,, || 1 11
16

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
This can be computed via accumulating small batches for both numerator and denominator.
B. Full Evaluation Results
We implement comprehensive evaluation to validate the effectiveness of collaboration pruning on model performance with 23 popular benchmarks.
For OLMoE, we provide the results Table where the raw off-the-shelf model performs best on 5 benchmarks. Apart from these tasks, router-based pruning outperforms similarity based pruning on 12 benchmarks in total, while similarity-base pruning outperforms router-based pruning on 8 benchmarks.
Table 3. Full Evaluation Results for Collaboration Pruning with OLMoE.
Router- bused Prune, I GPU
Router based
Sim aan-based
Sim- naseb Prune, GPUs
Standard SPT
Router based
Sim based
Method
Metric
Raw
Prune, GPUSs
Prune, GPU
Better
Better
(Lovesque
accuracy
83.52
75.09
86.45
83.52
Winogrande (Sakaguchi et al., 2021)
accuracy
68.35
61.40
68.11
61.56
68.82
69.53
al., 2020)
accuracy
4.56
8.53
OpenBookQA (Mihaylov al., 2018)
accuracy
33:20
26.00
34 60
23.00
32.40
34.20
PIQA (Bisk et al., 2020)
accuracy
80.09
73.45
80.41
71.98
80.47
80.85
HellaSwag (Zellers
2019)
accurncy
58.07
48.14
57.33
47.14
56.90
59.91
SST-2 (Socher et al., 2013)
accuracy
62.61
75.92
86.35
74.66
85.44
83.94
MuluNLI (Williams
accuracy
10.96
35.82
33.61
13.23
QASPER (Dasigi al., 2021)
Fl score
90.53
93.88
97.00
97.03
97.03
97.79
MRPC (Dolan & Brockett, 2005)
accuracy
51.96
66.18
67.65
68 08.14
66.43
MRPC (Dolan & Brockett, 2005)
score
59.84
79.09
80.65
81.34
81.06
79.35
MalaRC (Khashahi al., 2018)
accuracy
57.16
57.22
57.24
56.70
57.10
57.16
WNLI (Wang al., 2018)
accurncy
54.93
53.52
60.56
59.15
67.75
53.52
RTE (Rentivagli et al., 2009)
accuracy
55.60
57.76
58.12
60.65
61.73
63.18
QNLI (Wang et 2018)
accuracy
52.65
50.91
50.28
19.99
MMLU (Hendrycks et al., 2021)
accuracy
50.54
38.30
47.86
37.59
46.88
49.46
/
RACE (Lai al., 2017)
accuracy
37.22
39 29.23
39. 39.62
37.42
38.28
39 62
MathQA (Amini ot 2019)
accuracy
29.92
27.67
29.85
26.93
30.62
31 02
SciQ (Weibl et al., 2017)
accuracy
94.60
91.20
94.40
91.00
94.20
94.70
PROST (ueel-Oulllette
2021)
accurecy
28.24
25.80
28.90
26.80
27.84
29.16
BoolQ (Clazk et al., 2019)
accuracy
71.50
71.96
76.61
68.59
75.96
77.77
COPA
accuracy
85.00
75.00
86.00
LogiQA (Liu et al., 2021)
accuracy
23. 35
23.04
20.04
22.73
20.89
22.73
COQA (Reddy al., 2019)
exact match score
56.42
50.90
56.37
47.63
53.23
56.35
COQA (Redyy al., 2019)
score
71 20
66.17
70.99
63.28
08.85
70.84
17

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
For Qwen-MoE, we provide the results in Table 4, where the off-the-shelf model performs best benchmarks. Apart from these tasks, router-based pruning outperforms similarity- pruning 9 tasks, while similarity-basee pruning outperforms router-based pruning on 11 tasks.
Table 4. Full Evaluation Results for Collaboration Pruning with Qwen-MoE.
Router-based Prune, I GPU
Router-based Prune, GPUs
Sim-based Prunc. GPU
Sim-based Prunc, 2 GPUs
Standard SFT
Router-based Better
Sim-based Better
Methed
Metric
WSC (Levesque ol., 2012)
accurccy
82.05
79.49
$1.68
78.30
82.78
50.50
Winogrande (Sakaguchi et al., 2021)
accuracy
68. 43
68.11
69.06
66.06
70.24
70.01
ASDiv (Miao 2020)
accuracy
4.38
12:28
12.49
10.63
OpenDookQA (Mihaylov al., 2018)
accuracy
30.40
27.90
29.40
28.60
32. 22.20
29.60
PIQA (Bisk al., 2020)
accurncy
79:71
80.25
78.61
79.87
HellaSwap (Zellers et ol., 2019)
accurncy
57.90 96
55.29
57.19
04.73
07.74
67.42
SST-2 (Socher et al., 2013)
accuracy
08:46
78.21
77:41
73.17
87.01
82.22
MalaNLJ (Williams
2018)
accurncy
49.77
49.83
51.88
46.41
54.40
52.50
QASPER (Dasigi al., 2021)
F1 score
90.81
98.78
97.03
98.04
86.66
97.51
MRPC Brocket, 2005)
accuracy
76.47
71.57
78.19
76.72
75.49
MRPC (Dolan & Brockett, 2005)
Fl score
85.08
62.48
85.19
63.76
84.52
63.10
MalriCC (Khashahi al., 2018)
accuracy
37
43.35
39.62
38.77
40.14
36.67
WNLI WWang 2018)
accuracy
57.75
43.06
51.93
56.94
56.31
51.93
RTE (Bentivogli et al., 2009)
accuracy
08.23
69.08
69:68
63.31
75.45
71.12
QNLI (Wang 2018)
accurncy
57.44
53.25
58.05
52.00
51.18
54.99
MMLU (Headrycks et al,, 2021)
accuracy
60. 86
56.85
60.13
56.23
57.25
59.99
RACE al., 2017)
accuracy
41.06
MathQA (Amini et al., 2019)
accuracy
36.15
36.48
38.90
37.09
41.31
39.26
SeiQ (Welbl et al., 2017)
accuracy
94.40
95.40
90.00
96.70
96.50
96.20
PROST (Aroca: (Aroca-Ouellette 2021)
accuracy
30.50
30.01
31.22
30.08
32.47
51.41
BoolQ (Clark et al., 2019)
accuracy
79.57
78.38
79.88
77.22
81.41
80:40
COPA (Roemmele al., 2011)
accurncy
84.00
81.00
86.00
79.00
83.00
$3.00
LogiQA (Liu et al., 2021)
accuracy
30:41
30.26
31.03
31.80
$1. 95
29.03
(Roemmele 2011)
malch
61.40
66.77
66.22
61.48
COQA (Roemmele et al., 2011)
Fl score
78.60
78.15
80.04
79.30
77.53
79.48
/
18

---

Optimizing Collaborative Communications across Experts for Accelerated Parallel MoE Training and Inference
For Dee, we provide the results in Table 5, where the off-the-shelf model performs best 7 benchmarks. Apart from these tasks, router-based pruning outperforms similarity-ban -hased pruning 11 benchmarks, while similarity-basee pruning outperforms router-based pruning on benchmarks.
Table 5. Full Evaluation Results for Collaboration Pruning with DeepSeek-
Router-based Prune, I GPU
Router-based Prune, GPUs
Sim-based Prunc. GPU
Sim-based Prunc, GPUs
Standard SFT
Router-based Better
Sim-based Better
Methed
Metric
WSC (Levesqu ol., 2012)
accurccy
84.98
78.75
$1.88
82.05
84.82
$4.62
Winogrande (Sakaguchi et al., 2021)
accuracy
70.48
66.30
70.40
66 66.22
70.72
71.19
ASDiv (Miao 2020)
accuracy
0.91
1.65
3,08
OpenDookQA (Mihaylov al., 2018)
accuracy
32.20
29.00
33.40
29.20
34.20
34.20
PIQA (Bisk 2020)
accurncy
78-73
77:15
79.92
77.26
79:76
79:11
HellaSwap (Zellers et ol., 2019)
accurncy
58.06
64.05
58.36
63.71
08.31
58.49
/
SST-2 (Socher et al., 2013)
accuracy
61.68
63.71
76.72
89.40
73.05
78.33
MalaNLJ (Williams
2018)
accurncy
42.30
36.77
45.65
38.01
44.42
41.55
QASPER (Dasigi al., 2021)
F1 score
93.06
90.21
96.74
98.00 98.04
97.03
91.67
MRPC Brocket, 2005)
accuracy
68.38
67.65
38
MRPC (Dolan & Brockett, 2005)
Fl score
81.29
60.71
81.17
80.70
81.22
81.22
MalriCC (Khashahi al., 2018)
accuracy
64.08
56.72
66.00
06:88
67.01
WNLI WWang al., 2018)
accuracy
50.70
43.06
45.07
45.07
52.11
47.89
RTE (Bentivogli et al., 2009)
accuracy
62.82
61.01
62:09
61.01
65.34
60.29
QNLI (Wang 2018)
accurncy
49.50
53.12
50.14
50.17
50.80
49.83
MMLU (Headrycks et al,, 2021)
accuracy
37.95
36.16
42.
34.69
41.91
38. .6
RACE al., 2017)
accuracy
38
39.81
38.47
39.62
MathQA (Amini et al., 2019)
accuracy
31.19
29.21
39.50
30.25
33.63
33.77
x
SeiQ (Welbl al., 2017)
accuracy
92 82.80
83.30
93.50
94.70
93350
93.40
PROST (Aroca: Quellette 2021)
accuracy
28.72
28.26
29.00
28.91
28.79
26.61
BoolQ (Clark et al., 2019)
accuracy
72.39
68.09
68.87
73.39
70.15
71.93
COPA (Roemmele al., 2011)
accuracy
90.00
84.00
89.00
86.00
87.00
88.00
LogiQA (Liu et al., 2021)
accuracy
25.35
24.12
25.
25.65
25.65
25.90 35.96
(Reddy
malchh
641.1
64.02
63.83
62.25
62.80
COQA (Reddy al., 2019)
Fl score
77.79
77.08
77.33
73.46
77.05
76.72
/
19


Reviews:
[{"review_id": "c4AhlJcZBJ", "reviewer": "Reviewer_tR7y", "summary": "This paper introduces Occult, an algorithm-system co-design approach aimed at reducing the communication overhead of Mixture-of-Experts (MoE) large language models (LLMs). Specifically, the authors first propose BRIM, a data structure designed to support fundamental MoE operations efficiently. Next, they optimize expert placement based on calibration data to minimize inter-device communication. Finally, by replacing experts on remote devices with similar experts on the N_d devices, they further reduce collaborative communication overhead. Experimental results demonstrate that the proposed method significantly improves efficiency across various scenarios, including prefilling, decoding, and training, compared to existing frameworks.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**[Experiments & Analysis 1: Larger MoE and fewer-expert models]**\n\nOccult 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. \n\nWe'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.\n\n|Model|# Experts|# Devices for EP|Method|Avg Latency Per Step|Speedup|\n|-|-|-|-|-|-|\n|Mixtral-8x7B|8|4|Megablocks|9.64|1.0|\n|Mixtral-8x7B|8|4|Occult, w/o prune|9.08|1.06|\n|Mixtral-8x7B|8|4|Occult, prune within 1 GPU|8.34|1.13|\n|DeepSeek-MoE|64|8|Megablocks|16.56|1.0|\n|DeepSeek-MoE|64|8|Occult, w/o prune|12.10|1.37|\n|DeepSeek-MoE|64|8|Occult, prune within 3 GPUs|10.95|1.51|\n|DeepSeek-MoE|64|8|Occult, prune within 2 GPUs|9.31|1.78|\n|DeepSeek-MoE|64|8|Occult, prune within 1 GPU|8.50|1.95|\n\n\nWe also provide the evaluation results for 8-way expert parallelized DeepSeek-MoE with Occult to demonstrate its effectiveness:\n\n|Task|Strategy|No Tune|Prune within 1 GPU|Prune within 2 GPUs|Prune within 3 GPUs|Prune within 4 GPUs|Prune within 5 GPUs|No Prune|\n|-|-|-|-|-|-|-|-|-|\n|MMLU|Router-based|37.95|35.04|40.41|41.34|41.43|41.19|38.66|\n|MMLU|Sim-based|37.95|33.68|39.80|**41.74**|41.40|41.48|38.66|\n||\n|OpenBookQA|Router-based|32.20|33.8|36.2|37.2|**37.8**|37.2|34.20|\n|OpenBookQA|Sim-based|32.20|33.4|36.4|36.8|**37.8**|37.2|34.20|\n||\n|MathQA|Router-based|31.19|32.93|35.08|34.97|35.95|**36.08**|33.77|\n|MathQA|Sim-based|31.19|33.17|34.94|35.51|35.24|35.61|33.77|\n\n\n**[Experiments & Analysis 2: Performance improvement explanation]**\n\nThe 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. \n\n**[Question: Difference with existing method for MoE]**\n\nOur framework differs from existing MoE libraries in several fundamental ways, making it more communication-efficient:\n- **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.\n- **State-based token representation**: As stated in lines 206-216, \"we outline them as 3 states: Original ($\\texttt{ORI}$)... Simplified ($\\texttt{SFD}$)... Expanded ($\\texttt{EPD}$),\" allowing us to maintain tokens in \"the token counts across states follow $\\texttt{ORI}<\\texttt{SFD}<\\texttt{EPD}$.\" This contrasts with existing methods that simply replicate $k$ times for each token, regardless of collaboration patterns.\n- **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.\n\nThank you again for your questions. We will emphasize these distinctions in our revised draft.\n\n\n\n[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.\n\n[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.", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "No", "experimental_designs_or_analyses": "Issues:\n\n- The experiments can be further improved. It is better for the authors to conduct experiments on larger MoE and MoE with fewer number of experts (e.g., Mixtral), which will demonstrate the generalization and scaling ability of the method.\n\n- According to Table 3 and the introduction section, the proposed method can achieve higher performance than the original model. However, the authors only provide explanation on why two-device pruning is better than single-device pruning in Sec 5.2. More explanation is needed regarding the former one.", "relation_to_broader_scientific_literature": "No", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "**Strengths**\n1. The motivation of the paper is clear. Collaborative communication is a major overhead in MoE computation, and optimizing it through co-design of hardware and algorithms is a reasonable approach.\n\n2. Experiments show that the proposed method achieves faster speed and higher performance compared to baseline methods.\n\n**Weakness**\nPlease refer to \"Experimental Designs Or Analyses\"", "other_comments_or_suggestions": "No", "questions_for_authors": "The authors provides a new framework for all-to-all communication in Sec 4.1. What's the difference between this framework and existing method for MoE?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "80ypwF9sdf", "reviewer": "Reviewer_geV9", "summary": "In this paper, the author proposes Occult, an MoE training and inference framework designed to reduce communication costs by effectively managing intra- and inter-collaboration among experts. The evaluation results demonstrate that the proposed method achieves significant speedup compared to the state-of-the-art MoE framework.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**[Experiments & Analysis 1: Model configuration]**\n\nThanks for the constructive feedback. We've added a supplementary table detailing model configurations:\n\n|Model|# Params|# Experts|Top-K|# MoE layers|Hidden size|Expert intermediate size|\n|-|-|-|-|-|-|-|\n|OLMoE|6.92B|64|8|16|2048|1024|\n|Qwen1.5-MoE|14.3B|60|4|24|2048|1408|\n|DeepSeek-MoE|16.4B|64|6|27|2048|1408|\n\n**[Experiments & Analysis 2: Pruning strategies]**\n\nThe performance differences between router- and similarity-based pruning arise from how they handle expert replacement:\n- Router-based pruning uses only routing scores, which may not involve expert attribute\n- Similarity-based pruning considers inter-expert similarity, obtained from a profiling dataset\n\nAs 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.\n\nSimilarity-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. \n\n**[Experiments & Analysis 3: throughput, TTFT, and TPOT]**\n\nFollowing 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:\n\nPrefilling:\n\n|Model|Method|# Tokens|Throughput|TTFT (s)|\n|-|-|-|-|-|\n|OLMoE|Occult (Prune, 2 GPUs)|16384|1777.01|9.22|\n\nDecoding:\n\n|Model|Method|# Generated Tokens|Throughput|TPOT (ms)|\n|-|-|-|-|-|\n|OLMoE|Occult (Prune, 2 GPUs)|512|682.67|1.46|\n\n**[Strengths And Weaknesses 1: Cost of pruning]**\n\nPruning does not require more time to converge because \n- 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. \n- The average latency per training step can be greatly reduced with Occult.\n\nWe report the latency and memory cost of pruning for DeepSeek-MoE here, with batch size 16 and 8 GPUs:\n\n||Megablocks|Occult (w/o Prune)|Occult (Prune, 4 GPUs)|Occult (Prune, 3 GPUs)|Occult (Prune, 2 GPUs)|Occult (Prune, 1 GPUs)|\n|-|-|-|-|-|-|-|\n|Memory (GB)|43.12|40.68|36.93|34.46|32.18|30.73|\n|Avg Latency per Step (s)|8.36|6.43|6.02|5.54|4.44|3.93|\n\nOur prune algorithms can accelerate clock-time latency and reduce memory costs for training.\n\n**[Strengths And Weaknesses 2: Bandwidth impact]**\n\nWe 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:\n\n|Bandwidth|Speedup (Occult vs. Megablocks, w/o Prune)|Speedup (Occult vs. Megablocks, Prune, 2 GPUs)|Speedup (Occult vs. Megablocks, Prune, 1 GPU)|\n|-|-|-|-|\n|18GB/s|1.37|1.78|1.95|\n|46GB/s|1.12|1.43|1.76|\n\n**[Strengths And Weaknesses 3: Scalability]**\n\nOur approach can remain effective at a very large scale with a hierarchical strategy:\n- 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.\n- Inter-node optimization: Cross nodes, expert placement rescheduling & collaboration pruning can be applied to minimize cross-node communication, which is typically the latency bottleneck.\n\nThis 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.\n\n**[Comments: Notation issue]**\n\nThanks for the constructive suggestion on “faster” and “speedup”, we will fix it in the revised paper.\n\n**[Question 3: Integration with DeepEP]**\n\nDeepEP mainly contains 2 stages for hardware-tailored MoE communication:\n- Inter-node all-to-all: sending tokens to the corresponding GPU on the target node\n- Intra-node all-to-all: sending tokens to the target GPUs containing the target expert\n\nIn 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.\n\nThank you again for your questions. We will include these additional experiments and analyses in our revised draft.", "claims_and_evidence": "Overall, the evidence is sufficient to support the claims.", "methods_and_evaluation_criteria": "Overall, the baselines and datasets are appropriate.", "theoretical_claims": "Overall, the theoretical claims are correct.", "experimental_designs_or_analyses": "1. Model configurations, such as the number of experts, are critical but not provided. \n\n2. It appears that different pruning strategies significantly affect performance (e.g., similarity-based and router-based). However, there is no empirical or theoretical analysis to determine which strategy should be utilized under specific scenarios. \n\n3. Regarding Figures 9 and 10, presenting throughput (tokens/s), TTFT, and TPOT under varying batch sizes and sequence lengths may be more common, useful, and informative for evaluating performance.", "relation_to_broader_scientific_literature": "This work enhances the training and inference efficiency of MoE models, which is crucial for future LLM deployment.", "essential_references_not_discussed": "The references provided are appropriate and sufficient.", "other_strengths_and_weaknesses": "1. There is no discussion on the pruning cost or the impact of bandwidth. \n\n2. The performance may heavily depend on the number of GPUs. Since the evaluations are conducted with 4 GPUs, there is a greater opportunity to increase intra-collaboration. However, for larger models trained with more GPUs, such as 1,000 GPUs, it will be more challenging to place experts on the same device, or model performance may be compromised to achieve this.", "other_comments_or_suggestions": "The notation \"0.66x faster\" and \"0.55x speedup\" seems to indicate a slowdown rather than a speedup. Maybe it should be 1.66x faster and 1.55x speedup.", "questions_for_authors": "1. What is the cost of pruning? Does it require more time to converge? \n\n2. What is the impact of bandwidth? It seems that the improvement is more significant in a low-bandwidth system. What would the improvement be in a high-bandwidth system? \n\n3. DeepSeek recently released DeepEP, a communication library tailored for MoE and EP. Is your optimization orthogonal to DeepEP? Would your method benefit from DeepEP or achieve even better performance if integrated with it? I believe this is a parallel work, so the discussion is not required but welcome.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lmHK7bvTkS", "reviewer": "Reviewer_1GwW", "summary": "The paper introduces Occult, an algorithm-system co-design approach to optimize collaborative communication in MoE models for large-scale training and inference. The key idea is to reduce inter-device communication costs by maximizing intra-device expert collaboration, using expert placement rescheduling and collaboration pruning strategies. The paper shows that Occult achieves over 50% speedup across various MoE-based LLMs while maintaining or even improving model quality. The authors provide theoretical justifications, empirical validation, and extensive benchmarking against state-of-the-art frameworks like DeepSpeed, Tutel, and MegaBlocks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n**[Potential concerns 1: Different task types]**\n\nWe'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:\n\nHumanEval:\n\n|Method|No Tune|Prune, 1 GPU|Prune, 2 GPUs|No Prune|\n|-|-|-|-|-|\n|Router-based|26.83|17.68|**27.44**|22.56|\n|Sim-based|26.83|17.07|26.22|22.56|\n\nGSM8K ( flexible-extract ):\n\n|Method|No Tune|Prune, 1 GPU|Prune, 2 GPUs|No Prune|\n|-|-|-|-|-|\n|Router-based|16.91|7.28|15.31|**17.97**|\n|Sim-based|16.91|11.22|16.00|**17.97**|\n\n**[Potential concerns 2: Scalability]**\n\nOur 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:\n\nAvg training latency per step (s):\n\n|Setting|Occult ( 1 GPU )|Occult ( 2 GPUs )|Occult ( 3 GPUs )|Occult ( 4 GPUs )|MegaBlocks|\n|-|-|-|-|-|-|\n|8 GPUs|8.50|9.31|10.95|11.92|16.56|\n|16 GPUs||9.55|10.25|10.93|14.97|\n\nOccult acceleration can be more apparent on a better-grouped expert placement, thereby Occult can better improve training efficiency on 8-way expert parallelism.\n\nWe also evaluate 8- and 16-way EP on MMLU and MathQA:\n\n8-way EP:\n\n|Task|Strategy|No Tune|Prune within 1 GPU|Prune within 2 GPUs|Prune within 3 GPUs|Prune within 4 GPUs|Prune within 5 GPUs|No Prune|\n|-|-|-|-|-|-|-|-|-|\n|MMLU|Router-based|37.95|35.04|40.41|41.34|41.43|41.19|38.66|\n|MMLU|Sim-based|37.95|33.68|39.80|**41.74**|41.40|41.48|38.66|\n||\n|MathQA|Router-based|31.19|32.93|35.08|34.97|35.95|**36.08**|33.77|\n|MathQA|Sim-based|31.19|33.17|34.94|35.51|35.24|35.61|33.77|\n\n16-way EP:\n\n|Task|Strategy|No Tune|Prune within 2 GPUs|Prune within 3 GPUs|Prune within 4 GPUs|Prune within 5 GPUs|No Prune|\n|-|-|-|-|-|-|-|-|\n|MMLU|Router-based|37.95|39.69|40.37|41.23|**41.62**|38.66|\n|MMLU|Sim-based|37.95|39.23|40.25|41.31|41.61|38.66|\n||\n|MathQA|Router-based|31.19|35.61|35.14|35.21|**35.78**|33.77|\n|MathQA|Sim-based|31.19|34.84|35.21|35.68|35.71|33.77|\n\n**[Strengths and Weaknesses 2: Collaboration pruning trade-offs]**\n\nWe'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:\n\n||Megablocks|Occult (w/o Prune)|Occult (Prune, 4 GPUs)|Occult (Prune, 3 GPUs)|Occult (Prune, 2 GPUs)|Occult (Prune, 1 GPUs)|\n|-|-|-|-|-|-|-|\n|Memory (GB)|43.12|40.68|36.93|34.46|32.18|30.73|\n|Avg Latency per Step (s)|8.36|6.43|6.02|5.54|4.44|3.93|\n\n**[Comments 1: To multi-node setting]**\n\nThanks 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:\n- 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 )\n- Expert parallelism can be organized across different nodes, placing the grouped experts in Occult on the same node\n- Tensor parallelism can be adopted in intra-node GPUs for very large expert since the intra-node bandwidth is usually abundant\n\n**[Comments 2: RL-based expert assignment]**\n\nEnd-to-end rescheduling may be impracticable, as dynamically tuning expert placement requires huge additional bandwidth, although it can be asynchronous.\n\nWhile 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.\n\nThank you again for your questions. We will incorporate the refinements in our revised draft.", "claims_and_evidence": "The paper claims that optimizing collaborative communication via expert placement rescheduling and collaboration pruning can significantly reduce all-to-all communication costs in MoE models, leading to faster training and inference with minimal quality degradation.\n\n**Support:**\n\n(a) Theoretical derivations: The authors introduce a collaboration graph-based formulation to quantify inter- and intra-collaboration.\n\n(b) Algorithm side: The expert placement rescheduling algorithm is validated through profiling experiments showing up to 20% reduction in communication budget.\n\n(c) Empirical evaluation: Occult achieves up to 8.66× speedup for inference and up to 10× faster training over baseline MoE frameworks.\n\n**Potential concerns:**\n\n(a) The impact of aggressive collaboration pruning on model quality could be better explored for different task types.\n\n(b) No direct discussion on scalability beyond 4 GPUs—it would be useful to see how Occult performs on larger clusters and models.", "methods_and_evaluation_criteria": "The proposed methods make sense for optimizing communication bottlenecks in MoE training and inference. Evaluation criteria are appropriate, leveraging: (a) Latency benchmarks for training and inference (Figures 9–11). (b) Accuracy and task performance metrics across multiple NLP benchmarks. The concern is that the paper does not explicitly evaluate multi-node scaling performance, and the evaluated model is relatively small scale.", "theoretical_claims": "The paper does not introduce new mathematical proofs, but it provides well-founded theoretical insights, including quantification of communication cost and its relationship with intra- and inter-collaboration; and the optimization bounds for communication overhead in MoE expert parallelism.", "experimental_designs_or_analyses": "Experimental setup is generally robust, with:\n(a) Multiple MoE architectures (OLMoE, DeepSeek-MoE, Qwen-MoE) (2) Comparison against strong baselines (DeepSpeed, Tutel, MegaBlocks). (3) Latency benchmarks covering different workloads (training, inference, decoding).\n\nLimitations or improvements:\n(a) Limited scalability analysis beyond 4 GPUs. (b) The selection of expert placement rescheduling is heuristic—would an end-to-end learned approach perform better?", "relation_to_broader_scientific_literature": "The paper is well-grounded in prior work, citing relevant MoE frameworks (DeepSpeed-MoE, Tutel, MegaBlocks). The connection to general parallel computing strategies could be stronger—Occult shares similarities with load balancing and scheduling techniques in distributed systems.", "essential_references_not_discussed": "There are no specific references currently on my mind that have not been discussed.", "other_strengths_and_weaknesses": "Pros:\n1. The paper addresses a key limitation in MoE scalability. Communication is a well-known challenge in distributed MoE training, and Occult provides a practical, well-validated optimization.\n2. The paper presents strong empirical results, a 50%+ speedup in multiple MoE workloads is a compelling result.\n3. A lightweight, heuristic expert placement rescheduling was proposed that provides significant efficiency gains.\n\nCons:\n1. Limited discussion on multi-node scalability. The paper mainly explores single-node, multi-GPU setups, leaving open questions about large-cluster scaling.\n2. Trade-offs in collaboration pruning are not well studied. More ablation studies would help clarify when pruning impacts accuracy.", "other_comments_or_suggestions": "1. What are required to be done to scale up Occult when moving beyond single-node multi-GPU setups?\n2. Instead of heuristics, could an RL-based expert assignment policy be more effective?", "questions_for_authors": "Please check \"Other Comments Or Suggestions\"", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "azycsUS8SU", "reviewer": "Reviewer_NL62", "summary": "All-to-all communication is a major bottleneck in training and inference for mixture-of-experts (MoE) large language models. While existing MoE kernels have improved computational efficiency, all-to-all communication remains a bottleneck. The authors propose Occult, which aims to (1) reduce redundant communication and (2) encourage in-device collaboration via router update to further optimize latency.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank reviewer NL62 for the dedicated and professional comments. To address your concerns, we provide detailed pointwise responses below:\n\n**[Claims and Evidence]**\n\nWe provide code at https://anonymous.4open.science/r/Occult-D802.\n\n**[Theoretical claim 1: Communication complexity in Fig. 2]**\n\nThank 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.\n\nTo clarify the communication complexity in Fig. 2:\n- Fig. 2(a): Each token is repeated twice, yielding $C_{\\mathcal{T}}$ = 2\n- Fig. 2(b): Red and green tokens are repeated once while yellow and blue tokens are twice, yielding $C_{\\mathcal{T}}$ = 1.5\n- Fig. 2(c): Each token is repeated once, yielding $C_{\\mathcal{T}}$ = 1\n\nThe 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.\n\n**[Experiments & Analysis 1: Multi-node Training]**\n\nOur 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.\n\n**[Experiments & Analysis 2: Fine-tuning baseline for fair comparison]**\n\nIn Sec. 5.2 and Fig. 7, we provide the evaluation results for \n- Original model (no tuning), shown in yellow dashed line\n- Tuned model with pruning, shown in brown dots ( similarity-based pruning ) and green stars ( router-based pruning ).\n- Standard tuning, shown in pink diamond. \n\nIn 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.\n\n**[Strengths and weaknesses 1: Fig. 2]**\n\nThank 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.\n\n**[Strengths and weaknesses 2: Communication complexity]**\n\nWe 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:\n- 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.\n- 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.\n\nWe appreciate you highlighting this concern and will expand these explanations in Section 3.2 to provide better clarity.\n\n**[Strengths and weaknesses 3: Figure 4]**\n\nWe apologize for the missing in-text reference to Fig. 4. We will add detailed descriptions, including:\n- $\\texttt{SFD}$ tokens serve as the all-to-all content, constructed from $\\texttt{ORI}$ tokens based on $BRIM_0$\n- $\\texttt{FFN}$ utilizes $BRIM_1$ as an auxiliary input to guide the token layout of the output tensor\n- Intermediate tokens are organized densely in the $\\texttt{EPD}$ state\n\nI hope they can help readers better understand the token state transitions and the role of $BRIM$s in our framework.\n\n**[Other comments]**\n\nThanks for pointing out. We will fix these typos carefully in our revised draft.\n\n**[Question: Huggingface Latency]**\n\nPrefilling 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:\n- Limited GPU amount ( our experiments use only 4 GPUs )\n- Deep model architectures ( the three evaluated models contain 16, 24, and 28 layers, respectively )\n\nThereby, 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.\n\nThank you again for your comments. We will incorporate these refinements in our revised draft.", "claims_and_evidence": "Most claims are well-supported by evidence or align with well-established knowledge in the field. However, some claims lack sufficient clarity or appropriate experimental validation:\n\n- The claim that the proposed algorithm benefits training is not adequately evaluated. (See below)\n- The paper does not provide the corresponding code, making it difficult to verify reproducibility.", "methods_and_evaluation_criteria": "The proposed method and evaluation criteria are well-aligned with the problem at hand.", "theoretical_claims": "Equation (5) states that communication complexity is lower-bounded by the number of devices needed to fill k and upper-bounded by the minimum of k and the number of devices. While this holds in terms of send/receive operations, this definition differs from the one used in Figure 2 and Section 3.1, creating inconsistency.", "experimental_designs_or_analyses": "The experimental design appears insufficient to fully support the paper’s claims:\n\n- The stated goal is to reduce all-to-all latency, which is especially problematic in multi-node multi-gpu setups. However, the evaluation only uses four GPUs connected via PCIe, which has limited bandwidth. This setup leads to a high number of experts per GPU and does not reflect realistic multi-node scenarios. While the resource constraints are understandable, without larger-scale results, the claim about latency reduction remains unconvincing. The author should provide more results, either from experiments or projections, to support this claim. Without the new evidence, the proposed method might be more applicable for **inference**, and the paper should consider revising its scope to emphasize this contribution.\n\nFor section 5.2:\n- The comparison is misleading: fine-tuning naturally improves model performance, so a fine-tuned model outperforming a baseline is expected. If the goal is to show that \"expanding to two devices achieves comparable or superior quality to standard training,\" the baseline should also be fine-tuned for a fair comparison.\n\n\n**I hope the author could address these concerns, and I'm more than willing to update my evaluation if my concerns are addressed.**", "relation_to_broader_scientific_literature": "The proposed optimization reduces all-to-all latency, a well-known bottleneck in MoE-based LLMs. The work contributes to this active research area.", "essential_references_not_discussed": "I'm not aware of any essential references that were not discussed.", "other_strengths_and_weaknesses": "While I think the paper do provide some interesting insight on MoE, the writing of this paper has clarity issues, and should be significantly improved. Sections 3 and 4 are difficult to follow and should be rewritten for clarity. Here are some examples:\n\n- Figure 2\n    - This figure is very confusing. Readers unfamiliar with MoE literature will struggle to interpret it. The term device D-0/1 appears to refer to physical devices, but it actually denotes expert assignments of token. The figure needs revision.\n\n- Communication Complexity in Section 3.2\n  - The term \"communication complexity\" is misleading. A more accurate term would be \"communication volume\" to reflect what is actually being measured, but it seems to be different from the definition used in Figure 2.\n\n- Figure 4\n  -There is no in-text reference to Figure 4. The figure itself is unclear and needs better explanation.", "other_comments_or_suggestions": "Line 218 right column: continuously -> contiguously\nLine 383: device amount -> device count", "questions_for_authors": "Why is Hugging Face’s decoding latency lower than the proposed method when batch size is small?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: All four reviewers converge toward a positive recommendation, with two accepts and two weak accepts. The authors provided comprehensive responses that effectively addressed concerns regarding scalability and experimental clarity. The consensus among the reviewers, along with the paper’s clear motivation in tackling communication costs in MoE models, supports acceptance.
Decision: Accept (poster)