NetMoE: Accelerating MoE Training through Dynamic Sample Placement
Authors: Xinyi Liu
Yujie Wang
Fangcheng Fu
Xupeng Miao
Shenhan Zhu
Xiaonan Nie
Bin CUI
Venue: ICLR 2025 Spotlight (2025)
Keywords: Mixture of Experts
All-to-All communication
Distributed training
Primary area: infrastructure, software libraries, hardware, systems, etc.

Abstract: Mixture of Experts (MoE) is a widely used technique to expand model sizes for better model quality while maintaining the computation cost constant. In a nutshell, an MoE model consists of multiple experts in each model layer and routes the training tokens to only a fixed number of experts rather than all. In distributed training, as experts are distributed among different GPUs, All-to-All communication is necessary to exchange the training tokens among the GPUs after each time of expert routing. Due to the frequent and voluminous data exchanges, All-to-All communication has become a notable challenge to training efficiency.

In this paper, we manage to accelerate All-to-All communication in MoE models from the training sample perspective, which is unexplored so far. In particular, we put forward the observation that tokens in the same training sample have certain levels of locality in expert routing. Motivated by this, we develop NetMoE, which takes such locality into account and dynamically rearranges the placement of training samples to minimize All-to-All communication costs. Specifically, we model the All-to-All communication given the sample placement and formulate an integer programming problem to deduce the optimal placement in polynomial time. Experiments with 32 GPUs show that NetMoE achieves a maximum efficiency improvement of $1.67 \times$ compared with current MoE training frameworks.

Published as a conference paper at ICLR 2025
NETMoE: ACCELERATING MoE TRAINING THROUGH DYNAMIC SAMPLE PLACEMENT
Xinyi Liu' Yujie Wang',
Fangcheng Fu'
Xupeng Miao?
Shenhan Zhu'
Xiaonan Nie! Bin Cui
Schooo of CS & Key Lab of High Confidence Software Technologies (MOE), Peking University Purrdue University Institute of Computational Social Science, Peking University (Qingdao)
{xy. 1iu, al alfredwang, c @pku. edw.
xupeng@purdue edu, {shenhan. zhu, xiaonan. nie, bin.
ABSTRACT
Mixture of Experts (MoE) is widely used technique expand model sizes for better model quality while maintaining the computation cost constant. In a nut- shell, MoE model consists of multiple experts in each model layer and routes the training tokens to only a fixed number of experts rather than all. In distributed training, as experts are distributed among different GPUs, All-to-All communica- tion is necessary to exchange the training tokens among the GPUs after each time of expert routing. Due to the frequent and voluminous data exchanges, All-to-All communication has become notable challenge to training efficiency.
In this paper, manage accelerate All-to-All communication in MoE models from the training sample perspective, which unexplored so far. In particular, we put forward the observation that tokens in the same training sample have cer- tain levels locality in expert routing. Motivated by this, we develop NetMoE, which takes such locality into account and dynamically rearranges the placement of training samples to minimize All-to-All communication costs. Specifically, we model the All-to-All communication given the sample placement and formulate an integer programming problem to deduce the optimal placement in polynomial time. Experiments with 32 GPUs show that NetMoE achieves maximum effi- ciency improvement of 1.67 compared with current MoE training frameworks.
Device,
Device
Device,
INTRODUCTION
MHA
MHA
MHA
Data Parallelism
input,
input,
Input,
In recent years, large language models (LLMs) have shown impressive performance in lan- guage understanding and generation (OpenAl
gating
gating
gating
All-to-All Scatter
2023] Touvron et al. 2023; Zhou et al. 2024
Dubey al. 2024 Shao al. 2024 Zhang et al.; 2024a) due to the increasing model size. However, larger models often come with greater computational costs. To address this, Mixture of Experts (MoE) models have been in- troduced to expand the model size greatly with- increasing the computational cost. Com- models
Expert,
Expert,
Expert,
Model Parallelism
output
output
output
bining MoE with Transformer-haace can yield outstanding performance across var- ious tasks, including natural language process- ing (Lepikhin et al. [021] Fedus et al. 2022).
An example of expert parallelism applied an MoE model with devices and E 2. experts (each device has two different experts).
Figure 1:
computer vision (Riquelme et al. 2021 Liang et al. 2022
recommendation systems (Tang al.
2020
Zou et al. 2022), and speech recognition (You et al. 2022 Kwon & Chung 2023
MoE models often replace the feed-forward network (FFN) layer with the MoE layer, which
sists of a gating network and several small FFNs, representing different experts. In the MoE layer, each token is routed by the gating network only few selected experts, and the final output is

---

Published as a conference paper at ICLR 2025
(a) An overview of' MoE. layer example.
All--AIl
0-3
10
1.0
111
122
13
2-0
2:1
2-
23
RAUS IMAI
240
2-1
242
3-3
222
NodeO
Nedel
SampleO
Espent)
Sample?
Deviced
Jdba mode
intin sode
Inkeess
Experi
Expert
Device
1-0
Alli-ear All Gather
Devicel
1-0
porquagu
jiter
Giahus
Toleet
Nodel
tokenss
Exper12
Sample?
Expert
Semple
Deviced
8-
2-1
Device2 [-]
2-01
0 token
lokeus
lokers
lokens
Espert3
Sample3
Expor:
Sampled
Devios
12
Devicel
[-]
1-2
(b) gaiber operation the MoE layer without adjusting gather operation the MoE layer after sumple placement sample placement.
adjustment enabled.
An example of sample exchange. The figure illustrates the All-to-All gather operation in a MoB layer with two nodes, each containing two devices, and each device having one expert. Different colors represent tokens to different experts, and i-j denotes the j-th token in the i-th sample. Fig. 2la) illustrates the complete process
Figure 2:
MoE layer during forward propagation. Fig. [bb shows the All-to-All gather operation
in the MoE layer without adjusting the sample placement, where the inter-node communication volume of each node is 5 tokens. Fig. ([c) displays the All-to-All gather operation after sample placement adjustment is enabled the positions of samples the devices change (samples O and are exchanged), reducing the inter-node communication volume tokens per node.
obtained by a weighted sum of the computations from the selected experts. By such means, we can increase the number of experts expand the model size for better performance, while keeping the computation complexity constant.
Despite the above benefit, given the potentially large number of experts, the memory capacity of single device is often insufficient. As result, expert parallelism (Lepikhin et al.) [021 Fedus et al.) 2022) is a common technique to facilitate the distributed training MoE models. As shown in Fig. each device holds only subset of the experts to reduce memory consumption. Meanwhile, other model parameters are replicated and stored on all devices, and the training data assigned to each device are different. In each MoE layer, based the routing result of the gating network, each token is sent to the device where the selected expert is located. The output from the expert is then sent back to the original device of the corresponding token. This involves two communication operations, namely the All-to-All scatter and All-to-All gather (He et al. 2021), respectively.
Due to the dynamic nature of routing, training MoE models efficiently faces several challenges, with the All-to-All communication, being the most significant one. Particularly, the All-to-All commu-
nication time can account for up 80% of the total training time (Hwang al. 2023 Liu al.
2023
He et al. 2022; Li et al., 2023 Yu et al.) 2024). One reason is because all tokens need to participate in the All-to-All operation. leading to a high communication volume. Another reason the communication frequency. Considering both the forward and backward propagation, each MoE
layer requires four All communications per training iteration. Such frequent and extensive communication incurs significant time costs. Therefore, accelerating All-to-All communication is essential to improve training efficiency.
Motivation: Recent studies have demonstrated that expert routing exhibits certain degree of data localiry. To be specific, input tokens may have distinct preferences for experts, and the correspond- ing distribution is often skewed (He et al., 2022, Nie al.) 2023 Xue et al., 2024] Jiang et al., 2024). In other words, the All-to-All operation in MoE models can be highly unbalanced across different devices, thus bounded by the device pair with the highest communication volume. Mean-
2

---

Published as a conference paper at ICLR 2025
while, is well known that network locality is an inherent characteristic of modern clusters for deep learning training. In particular, there are various communication channels in modern clusters, e.g. intra-node devices usually communicate via PCIe or NVLINK, while inter-node devices net or InfiniBand, with intra-node communication usually faster than inter-node ones. To achieve load balancing, existing methods propose techniques from the model perspective (Nie et al., 2023 Lewis et al., 202I). Typically, they either dynamically adjust the model placement but introduce a lot of additional communication, or modify the model definition but sacrifice the model performance (see $2.2 for more discussion). Yet optimization from the data perspective is under explored. In- spired by this, we propose NetMoE, which accelerates the All-to-All communication by combining the data locality in expert routing with the network locality among training devices. The essential idea of NetMoE is to dynamically adjust the placement of data samples during training based on expert routing results so that more tokens will be transmitted through high-speed channels rather than low-speed ones. As illustrated in Fig. 2 sampleo shows preference for expertz that re- sides node1, while samples favors experto that resides on nodeo. Using the vanilla All-to-All communication method would result in significant inter-node communication overhead, as shown in Fig. (b) However, by swapping the positions of sample, and samples depicted in Fig. ()),, part of the inter-node communication can be converted into intra-node communication or even intra- device memory copying, significantly reducing the time cost (detailed in $3.1). In this way, we can accelerate the All-to-All communication without affecting the computing results.
Ether-
However, it is non-trivial achieve dynamic sample placement. For one thing, how adjust the placement to maximize efficiency is complex and unexplored question. For another, since the ad- justment should be done for every layer in every iteration, it is vital devise an efficient algorithm to deduce the placement on the fly. To address these problems, we first revisit the cost modeling for All-to-Al to-All communication and formulate the dynamic sample placement problem into a combinato- rial optimization problem. Subsequently, we split it into two stages to ease the solving and design corresponding polynomial-tine algorithm to ensure timely solution.
In short, the technical contributions of this work are summarized as follows:
We propose NetMoE, the first effort that leverages both the data locality and network locality to accelerate the All-to-All communication through dynamic sample placement.
We formulate the dynamic sample placement problem as a combinatorial optimization problem, which aims to find the best sample placement that maximizes efficiency given the expert routing.
We dissect the problem into two stages and develop polynomial-tine solution to efficiently derive the sample placement during training.
We conduct experiments with various models on 32 NVIDIA A800 GPUs. Results show that NetMoE outperforms current MoE training systems by up to 1.67x in terms of training efficiency.
2
PRELIMINARY
2.1
PARALLELISM IN DISTRIBUTED TRAINING
Data and Model Parallelism: In data parallelism (Li et al. 2020: Sergeev & Balso 2018 Wang
al.) [2023 Zhang et al. (004bb), each device maintains a complete copy of the model parameters, while different training samples are assigned to each device. After the backward computation is completed, the model gradients from all devices are aggregated before updating the model param- eters. In model parallelism (Naaaanan et al. 2021bb Huang et al.) 2019 Narayanan et al. 2021a Guan al.) 2024), model parameters are distributed across multiple devices, with each device re- sponsible for only portion of the model. Communication operations are necessary to transmit the intermediate results (a.k.a. forward activations and their backward gradients) to accomplish the forward and backward propagation.
Expert Parallelism: As shown in Fig. [ expert parallelism (Lepikhin et al.) 2021 Fedus et al., 2022) can be regarded as combining model parallelism and data parallelism. It distributes expert parameters across different devices like model parallelism, while replicating other parameters on all devices like data parallelism. In each MoE layer, each token will be routed by the gating network to top K different experts for processing, where K is a hyperparameter, typically a small value, such or 2, which helps to reduce the computational complexity. After the MoE layer obtains the
3

---

Published as a conference paper at ICLR 2025
gating routes, tokens sent to the devices where the corresponding experts are located based the routing. The results from the expert computations are then sent back to the original devices where the tokens are located. Since the experts distributed across different devices, communication during this process involves all devices sending and receiving messages with one another, leading to what known All-to-All communication.
2.2 DISTRIBUTED TRAINING ACCELERATION TECHNIQUES FOR MoE MODELS
Dynamic Expert Placement: The efficiency of MoE models is constrained by the extensive and frequent All-to-All communication required during training. response to this issue, some studies have observed that data tends to show a preference for certain experts during training. Then, based on this observation, they further propose to dynamically adjust the placement of experts to reduce communication volume (He et al. 2022 Nie et al. 2023: 7Zhai et al. 2023). For instance, popular experts can be placed on more devices in the data parallel manner, so that the communication volume related to them would decrease. However, due to the growing size of experts, these approaches incur substantial overhead of transmitting expert parameters among the devices, so they cannot adjust the expert placement for every iteration, leading to sub-optimality. In contrast, the communication volume from a different perspective: we dynamically adjust the placement of samples in every iteration accelerate the All-to-All communication. To be specific, an optimization problem to deduce the best sample placement that minimizes the time cost of All- to-All communication. As we will evaluate in $4 our work outperforms existing works based on dynamic expert placement when training MoE models.
work tries to reduce
formulate
Modification in Model Definition: To achieve better workload balance in MoE training, there are many existing works developed to modify the model definition (e.g., routing mechanisms, model architectures). Some approaches modify the routing mechanism balance the load across experts, which helps reduce synchronization time between devices (Lewis et al. 2021). Recognizing the network locality in distributed training, several works introduce routing topology loss to prior- itize routing tokens within the same node, thereby reducing inter-node communication (Li et al,, 2024 Chen et al., 2022). Other approaches (Zeng & Xiong 2023) map tokens to smaller hidden layer dimension before inter-node communication, further decreasing the communication load. SC- MoE (Cai et al. 2024) proposes feeding the output of the current attention layer directly into the next MoE layer, enabling parallel forward propagation with the current MLP layer in order to fully overlap All-to-All communication with computation. Although these methods improve training ef- ficiency, they inevitably impact model convergence.
When applying these methods, we usually need to run numerous trials to tune the hyper- parameters, such as adjusting the weight of the topology-aware routing loss (Chen al. 2022) or tuning the papeer-prameters for differ- ent communication channels (Zeng & Xiong 2023). Given that each trial of LLM training can take days or even months, their utility is inevitably hampered. In contrast, our work fo- cuses on how to accelerate All-to-All commu- nication without affecting model convergence.
Problem Formulation
Two-Stage Dissection Stage
2 Stage
2 Stage
ILP Problem (Eq. 5] $3.1
$3.2
Implementation
Polynomial-time Solver
Scatter
1" Stage
gnd Stage nodes
Residual
Inlining
FFN
&
(Eq.7) & (0.1)-IPP (Eq.11)
Add
Add
[0.11-ILP (Eq:20) + Bipartite Graph
Offloading Solver
Scatter
FFN
Bipartito Graph
NETMOE
Solver
$3.3
$3.2
In this section, we introduce NetMoF, novel framework designed to optimize distributed training for MoE models by considering both data and network locality. Given target MoE model and the hardware environment, NetMoB aims to minimize the All-to-All communication cost. Its core innovation lies in optimizing the placement of samples within each MoE layer to maximize the utilization of faster intra-node bandwidth, thereby reducing the communication volume slower inter-node connections. Specifically, NetMoE swaps the samples across devices during each MoE layer, enabling more tokens to communicate within the node during All-to-All communication.
Figure 3: The overview of the method of NetMoE.
4

---

Published as a conference paper at ICLR 2025
Table I: Notations used throughout this work. We assume I is divisible by
J, and J is divisible by N. which are common in distributed training.
Table 2: Bandwidth of cach channel of the NVIDIA A800 GPU cluster used in our ex- periments.
L
The number of tokens per sample.
H
The hidden size for each token.
E
The number of experts in the MoE layer.
K
The number of experts to be routed per token.
I
The number of samples per iteration (a..a. global batch size).
J
The number of devices (i.c., GPUs).
Channel
Bandwidth
N
The number of nodes (machines).
Intra-device
~2TB/
The indicator function.
Intra-node Inter-node
400GB/s
The set of natural numbers less than n., i... {0, 1, .... n 1}.
100GB/s
Fig. fillustrates the overview of this section. We begin by introducing the modeling of All-to-All
communication in MoE training and formulate our optimization problem in $3.1 We then illustrate how to solve the problem in $3.2 with the detailed algorithm shown in Alg. We also present our
implementation details in 3.3 For clarity, the frequently used notations are listed in Table
3.1 PROBLEM FORMULATION
Communication Modeling: We first discuss the mathematical modeling of All-to-All communica- tion, which is the optimization target of NetMoE. We use the a-3 model (Sarvotham et al. (001) to analyze All-to-All communication, where represents the latency cost and represents the band- width cost. Specifically, we classify communication into three categories: intra-device, intra-node, and inter-node communication, each using different channels. Table lists the bandwidth of each channel used in our experiments. Since intra-device communication is typically achieved via mem- ory copying, it is significantly faster than the other two categories and thus considered in our modeling. Therefore, the communication time is determined by the maximum time required for data transfer across the intra-node and inter-node channels. The bandwidths of these channels are represented by Vintras and Uinter, respectively. Thus, for each All-to-All communication, its time cost can be expressed by the following formula, where 8. represents the communication volume for the corresponding channel.
t max (Lintra, linter); where
lintra = Oljnttra Sentra Bintra Uintray
Linter Qinter Binter Binter 1/ Vinter
The bandwidth (v.) and latency (a.) can be obtained by profiling the hardware environment before training, while the communication volume (s.) needs to be dynamically determined based on the routing results within the MoE layer. We then analyze how to calculate the communication volume.
Let roule E NIxL*K be the token routing results of the gating network, which represents the K experts that each token will be sent Then, the number of tokens that the i-th sample needs send to the e-th expert can be counted as
Σ Troutes, e = e] for M..e E [T]
nummi
Next, num N/XE can be used to model the communication volume across different channels. Let ExpDev(e) be the device index of the e-th expert, SmpDev(i) the device index where the i-th sample should be routed to, and Node(j) the node index of the j-th device. By considering the communication volume as the number of tokens that need to be transmitted, we have
Σ
Sintra
numiee
Sinter
numii
(G.))
(s.))
where Sintra and Sinter can be calculated via the device indices experts and samples:
Sintra {(t, e) (i)pee)))) Node(RxpDev(e))
A SmpDev(i) * ExpDev(e)) Sinter = {(i, e) Node (SmpDe(()))) * Node(ExpDev(e))
(4)
Rationality of Dynamic Sample Placement: that the time cost of All-to-All communication is highly related the placement of experts and
Given the aforementioned modeling, there is no doubt
5

---

Published as a conference paper at ICLR 2025
samples. In practice, dynamically adjusting the placement does not affect the training results the All-to-All communication is still correctly performed. Combining with the common fact of network locality that Vintra > Vinter, we can adjust the placement of samples and/or experts to reduce the inter-node communication volume, even the intra-node communication volume becomes slightly higher. In fact, with similar goal, several existing works have proposed to dynamically adjust the placement of experts based on their popularities (He et al. 2022 Nie et al., 2023 Zhai et al.) 2023), as introduced in $ However, all these works overlook the data locality
tokens of the
same sample are usually routed to the same expert (Xue et al., 2024 Jiang et al., 2024). thereby missing the optimization opportunity of dynamic sample placement. More importantly, the size of the parameters of experts is usually much larger than the size of the samples. This prevents previous works from adjusting the expert placement in every iteration. In contrast, the adjustment of sample placement can be fused with the All-to-All communication by nature (detailed below), requiring zero extra communication. Consequently, this work focuses on the unexplored aspect, aiming to accelerate MoE training by dynamically adjusting the sample placement
To help readers better understand the strength of dynamic sample placement, we take Fig. as an example, where I 4, L = 4, E = 4, K 1, and both the experts and samples placed sequen- tially, i.e., ExpDev = [0, 1, 3], SmpDev = [0, 2, 3]. Fig. 2[b) shows the communication with- changing the placement of samples. According to Eq. 3 and Eq. if we only consider the send- ing volume of node 0, then Sinter = {(0, 2), (0, 3). (1 2) 1, ()3,), indicating that Sinter = 5. How-
ever, after optimizing the placement of samples as Fig. c), i... SmpDev = [3, 2, 0). the cor- responding inter-node communication volume changes into Sinter = {33, 2), (3, 3), (1, 2), (1, which gives Sinter = 2. Furthermore, it is worth noting that the sample placement adjustment can be combined with the All-to-All gather operation. To be specific, instead of restoring tokens to their original positions, they are directly placed in their new positions according to the altered sample placement. This method directly optimizes the current communication operation without introducing any extra communication.
Problem Formulation: After the sample placement adjustment is determined, it can be seen that altering SmpDev affects two All-to-All operations: the gather operation of the current MoE layer and the scatter operation of the next MoE layer. Thus, our optimization targets these two operations. For the 1-th layer, the optimization problem can be written as follows.
(lgatheer) + z(+ (cuuter)))
arg min SmpDev(() [J] for ie [?]
max
(lgathee)) intra
+ max
Cnntra
(lltla ((+tterateer) (5) inttr
s.t. Σ I[mmDDvv(i) = 1/J for [J] l
Since single change in sample placement affects two All-to-All operations, both communication times are included in the optimization objective. Additionally, to ensure computational and memory balance across devices, each device should retain the same number of samples before and after the sample placement adjustment. This forms the basis for the constraints our optimization.
3.2 PROBLEM SOLVING
5q is a complex combinatorial optimization problem, which cannot be solved optimally in polyno- mial time. As the cluster size increases, even finding an approximate solution may take significant amount of time. Since this problem needs to be solved before each gather operation, solving directly would result in unbearable additional time costs. To address this, design efficient method to obtain approximate solutions. In particular, we first dissect the optimization problem into two stages and develop polynomial-time algorithm to achieve the solution, as introduced below.
Two-Stage Dissection: Although Eq. D takes the maximum value of the two kinds of communica- tion cost, practice, due the significant bandwidth difference between the inter- and intra-node
Our work is fully compatible with the dynamic expert placement technique. Specifically, in the problem formulation and solving of NetMoE, we do not make any assumption on the expert placement. Instead, it is treated as input. Thus, we can dynamically adjust the expert placement like previous works, and NetMoF. still deduce the optimal sample placement. We would like leave the combination future work.

---

Published as a conference paper at ICLR 2025
Nodel
Nodal
Expert2
All-to-A All Gather
Sample2
Espert2
All-to-A1I Gather
Sample2
Device2
0-2
Deviced
[10
0-2
1-3
2-1
0-0
0.1 0.2
NVLod
lokems
intra node
swap
sokens
tokeen
sample
Experts
Sampled
Experts
Samplet
Devices
Device
0-1
1-2 240
2-0 2-1
2-2
(a) Sample placement after the first stage.
(b) Sample placement after the second stage.
Figure 4: An example of the second stage optimization. Fig. a)) shows the MoE layer in Node; after the first stage optimization in Fig. (]... By applying the second stage optimization within the node, the intra-node
communication can be reduced by token (by swapping samples and samplez), shown in Fig. 4 b).
connections, the most time-consuming term is usually the inter-node one. Therefore, we propose two-stage solving strategy: the first stage optimizes tinter the global scale, while the second stage minimizes finra within each node, without affecting finter. Formally, suppose there are N nodes and each node consists of J/N devices, then the optimization formula of the first stage can be written as the following integer linear programming (ILP) problem:
(.gather) ((+1aeenter) inter
arg min (SmpllS (s) for iE
Ginter
s.t.
Σ 1[Node((mpDev(0) = n] = I/N for n E [N] SE
The constraint balance across devices in Eq. turned into the balance across nodes since we focus on inter-node communication in the first stage. After obtaining the optimal solution of the first stage, the second stage considers rearranging the samples within each node individually. For the n-th node, denote [T]; S [m] the set of samples appointed to it after solving Eq. 6 And let []], = {ili E [J] A Node(j) n be the set of experts reside on it ([]], is determined by the device placement rather than obtained by Eq. 6). Then, to optimize for the n-th node, we should solve the following ILP problem:
min [, for KE
(Lgather)
intrr
+
1 scatte)
inntra
s.t. Σ I[SmpDev(i) = = I/J for e [J] (7) SE(/
Specifically, Fig. can be regarded as the optimization of the first stage, while Fig. demonstrates the second stage of optimization built upon it. Although the second stage consists N ILP prob-
lems, each for one node, they are independent and can be solved concurrently.
Polynomial-time Solver: By dissecting the original combinatorial optimization problem, we obtain N + ILP problems, which can be solved via existing libraries like PuLP (Mitchell et al. 201T.. However, recall that we need to solve these problems for each layer in each training iteration, the efficiency of mobblem-solving is vital. Unfortunately, since ILP problems are NP-hard, when try to solve them via PuLP, the time cost of solving exceeds the time cost of seatter communica- tion and experts' computation (as evaluated in $4.4). making it impractical. Given the fact that each sample must be assigned to one device, we reconsider the ILP problems as assignment prob- lems by transforming them into weighted bipartite matching problems, and subsequently develop polynomial-ting solver based on the widely used Kuhn-Munn Munkres (KM) algorithm.
We first introduce how transform the ILP problems into weighted bipartite matching problems. Let Cin and c, represent the inter- and intra-node communication volume when placing the i-th sample on the j-th device in the n-th node, They can be calculated using the following formulas:
Σ numaee: 4. =
reamsse where
ees'
S (elxxode(ExDDev(e)) # n). S' = (e)xoode(xxDee(ee) Node(j) A ExpDev(e) + }j).
To make the expression clearer, let pi,n, Pl,,, {0, 1} indicate whether the i-th sample is placed on the n-th node and the j-th device, respectively. Then, the optimization objective can be expressed as
finter = Otnter + Binter
Σ
Ct.uPi.y
tintra = Quntra + Iintra
Σ
dipphss
where
(9)
Pia = [[Nodc(SmDDev(0 = n), p'., = [[SmpDov(i) =
for e [
7

---

Published as a conference paper at ICLR 2025
Algorithm NetMoE Optimization
1: function So Solvel Solve(num)
2: Get c, via Eq. 8 and build bipartite graphs
3: Get the optimal solution p' via the Munn--Mu (KM) algorithm
4: return the optimal sample placement according to p*
5: The Main Training Process:
6: for submodule in model do
7: if submodule is MoE layer then
8: Get route from the gating network and calculate num via Eq.
9: Invoke Solves ve(uum) in a background thread
> Offloading solving process
10:
Get impul from All-to-All scatter
Get output from expert computation
12:
output = imput + output
> Expert residual inlining
Get the optimal sample placement from the background thread
14:
Perform All-to-All gather with the optimal sample placement
15:
else
16:
Coomodule.forwardde
After modifying the corresponding constraints, transform the ILP problems into (0,1)-ILP prob- lems. For instance, below presents the transformed problem for the first stage (i.e., Eq. 6
arg min Ps.m ne]]
Clenter + Binter
Σ
(.gather)
+ G.n
((1..catter))
pi.
MELN
(10)
s.t.
Σ
Pin I/N for [N],
Σ nE[N]
Pin for
ME
This (0.I)-ILP problem can modeled as a weighted bipartite matching problem. In particular, consider bipartite graph with two sets of graph nodes, P and Q. The set P represents all train- ing samples, and |P| = I. The set Q represents all training nodes (machines), where each training node can handle B == I/N train- ing samples. To model this, each graph node in is duplicated times, resulting in IQ1 = I. A weighted edge exists between every pair of graph nodes from P and Q. Let P. represent the i-th training sample and Qn the nn/ B) -th training node. The weight of the edge (l.gaahhe))
Samples
Nodes di glonk
0
Hikey
ouken tokens
0
(fl+t,ecatter) between P; and Qn is denoted as Wt,,. L, BJ + + LLL/BB This transformation reduces the problem of finding a minimum weight perfect matching in this bipartite graph, which ficiently solved to optimality in polynomial time using the Kuhn- Munkres (KM) algorithm. Fig. illustrates structing a bipartite graph during the first stage in Fig. 2 The graph nodes the left represent set P. and the graph nodes the right
be ef-
example of con-
Figure An example of bi- partite graph.
represent set Q Each pair of graph nodes is connected by weighted edge, depicted by dotted
line. The red edges indicate the final matching scheme, where the total weight of all matched edges is minimized.
3.3 IMPLEMENTATION
NetMoE implemented on top of PyTorch (Paszke al. 2019), with custom operations (e.g., the calculation of num,, c', and the KMM algorithm) implemented in C++ and CUDA. The complete workflow of NetMoE is presented Alg. [ In addition to the problem-solving introduced in $3.2 NetMoE has been optimized in the following ways.
Expert Residual Inlining: In classic MoE models, residual connections are independent of the MoE layers. However, in NetMoE, the position of the training data changes after the All-to-All
"The problems of the second stage (Eq. 7 can also he transformed into 10.I)-II.P problems and solved in polynomial time similarly. We amit them here due the space constraint and only discuss the first stage.
8

---

Published as a conference paper at ICLR 2025
Table Time cost of different solvers vs. the summed time cost of All-to-All scat- ter and expert computation in milliseconds (MoF-GPT-S, J 16).
Table 3: Configurations of the evaluated models.
Model Name
Base
S
H
E
K
MoB-GPT-S
GPT-2
4
1024
768
2
2
3
KM
PuLP
Scatter + Computation
MoE-GPT-M
GPT-2
1024
1024
2
2
0.08
42.8
3.69 (2.63 + 1.06)
MoE-GPT-L
GPT-2
1024
1280
2
4
0,48
50.1
7.13 (5.34 1.79)
MoE-GPT-XL
GPT-2
1024
1600
2
1.48
72.9
13.50 (10.31 + 3.19)
MoE-GPT-XXL
GPT-3
4
1024
4096
2
2
10.82
143.7
27.31 (21.49 + 5.82)
24
31.09
266.5
41.65 (33.82 + 7.83)
gather operation, while the samples in the residual connections remain in their original positions. To ensure the correctness of the model, inline the residual connections into the expert computation, as shown in line 12 of Alg. [ This optimization ensures consistency in model accuracy before and after applying the algorithm. More details about inlining is elaborated in Appendix A.T
Offladingg Solving Process: The KM algorithm is hard to parallelize, making it unsuitable for highly parallelized accelerators like GPUs, so we perform the solving process the CPU. As shown in line 9 of Alg. [ after obtaining the routing results for the current layer, each device calculates and transfers nurn the CPU memory. The routing results for the next layer, required by the optimization algorithm, can be predicted by directly passing the current layer's input to the router of the next layer (Eliseev & Mazur, 2023] Tang et al. 2024). The solving process only needs to provide the new sample positions before the All-to-All gather operation. In this way. the solving process can be overlapped with the All-to-All scatter and expert computation. As we will show in $4.4 the solving time is fully hidden and thus introduces zero overhead. More discussion of algorithm selection and overlap potential is described in Appendix
4 EXPERIMENTS
4.1
EXPERIMENTAL SETUPS
We compare NetMoE with state-of-the-art methods based on dynamic expert placement, including FasterMoE (He et al.] 2022) and SmartMoE (Zhai et al.) 2023). We also included FastMoE (He ct al. 202I) to represent a baseline without adjusting the placement of experts or samples. All experiments are conducted
cluster consisting of 4 nodes, each equipped with 8 NVIDIA A800-
SXM4-40GB GPUs. As listed in Table 2 the GPUs within each node are connected via NVLink with a 400 GB/s bandwidth, while the nodes are interconnected via InfiniBand with a 100 GB/s bandwidth. The configurations of the evaluated models are listed in Table 3 We select the GPT model architecture (Radford al. 2019. Brown al. 2020 as the backbone and replace all FFN layers in each model with MoE layers. In particular, since SmartMoE requires at least 2 experts on each device, we the number of experts as E' the corresponding experiment, and we fix the number of selected experts for each token as K = 2. By default, we utilize 8 GPUs per node to carry out the experiments, and we present the results for scenarios with fewer GPUs per node in Appendix B All results are averaged over 50 iterations.
J. where J the number of GPUs in
4.2 END TO END PERFORMANCE
As shown in Fig. 6 NetMoE demonstrates up to a 1.67x speedup over FastMoF, a 1.37 x speedup over FasterMoE, and 1.33x speedup over SmartMoE. FasterMoE achieves significant optimiza- tion by overlapping expert computation and supporting dynamic expert placement. However, as the model's hidden dimension increases, the cost of communicating with experts rises, making difficult for it to maintain the same level of acceleration. This leads to performance gap between FasterMoE and NetMoE. On the other hand, SmartMoB outperforms FasterMoE, which is expected since SmartMoB adjusts expert placement to ensure load balancing top of FasterMoE's opti- mizations. However, SmartMoE primarily focuses on balancing the computational load, without emphasizing communication efficiency. When communication becomes the primary bottleneck, the benefits of load balancing are less pronounced. Consequently, by dynamically adjusting the sam- ple placement, NetMoE consistently outperforms the state-of-the-art systems. Last but not least,

---

Published as a conference paper at ICLR 2025
FastMoE
FasterMoF
SmartMoF
NetMoE
16GPUs
speedup
1.5
1.0
2015
0.5
0.0
32GPUs
Speedup
1.5
#7/4
1.0
2010
MoE-GPT-S
MoE-GPT-M
MoE-GPT-L
MoE-GPT-XL
MoE-GPT-XXL
Figure 6: End-to-end speedup (mean and standard deviation) of different methods.
Theory
Speedup
16 GPUs
4.4G
1.61x
32 GPUs
1.42x
3. 1.60x
1,00x
1.00x
1.00x
1.00x
1.00x
McE-GPPP GPT-GP
Figure 7: The actual theoretic speedup terms All-to-All communication cost.
it is noteworthy that method is compatible with dynamic expert placement. By adjusting the ExpDev(.) that is fed to our solver, NetMoE can be combined with dynamic expert placement to achieve even higher efficiency. We plan explore this integration in future work.
4.3 ALL-TO-ALI PERFORMANCE
As shown in Fig. 7 we conducted experiments on three kinds of model to observe the differences in All-to-All communication before and after applying NetMoF and compared these results with the thcoretical optimization values provided by the solver. It can be seen that the actual speedup in All- to-All communication slightly less than the theoretical values. This discrepancy reasonable, as our modeling of All-to-All communication assumes ideal conditions and does not account for potential routing conflicts or pardware-induces errors. In Appendix C we have provided more experimental results to analyze the acceleration of All-to-All communication.
4.4 SOLVER PERFORMANCE
To verify the efficiency of the solver, we compared the solving time under different scales with the summed time cost of All-to-All scatter and expert computation, shown in Table KM represents the algorithm used NetMoE, while PuLP (Mitchell et al., 2011) refers to the commonly used toolkit for solving linear programming problems. It can be observed that although the solving time exhibits super-linear growth with the increase in 7, the solving process is consistently hidden by the All-to-All scatter and expert computation for various scenarios. In contrast, PuL.Ps solving time difficult to get overlapped. This highlights the necessity of designing specialized optimization methods in scenarios with high real-time performance demands.
CONCLUSION
We proposed NetMoE to optimize All-to-All communication, which is the primary bottleneck in training MoE models. By leveraging data and network locality, our method dynamically adjusts the placement of training samples during training, transforming inter-node communication into intra-node communication to enhance All-to-All communication efficiency. We modeled the All- optimization problem and designed
to-All communication time and the sample placement a polynomial-tine time approach to solve it. Empirical results demonstrate that NetMoE outperforms existing MoE training systems by up 1.67 in terms of training efficiency.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
This work is supported by National Science and Technology Major Project 16602200016315)) Na- tional Natural Science Foundation China (U22B2037, U23B2048, 62402011), Beijing Municipal Science and Technology Project (Z231 100010323002), China National Postdoctoral Program for Innovative Talents (BX20230012). China Postdoctoral Science Foundation (2024M750103) Bei- jing Natural Science Foundation (4244080), the Fund of Kunpeng and Ascend Center of Excellence (Peking University), and High-performance Computing Platform of Peking University. Bin Cui is the corresponding author.
REFERENCES
Bo Adler, Niket Agarwal, Ashwath Aithal, Dong H. Anh, Pallab Bhattacharya, Annika Brundyn, Jared Casper, Bryan Catanzaro, Sharon Clay, Jonathan M. Cohen, Sirshak Das, Ayush Dattagupta, Olivier Delalleau, Leon Derezynski, Yi Dong, Daniel Egert, Ellie Evans, Aleksander Ficek, Denys Fridman, Shaona Ghosh, Boris Ginsburg, Igor Gitman, Tomasz Grzegorzek, Robert Hero, Jining Huang. Vibhu Jawa, Joseph Jennings, Aastha Jhunjhunwala, John Kamalu, Sadaf Khan, Oleksii Kuchaiev, Patrick LeGresley, Hui Li, Jiwei Liu, Zihan Liu, Eileen Long, Ameya Sunil Maha- baleshwarkar, Somshubra Majumdar, James Maki, Miguel Martinez, Maer Rodrigues de Melo, Ivan Moshkov, Deepak Narayanan, Sean Narenthiran, Jesus Navarro, Phong Nguyen, Osvald Nit- ski, Vahid Noroozi, Guruprasad Nutheti, Christopher Parisien, Jupinder Parmar, Mostofa Patwary, Krzysztof Pawelec, Wei Ping, Shrimai Prabhumoye, Rajarshi Roy. Trisha Saar, Vasanth Rao Naik Sabavat, Sanjeev Satheesh, Jane Polak Scowcroft, Jason Sewall, Pavel Shamis, Gerald Shen, Mo- hammad Shoeybi, Dave Sizer, Misha Smelyanskiy, Felipe Soares, Makesh Narsimhan Sreedhar, Dan Su, Sandeep Subramanian, Shengyang Sun, Shubham Toshniwal, Hao Wang, Zhilin Wang, Jiaxuan You, Jiaqi Zeng. Jimmy Zhang, Jing Zhang, Vivienne Zhang, Yian Zhang, and Chen Zhu. Nemotron-4 340b technical report. CoRR. abs/2406. 11704, 2024.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Horbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler. Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neu- ral Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Weilin Cai, Juyong Jiang, Le Qin, Junwei Cui, Sunghun Kim, and Jiayi Huang. Shortcut-connected expert parallelism for accelerating mixture-of-experts. CoRR, abs/2404.05019, 2024.
Chang Chen, Min Li, Zhihua Wu, Dianhai Yu, and Chao Yang. Ta-moe: Topology-aware large scale mixture-of-expert training. In Advances in Neural Information Processing Systems 35: Anmual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. Deepseekmoe: Towards ultimate expert specialization in mixture-of mixture-of-expprts language models. In Proceedings of the 62nd Annual Meeting of the Asso- ciation for Computational Linguistics (Volume I: Long Pupers), ACL 2024, Bangkok, Thailand, August 11-11 16, 2024, pp. 1880-129,,00 2024.
Ran Duan and Seth Pettie. Linear-time approximation for maximum weight matching. J. ACM, 61 (1):1:1-1:23, 2014.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshorn, Aobo Yang. Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurélien Rodriguez. Austen Gregerson, Ava Spataru, Baptiste Rozière, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris
11

---

Published as a conference paper at ICLR 2025
Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Graeme Nail, Grégoire Mialon, Guan Pang. Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, llyaan Zarov, Imanol Arrieta Ibarra, Isabel M. Kloumann, Ishan Misra, Ivan Evtimov, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang. Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, and et al. The llama 3 herd of models. CoRR, abs/2407.21783. 2024.
Artyom Eliseev and Denis Mazur. Fast inference of mixture-of-experts language models with of- floading. CoRR, abs/2312. 17238, 2023.
William Fedus, Barret Zoph, and Noam Shazeer. Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity. J. Mach. Learn. Res., 23:120:1-120:33 2022.
Lei Guan, Dong-Sheng Li, Jiye Liang, Wen-Jian Wang, Ke-shi Ge, and Xicheng Lu. Advances of pipeline model parallelism for deep learning training: An overview. J. Comput. Sci. Technol., 39 (3):6675584, 2024.
Jiaao He, Jiezhong Qiu, Aohan Zeng, Zhilin Yang. Jidong Zhai, and Jie Tang. Fastmoe: fast mixture-of-expert training system. CoRR, 0os/2103.13262 2021.
Jiaao He, Jidong Zhai, Tiago Antunes, Haojie Wang, Fuwen Luo, Shangfeng Shi, and Qin Li. Faster- moe: modeling and optimizing training of large-scale dynamic pre-trained models. In PPoPP '22: 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, Seoul, Re- public of Korea, April 6, 2022, pp. 120-134, 2022.
Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Xu Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoe V. Le, Yonghui Wu, and Zhifeng Chen. Gpipe: Efficient training of giant neural networks using pipeline parallelism. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 103- 2019.
Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang, Ze Liu, Han Hu, Zilong Wang, Rafael Salas, Jithin Jose, Prabhat Ram, Ho Yueu Chau, Peng Cheng, Fan Yang, Mao Yang, and Yongqiang Xiong. Tutel: Adaptive mixture-of -feeerrts at scale. In Proceedings of the Sixth Conference on Machine Learning and Systems, MLSys 2023, Miami, FL, USA, June 4-8, 2023, 2023.
Albert Q. Jiang, Alexandre Sablayrolles, Antoine Roux, Arthur Mensch, Blanche Savary, Chris Bamford, Devendra Singh Chaplot, Diego de Las Casas, Emma Bou Hanna, Florian Bressand, Gi- anna Lengyel, Guillaume Bour, Guillaume Lample, Lélio Renard Lavaud, Lucile Saulnier, Marie- Anne Lachaux, Pierre Stock, Sandeep Subramanian, Sophia Yang, Szymon Antoniak, Teven Le Scao, Théophile Gervet, Thibaut Lavril, Thomas Wang, Timothée Lacroix, and William El Sayed. Mixtral of experts. CoRR, abs/2401. .04088, 2024.
Yoohwan Kwon and Soo-Whan Chung. Mole : Mixture of language experts for multi-lingual au- tomatic speech recognition. In IEEE International Conference on Acoustics, Speech and Signal Processing ICASSP 2023, Rhodes Island, Greece, June 4-10. 2023, pp. 1-5.2023.
Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang. Maxim Krikun, Noam Shazeer, and Zhifeng Chen. Gshard: Scaling giant models with condi- tional computation and automatic sharding. In The ninth International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021, 2021.
Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. BASE layers: Simplifying training of large, sparse models. In International Conference on Machine Learning, ICML. 2021, 18-24 July 2021, Virtual Event, volume 139, pp. 6265-6272 2021.
12

---

Published as a conference paper at ICLR 2025
Jiamin Li, Yimin Jiang, Yibo Zhu, Cong Wang, and Hong Xu. Accelerating distributed moe train- ing and inference with lina. In Proceedings of the 2023 USENIX Annual Technical Conference, USENIX ATC 2023, Boston, MA, USA, July 10-12, 2023. pp. 945-959, 2023.
Jing Li, Zhijie Sun, Xuan He, Li Zeng, Yi Lin, Entong Li, Binfan Zheng, Ronggian Zhao, and Xin Chen, Locmoe: A low-overhead moe for large language model training. CoRR, abs/2401.13920, 2024.
Shen Li, Yanli Zhao, Rohan Varma, Omkar Salpekar, Pieter Noordhuis, Teng Li, Adam Paszke, Jeff Smith, Brian Vaughan, Pritam Damania, and Soumith Chintala. Pytorch distributed: Experiences on accelerating data parallel training. Proc. VLDB Endow., 13(12):3005-3018, 2020.
Hanxue L. Liang. Zhiwen Fan, Rishov Sarkar, Ziyu Jiang, Tianlong Chen, Kai Zou, Yu Cheng, Cong Hao, and Zhangyang Wang. vit Mixture-of-experts vision transformer for efficient multi- task learning with model-accelerator co-design. In Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022.
Juncai Liu, Jessie Hui Wang, and Yimin Jiang. Janus: A unified distributed training framework for sparse mixture-of-experts models. In Proceedings of the ACM SIGCOMM 2023 Conference, ACM SIGCOMM 2023, New York, NY, USA, 10-14 September 2023, pp. 486-498, 2023.
Stuart Mitchell, Michael OSullivan, and lain Dunning. Pulp: linear programming toolkit for python. The University of Auckland, Auckland, New Zealand, 65:25, 2011.
Deepak Narayanan, Amar Phanishayee, Kaiyu Shi, Xie Chen, and Matei Zaharia. Memory-efficient pipeline- pariinee-paallllllle DNN training. In International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp- 7937-7947, 2021a.
Deepak Narayanan, Mohammad Shocybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vi- jay Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, and Matei Zaharia. Efficient large-scale language model training on GPU clusters using megatron- Im. International Conference for High Performance Computing, Networking, Storage and Analysis, SC 2021, St. Louis, Missouri, USA, November 14-19, 2021, pp. 58, 2021b.
Xiaonan Nie, Xupeng Miao, Zilong Wang, Zichao Yang, Jilong Xue, Lingxiao Ma, Gang Cao, and Bin Cui. Flexmoe: Scaling large-scale sparse pre-trained model training via dynamic device placement. Proc. ACM Manag. Data, 1(1): 10:1-110:19 2023.
OpenAl. GPT-4 technical report. CoRR, abs/2303.087/ 087744, 2023.
James B. Orlin and Ravindra K. Ahuja. New scaling algorithms for the assignment and minimum mean cycle problems. Math. Program., 54:41-56, 1992.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Ed- ward Z. Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver. BC, Canada. pp. 8024 8035, 2019.
Pytorch. Gradient accumulation pytorch.
http:://giaa.gggtub..oommo 2019.
a
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9 2019.
Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Su- sano Pinto, Daniel Keysers, and Neil Houlsby. Scaling vision with sparse mixture of experts. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Infor- mation Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 2583-8595, 2021.
13

---

Published as a conference paper at ICLR 2025
Shriram Sarvotham, Rudolf H. Riedi, and Richard G. Baraniuk. Connection-level analysis and modeling of network traffic. In Proceedings of the 1st ACM SIGCOMM Internet Measurement Workshop, IMW 2001, San Francisco, California, USA, November 1-2. 2001, pp. 99-103, 2001.
Teven Le Scao, Angela Fan, Christopher Akiki, Ellie Pavlick, Suzana Ilie, Daniel Hesslow, Roman Castagné, Alexandra Sasha Luccioni, François Yvon, Matthias Gallé, Jonathan Tow, Alexan- der M. Rush, Stella Biderman, Albert Webson, Pawan Sasanka Ammanamanchi, Thomas Wang, Benoît Sagot, Niklas Muennighoff, Albert Villanova del Moral, Olatunji Ruwase, Rachel Baw- den, Stas Bekman, Angelina McMillan-Major, Iz Beltagy, Huu Nguyen, Lucile Saulnier, Samson Tan, Pedro Ortiz Suarez, Victor Sanh, Hugo Laurençon, Yacine Jernite, Julien I Launay, Margaret Mitchell, Colin Raffel, Aaron Gokaslan, Adi Simhi, Aitor Soroa, Alham Fikri Aji, Amit Alfassy, Anna Rogers, Ariel Kreisberg Nitzav, Canwen Xu, Chenghao Mou, Chris Emezue, Christopher Klamm, Colin Leong, Daniel van Strien, David Ifcoluwa Adelani, and al. BLOOM: A 176b- parameter open-access multilingual language model. CoRR, 0bs/2211.05111 2022.
Alexander Sergeev and Mike Del Balso. Horovod: fast and easy distributed deep learning in tensor- flow. CoRR, (0ss//80,,0599, 2018.
Yunfan Shao, Zhichao Geng, Yitao Liu, Junqi Dai, Hang Yan, Fei Yang, Li Zhe, Hujun Bao, and Xipeng Qiu, CPT: a pre-trained unbalanced transformer for both chinese language understanding and generation, Sci. China Inf. Sci., 67(5). 2024.
Hongyan Tang. Junning Liu, Ming Zhao, and Xudong Gong, Progressive layered extraction (PLE); A novel multi-task learning (MTL) model for personalized recommendations. In RecSys 2020: Fourteenth ACM Conference Recommender Systems, Virtual Event, Brazil, September 22-26, 2020, pp. 269-278, 2020.
Peng Tang, Jiacheng Liu, Xiaofeng Hou, Yifei Pu, Jing Wang, Pheng-Ann Heng, Chao Li, and Minyi Guo. HOBBIT: A mixed precision expert offloading system for fast moe inference. CoRR, abs/2411.01453 2024.
Tensorflow. Gradient accumulation tensorflow.
[er.ooffoww/pll//2353 2019.
Hugo Touvron, Louis Martin, Kevin Stone. Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cristian Canton-Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurélien Rodriguez, Robert Stojmic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models. CORR, 0ttp2237700928.. 2023.
Guozheng Wang, Yongmei Lei, Zeyu Zhang, and Cunlu Peng. A communication efficient admm- based distributed algorithm using wwo-dimensional torus grouping allreduce. Data Sci. Eng., 8 (1):61-72, 2023.
Fuzhao Xue, Zian Zheng, Yao Fu, Jinjie Ni, Zangwei Zheng, Wangchunshu Zhou, and Yang You. Openmoe: An early effort on open mixture-of-experts language models. CoRR, abs/2402.00139, 2024.
Zhao You, Shulin Feng, Dan Su, and Dong Yu. Speechmoe2: Mixture-of-expert model with im- proved routing. In IEEE International Conference on Acoustics, Speech and Signal Processing, ICASSP 2022, Virtual and Singapore, 23-27 May 2022, pp. 7217-7221, 2022.
Dianhai Yu, Liang Shen, Hongxiang Hao, Weibao Gong, HuaChao Wu, Jiang Bian, Lirong Dai, and Haoyi Xiong. Moesys: A distributed and efficient mixture-of-expertt training and inference system for internet services. IEEE Trans. Serv. Comput., 7(5):2626-2639 2024.
14

---

Published as a conference paper at ICLR 2025
Zhiyuan Zeng and Deyi Xiong. Scomoe: Efficient mixtures of experts with structured communica- tion. In The Eleventh International Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023, 2023.
Mingshu Zhai, Jiaao He, Zixuan Ma, Zan Zong, Runqing Zhang, and Jidong Zhai. Smartmoe: Ef- ficiently training sparsely-activated models through combining offline and online parallelization. In Proceedings of the 2023 USENIX Annual Technical Conference, USENIX ATC 2023, Boston, MA, USA, July 10-12, 2023, pp. 961-975, 2023.
Huangzhao Zhang, Kechi Zhang, Zhuo Li, Jia Li, Jia Li, Yongmin Li, Yunfei Zhao, Yuqi Zhu, Fang Liu, Ge Li, et al. Deep learning for code generation: a survey. Sci. China Inf. Sci., 67(),, 2024a.
Zhen-Xing Zhang, Yuan-Bo Wen, Han-Qi Lv, Chang Liu, Rui Zhang, Xia-Qing Li, Chao Wang, Zi-Dong Du, Qi Guo, Ling Li, Xue-Hai Zhou, and Yun-Ji Chen. Ai computing systems for llms training: a review. J. Comput. Sci. Technol., 2024b.
Xuanhe Zhou, Zhaoyan Sun, and Guoliang Li. DB-GPT: large language model meets database. Data Sci. Eng., 9()::022-11, 111, 2024.
Xinyu Zou, Zhi Hu, Yiming Zhao, Xuchu Ding, Zhongyi Liu, Chenliang Li, and Aixin Sun. Au- tomatic expert selection for multi-scenario and multi-task search. In Proceedings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval, SIGIR 2022, Madrid, Spain, July 11 15, 2022, pp. 1535-1544, 2022.
15

---

Published as a conference paper at ICLR 2025
A MORE DETAILS OF IMPLEMENTATION
A. DETAIL OF EXPERT RESIDUAL INLINING
As shown in Fig. the original residual addition method adds the attention output to the result obtained from the gather operation. In NetMoE, however, it is added after the scatter operation but before the gather operation. Such an inlining facilitates the adjustment of sample placement, and meanwhile ensures the correctness of computation.
input
input
MHA
=MHA[input)
attn MHA[[nputt]
MHA
moe. input Normlinput attn)
Add Norm
2. jog__nnut= Norm[inpuu attn)
Add Norm
route Gate|moe_ linput)
Gate
3. route Gate(mooejnpput
Gate
4. mmlp_inputt Scatter(route, moe input!
Scatter
4. mimp_ing Inpp_input= Scatter(route, moe input)
Scatter
Residual Inlining
5, mlppp_output output MLP(mip input)
FFN
FFN
output MLP((mpp_input
FFN
FFN
Add
Add
milp_output_ inllined mip_input mip_input
6. moe output Gather(mip_outpput
Gather
Gather
7. moe, outout
=utpptt= Norm(moe_jngut mno_out
Gather(mlp output, inlinee]
Add & Norm
output
Norm
Norm/moe outpur)
Figure 8: Illustration of the Transformer layer with and without the expert residual inlining.
A.2 DISCUSSION OF ALGORITIM SELECTION AND OVERLAP POTENTIAL
Our design adopts the KM algorithm based two practical factors: (1) Although the time com- plexity of the KM algorithm is ( the current training process commonly employs gradient accumulation (Tensorflow) 2019: Pytorch 2019) due to the limited GPU memory. Thus, the value of / is typically confined to an acceptable size, ensuring that the solving time can be effectively overlapped: (2) The algorithm's 's runtime is fully overlapped with communication phases, rendering further acceleration unnecessary for hiding the overhead of solver. While faster approximate solvers exist (Orlin & Ahuja 1992; Duan & Pettic 2014), their benefits would be marginal in current train- ing configurations where computatoon-ommmuicaton overlap already masks the optimization time.
16

---

Published as a conference paper at ICLR 2025
B END To END PERFORMANCE WITH FEWER GPUS PER NODE
Fig. Lillsstraee the end-to-end speedup for configurations with GPUs or GPUs per node. The re- sults demonstrate that NetMoE achieves the best performance across various experimental settings, which are consistent with the results obtained when there are 8 GPUs per node (as demonstrated in Fig.
It is worth noting that standard server configurations typically accommodate up to 8 NVIDIA GPUs per node. Thus, 8 GPUs per node represent standard setup for distributed training of large language models Dubey et al. (2024); Adler et al. (204)); Dai et al. (2024); Scao et al. (2022). Although superpods like the NVIDIA GB200 NVL72 support high-speed connections (e.g., NVLink) among more than 8 GPUs, they rely custom hardware and are prohibitively expensive. Training scenarios on superpods are rare significantly differ from the typical scenarios in GPU clusters or clouds. Therefore, this paper opts for experiments with configurations of up to 8 GPUs per node.
FastMoE
FasterMoE
SmartMoE
NetMoE
= 2,2 GPUS per node
1.5
Speedup
1.0
1331
$5.04
20.2
25.03 26:34
4.4
20.04
-0
20.22
0.5
MoE. GPTS GPTS
MoE GPTM
McE GPT
MoE 6PTXL
MoE GPTXXXL
8.N 2,4 GPUs node
Speedup
1.5-
1:75g
10.00
1.0
1.0%
11000
3:50y
20.3
1.00
0.0
MoE-GPT-S
MoE-GPT-M
MoE-GPT-L
MoE: DDE-GPT-XL
MoE-GPT-XXL
-= N 4, GPUs per node
Speedup
20.00
1.00
3.40%
1000
10
2010
LU0SX
MoE-GPTS
MoE-GPT-M
MoE-GPT-L
MoE-GPT-XL
MoE-GPT-XXL
per
1.5
1.5%
Speedup
1.0
20.00
1074
20.00
40
1.06
Tou
$7.00
2010
20
Aard 84-33
0.5
MoF-CPT-S
MoE-GPT-M
MoE-GPT-L
MoE-GPT-XL
MoF-GPT-XXXL
Figure 9: End-to-end speedup (mean and standard deviation) of different numbers of total devices (denoted as J)) and numbers of nodes (denoted as N).
17

---

Published as a conference paper at ICLR 2025
C DETAILED ANALYSIS OF ALL-TO-ALL COMMUNICATION OPTIMIZATION
To gain deeper insight into the source of NetMoE's optimization, we assess two kinds of statistics:
The proportions of training samples that are exchanged across nodes or across devices by NetMoF, respectively. A higher proportion indicates more samples are adjusted across nodes/devices.
The intra-node and inter-node communication volumes before and after applying NetMoE.
Firstly, Table S summarizes the mean and standard deviation across all iterations. After apply- ing NetMoE, great proportion of training samples are exchanged across nodes, leading to the reduction in the inter-node communication volume. It is noteworthy that although the intra-node communication volume accounts for large proportion (i.e., Sintra Fenera 1 Sunter increases) after applying NetMoF, it will not become the performance bottleneck since the inter-node communica- tion bandwidth is much lower. As a result, the All-to-All communication can be accelerated due to the reduction in inter-node communication volume brought by sample placement adjustment.
Santre
Secondly, since the routing result dynamically changes during the training of MoE models, to dis- cover the impact of routing distribution, in Fig. I0 we plot (I) the reduction in inter-node com- munication, and (2) the proportion of samples exchanged across nodes, across different iterations. Meanwhile, we follow prior works He et al. (2022); Nie al. (203) record the distribution of expert selection across different iterations in order to describe the routing distribution. It can be observed that the routing distribution changes during the model training process. However, NetMoE consistently reduces the inter-node communication by adjusting the sample placement given the dy- namic distributions. Consequently, the effectiveness of NetMoE is robust to the routing distribution. Table 5: Summary of communication volume and proportion of sequence adjustment. For commu- nication volume, provide the intra-node and inter-node communication volumes before and after applying NetMoE, with the increase or reduction given in parentheses. For the proportion of se- quence adjustment, "Across Nodes" indicates the proportion of sequences that exchanged across nodes, and "All" indicates the proportion of all sequences that are adjusted.
(a) nodes, 16 GPUs
Cremmunicatinn Volume
Proportion Sequence
wio NetMoE
NetMol
Nodis
MoF-GPT-S
168.45 5.43
191.07 + 5.41
162.24-1 11 3.69%)
116.37 + 11.69 (1 10%)
43.563 +: 3,560
11.304 + 1.319
McE- MPEM
222.89 1: 5.72
258.20
1: 5:72
214.31
7.42 3.85))
147.33
42,94%)
11.455 1: 2.122
91,929
1.163
MoE-GPTH
5.18
1,080
MoF-GPTTXI
402.06
5.68
313.00
8.30 9.91%)
256.04
36.18)
43,619
+:
1.369
MoE: XXL
922.10
1.:00:211 1.16
989.80
4.16
$72.00 1: 8.29 5.46%)
570.10
5.29
42.36%)
45.688 1: 3.000
92,169
1.294
(b) 4 nodes, 32 GPUs
Communication Volume (MB)
Propoction of Sequence Adjustment (9)
wio NotMo
w/ NotMoE
Across Nodes All
Pinza
Sisser
Finrs
Sunvr
MoB-GPTES
219.21
31.218)
351.15
39.0225)
96.427
0.544
13.00 28.705)
13,00
35.79%)
72.340
McF-GPPHH
250.56 6.74
958.44 + 6.74
376.66
10.47 (1 34.25%)
501.72
10:47
38.269))
72,603
+: 1:249
96.902
0.590
McE-GPPTXL
350.62 7.10
1199.
7.10
423.37
11,69
20.7595)
791.19
11,69
34.0296)
72.217 - 1.499
96.159
0.569
MoE-GPT-XXL
38.839)
18

---

Published as a conference paper at ICLR 2025
McE-GPT-S (16GPU5)
1.00
1.0
0.75
0.50
0.50
0.25
0.25
0.00
20
30
50
20
30
50
19
20
30
40
50
MoE-GPT-M (16GPUS)
1.00
1.00
Scbetbe movero
0.75
0.75
0.50
0.50
0.5
0.25
0.25
0.00
20
0..00
20
30
0.0
30
(16GPUs)
1.00
Lchore-Syreen
1..00
Sestaen Deaerre
3.0
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
30
50
30
50
20
30
McE GPTXL (16GPUs)
Sencien Porpoce
1.0
2.5
20
30
30
MoE-GPT-XXL
1.00
1.00
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
0.0
30
McE-GPT-S (32GPUs)
Sel:dan Treuuery
1.0
0.5
16
20
30
40
50
10
20
30
40
50
10
20
30
40
50
McE GPT-M (32GPUs)
1.00
Duhare saretinng
1.00
Senadton Prsuare
0.75
0.75
0.50
0.50
0.5
0.25
0.25
0.00
0.00
0.0
McE-GPT-L (32GPUs)
1.00
1.00
1.0
0.75
0.50
0.5
0.25
20
30
40
50
10
20
30
40
50
10
20
30
40
50
McE-GPT-LL (32GPUs)
1.00
1.00
0.75
0.75
0.50
0.50
0.5
0.25
0.25
0.00
0.00
0.0
30
MoE-GPTXXL (32GPUs)
1.00
1.00
1.0
0.75
0.75
0.50
0.50
0.25
0.25
0.00
0.00
20
30
50
20
30
50
20
30
Reraiins
Berations
Figure 10: Left: The reduction in inter-node communication volume. Middle: The proportion of
samples exchanged across nodes. Right: The distribution of expert selection (layer 0).
19


Reviews:
[{"review_id": "OSJVSDySzR", "reviewer": "Reviewer_v9Qs", "summary": "This paper presents NetMoE, a novel framework designed to optimize the routing of samples in Mixture of Experts (MoE) models by taking into account the actual inter an intro-node communication bandwidth. The goal is to minimize the *time* the routing process takes, which usually amount to minimize inter-node expert routing in the All-to-All communications, while being mathematically equivalent to the standard routing procedure. This paper formulates the problem as an integer linear programming optimization problem, and relaxes it so that an approximate solution can be found sufficiently fast dynamically at each stage of the MoE. Experimental results demonstrate that NetMoE outperforms existing MoE training systems, achieving up to a 1.67x speedup in training efficiency.", "questions": "see Weaknesses.\n\n**typo:** In table 1 *\"number of of nodes\"*.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "* **The problem is clearly motivated:** The challenges of routing samples in MoE are clearly written, making the goal of this paper feel natural after reading the first two sections.\n* **Challenges of ILP solving are made clear, and the proposed solution seem effective:** The building to the final approximate method is clear and well motivated through empirical results in Tab.4. The optimization gap between the optimal and the approximate solution seem reasonable in Fig.6.\n* **Non negligible empirical benefits of the method are demonstrated:** The speedup brought by NetMoE compared to Dynamic Expert Placement methods seem significant in the experiments displayed.", "weaknesses": "* **Notations and problem formulation hard to follow:** Many notations are introduced, making the reading of section 3 a bit cumbersome. Maybe putting some of the mathematical details and ILP formulations in Appendix could help lighten the section and make it more readable?\n* **No comparison with methods using a modification in the model definition:** While methods introduced in Sec. 2.2 change the convergence property of the model in terms of iterations, the fact that they allow for more iterations per time unit could counter this. Would it be possible to also compare NetMoE to these methods (e.g., in terms of *\"time to reach a certain level of perplexity\"*)?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "EZk8P9a12j", "reviewer": "Reviewer_EW5K", "summary": "The whole idea of NetMoE is that we want to reduce the All-to-All scatter & gather communications by reducing the amount of cross-node/device routing of tokens. To achieve this we will adjust the sample/sequence that would minimize the inter-node & intra-node communication volume. This is (approximately) solvable as a weighted bipartite matching / assignment problem between training samples and machines, as shown in Eqn 9 and 10. \n\nThe authors conduct experiments on GPT pretraining and compare with dynamic expert placement baselines as FasterMoE and SmartMoE. NetMoE generally has higher speedup (Figure 5) and the actual speedup is close to the theoretically optimal speedup (Figure 6).", "questions": "1. The sequence adjustment is done per iteration and per layer and composable of reducing the all-gather communication of this layer and all-scatter of next-layer (Eqn. 7). The reduction from all-gather is clear, but I don't understand how it is even possible to reduce the all-scatter costs of *next-layer* as we even don't know what is the routing probability due to an attention block before the MoE. \n\n2. I don't understand how does expert inline residual fix the position issues of residual stream (it might be helpful to give a diagram as line 12 in Algorithm 1 is not sufficiently clear)", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well motivated and the writing is pretty clear. I have no difficulty on understanding the overall idea of sample adjustment (from Figure 2) and the optimization challenges & solutions (Equation 5, 8, 10) upon the first time of reading. \n\n2. Clever design: reformulating the ILP to a weighted bipartite matching / assignment problem and using Hungarian algorithm that has shorter solving time than communication time (so we can have actual speedup).", "weaknesses": "I don't have strong opposition to the overall idea of sequence adjustments for MoE but I believe the scope and limitations should be more clearly defined:\n\n1. The authors should provide a summary statistics on how many sequences are actually adjusted across nodes/devices during training and how it is correlated with the MoE specialization / router probability. \n\n2. A small-scaled ablation experiment is definitely needed to show if this communication volume reduction is robust w.r.t. the choice of dataset mixtures, as the performance of NetMoE might be data dependent.\n\n3. Table 4 is concerning because the limit of KM algorithm to use less time than all-scatter is $I/J \\sim 24$ (24 is my scaling extrapolation of Table 4's $I/J = 16$ results as KM's time complexity scales cubically w.r.t. # nodes, and $(24/16)^3 * 1 > (24/16) * 2$). A batch size of 24 per device is not a sufficiently large number.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "Xl7mx0dFK4", "reviewer": "Reviewer_rdXW", "summary": "This paper proposes to use a dynamic sample placement to speed up the MoE training. Specifically, this paper adopts a mathematical model to simulate the number of inter-node communication and intra-node communication and solve the integer programming problem to figure out the best sample allocation of the sample to reduce inter-node communication inspired by the locality in networks. This paper successfully reduces the all2all gather communication in training and achieve speed up.", "questions": "See Weaknesses. And\n\nMoving the sample should incurs more movements compared a subset of the tokens in the sample. Why moving sample gives less communication overhead?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper tackles the MoE training efficiency from a novel perspective, that is the data locality perspective. It dynamically locates the data to reduce the inter-node communication in all2all gathering. \n2. The results shows improvements compared with baselines, signifying the effectiveness of the method.\n3. The modeling of the networking problem is inspiring to the reviewer.", "weaknesses": "1. The scalability of the method is questionable, e.g., the improvements for 32 GPUs is smaller then the improvements for 16 GPUs. This leads to the question that what will happen if we continue increasing the number of GPUs? Will the improve converges to zero? \n2. When there are more GPUs, the communication should take a larger portion in the total time? Why the method here, which primarily focuses on optimizing communication, have less significant improvements.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "SWIE9bA3X0", "reviewer": "Reviewer_JYL4", "summary": "The paper proposes a topology-aware sample placement scheduling approach to optimize All-to-All communication in MoE training.", "questions": "•\tHow does the data locality used in the experiment compare to typical training scenarios, and what impact might this have on expected performance?\n\n•\tWhy is inter-node expert parallelism favored over pipeline or other model parallelism techniques in this context?\n\n•\tIs an auxiliary loss mechanism incorporated to mitigate expert selection skew, and if so, does it affect the performance of NetMoE?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "•\tTheoretical Rigor: This paper is thorough in formulating the communication challenges and solution as an optimization problem, with clear problem modeling and a detailed, polynomial-time solution.\n\n•\tPracticality: The method can integrate with existing MoE training systems while enhancing training efficiency.\n\n•\tEmpirical Validation: Experimental results across various configurations validate NetMoE’s improvements in All-to-All communication and overall training efficiency.", "weaknesses": "•\tExperimental Context: The paper could benefit from a more comprehensive discussion on the \"data locality\" conditions required to achieve the claimed speedups in real-world setups. Also, details on the distribution of data locality across real-world training tasks (and the one used in experiment) would give more insight into NetMoE's practical performance.\n\n•\tDiscussion on Experiment Setup: Given that inter-node expert parallelism can incur heavy communication costs, it would help if the authors provided reasoning for prioritizing inter-node expert parallelism over potentially less intensive techniques like a hybrid one: intra-node expert parallelism + inter-node pipeline parallelism. \n\n•\tMore Baseline Comparisons: Additional baselines, particularly concerning dynamic expert placement, would highlight NetMoE’s comparative advantages and limitations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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": "m4l4gosFxJ", "reviewer": "Reviewer_RL6X", "summary": "Communication efficiency is a significant challenge for training efficiency in distributed Mixture of Experts (MoE) models. Unlike other papers that address this issue from a model perspective, this paper offers a solution from a data perspective. It introduces NetMoE, a method that reassigns data samples to different nodes based on locality to minimize all-to-all communication costs. The problem is formulated as an integer programming problem, and the authors derive a polynomial-time solution. Experimental results further validate the effectiveness of their approach.", "questions": "Same as Weaknesses.\nQ1. In what scenarios would one choose a data perspective approach over a model perspective approach?\n\nQ2. Please revise your solution to ensure it aligns with the stated assumptions.\n\nQ3. Explain why the Kuhn-Munkres (KM) algorithm with highest time complexity is the best choice for this problem.\n\nQ4. Conduct additional experiments to demonstrate the impact of node and device variables on performance.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The author demonstrates strong writing skills, clearly stating the problem and solution. The system diagram is also very clear.\n2. They offer a new perspective on communication efficiency in distributed MoE by exploring how data placement can impact efficiency.\n3. Experiments are provided to validate the effectiveness of their approach.", "weaknesses": "1. The motivation is not clearly articulated. In the motivation section, the authors mention that previous works focus on the model perspective and do not explore the data perspective, which does not convey the true motivation. Instead, it should emphasize that in certain scenarios, the model perspective may be insufficient, while a data-focused approach can achieve better efficiency.\n\n2. The problem formulation and subsequent assumptions appear contradictory and I suspect the effectiveness of method. In Equation (1), the communication cost is defined as the maximum of intra-node and inter-node costs. However, in Section 3.2, the authors assume the maximum is the inter-node cost and address it first. This raises questions for the reviewer: if the inter-node assignment is fixed but minimizing intra-node communication results in a higher total cost than inter-node, this may lead an undesirable solution.\n\n3. The authors transform this problem into a weighted bipartite matching problem and solve it using the Kuhn-Munkres (KM) algorithm. However, based on the reviewer's knowledge, KM is sensitive to the sample input and has a time complexity of \nO(N^3) which may not be ideal for large models. The authors should justify their choice of KM as the solver.\n\n4. The experiments do not fully validate the approach. The impact of node and device count on performance is not examined. For instance, if there are very few devices in each node but many nodes overall, inter-node communication may dominate the time. Conversely, if there are numerous devices within fewer nodes, intra-node communication could become the dominant factor in training time.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "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:

Decision: Accept (Spotlight)