FlowMoE: A Scalable Pipeline Scheduling Framework for Distributed Mixture-of-Experts Training
Authors: Yunqi Gao
Bing Hu
Mahdi Boloursaz Mashhadi
A-Long Jin
Yanfeng Zhang
Pei Xiao
Rahim Tafazolli
Merouane Abdelkader DEBBAH
Venue: NeurIPS 2025 poster ()
Keywords: Distributed Deep Learning
Large Language Model
Mixture-of-Experts
Pipeline Scheduling
Primary area: infrastructure

Abstract: The parameter size of modern large language models (LLMs) can be scaled up to the trillion-level via the sparsely-activated Mixture-of-Experts (MoE) technique to avoid excessive increase of the computational costs. To further improve training efficiency, pipelining computation and communication has become a promising solution for distributed MoE training. However, existing work primarily focuses on scheduling tasks within the MoE layer, such as expert computing and all-to-all (A2A) communication, while neglecting other key operations including multi-head attention (MHA) computing, gating, and all-reduce communication. In this paper, we propose FlowMoE, a scalable framework for scheduling multi-type task pipelines. First, FlowMoE constructs a unified pipeline to consistently scheduling MHA computing, gating, expert computing, and A2A communication. Second, FlowMoE introduces a tensor chunk-based priority scheduling mechanism to overlap the all-reduce communication with all computing tasks. We implement FlowMoE as an adaptive and generic framework atop PyTorch. Extensive experiments with 675 typical MoE layers and four real-world MoE models across two GPU clusters demonstrate that our proposed FlowMoE framework outperforms state-of-the-art MoE training frameworks, reducing training time by14%-57%, energy consumption by 10%-39%, and memory usage by 7%-32%. FlowMoE’s code is anonymously available at https://anonymous.4open.science/r/FlowMoE.

FlowMoE: A Scalable Pipeline Scheduling Framework for Distributed Mixture-of-Experts Training
Yunqi Gao', ,ing Hu' Mahdi Boloursaz Mashhadi', A-Long Jin", Yanfeng Zhang', Pei Xiao", Rahim Tafazolli", Mérouane Debbah
School of Information Science and Electronic Engineering, Zhejiang University
2SGIC & 6GIC, Institute for Communication Systems (ICS), University of Surrey
"School of Advanced Technology. Xi' an Blaotong-Liverpov University
"School of Computer Science and Engineering, Northeastern University
KKU 6G Research Center, Department of Computer and Information Engineering, Khalifa University (gaoyunqi1999 binghu] @zju.edu.cc.
im.boloursazmashhadii p.xiao, Itafazolli @ @swrrey.u...... au.u Xonggjin@xiluuuddu..nn zhangyf@ mail. mail..ne..ddie new.i....ed...n merouane.debbarn @ ku.ac.ae
Abstract
The parameter size of modern large language models (LLMs) can be scaled up via the sparsely-activater
offturre-of-xxpett (MoE) technique avoid exces- sive increase of the computational costs. To further improve training efficiency, pipelining computation and communication has become a promising solution for distributed MoE training. However, existing work primarily focuses on scheduling tasks within the MoE layer, such as expert computing and all-to-all (A2A) com- munication, while neglecting other key operations including multi-head attention (MHA) computing, gating, and all-reduce communication. In this paper, we pro- pose FlowMoE, scalable framework for scheduling multi-type task pipelines. First, FlowMo constructs a unified pipeline to consistently scheduling MHA com- puting. gating, expert computing, and A2A communication. Second, FlowMoE introduces a tensor chunk-based priority scheduling mechanism to overlap the all-reduce communication with all computing tasks. We implement FlowMoRR an adaptive and generic framework atop PyTorch. Extensive experiments with 675 typical MoE layers and four real-world MoE models across two GPU clusters demonstrate that our proposed Flow framework outperforms state-of-the-art MoE training frameworks, reducing training time by 13%-57% energy consump- tion by 10%-39% and memory usage by 7%-32%. FlowMoE's code is available at bttps://ge hgt..tttth...../ com Somm/ZJU-Oo CNLAB CNLAB/FIOMOO...
Introduction
Large language models (LLMs) have demonstrated remarkable performance natural language pro- cessing (NLP) tasks as model sizes increase (c.g., GPT-3 [1] with 175 billion parameters, LLaMA3.1 [2] with 405 billion parameters, and DeepSeek-RI [3] with 671 billion parameters). However, parameter scaling causes a linear increase in the computational cost. Currently, sparsely-activated LLMs with Mixture-of-FEi -operrssss layers (MoE models) have become extremely popular [4], where the MoE layer replaces the standard feed-forward layer in traditional transformer blocks. In MoE, gating function selects small subset of dense layers, known as experts, to be activated for the input sample tokens. This dynamic selection enables only a few experts to participate in computation during each iteration [5, 6). Therefore, MoE techniques can scale the model size with limited increase in computation [5, 7, 8]. e.g., Google's Switch Transformer expands parameters from a few billion to
*Corresponiigg author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Table 1: Time for different tasks of each iteration in training four MoE models on 16-GPU (NVIDIA RTX3090) cluster (with 100Gbps bandwidth) running vanilla expert parallelism [19]. 'MHA + Gating Time' indicates the computing time of the MHA layer and the gating function. 'All-Reduce Time' indicates the time of all-reduce communication. 'Ratio' indicates the ratio of the sum of 'MHA + Gating Time and 'All-Reduce Time' over the total time per iteration.
Model
MHA + Gating Time All-Reduce Time
Iteration Time
Ratio
UPT2-Tiny-Monn
23.5ms
32.6ms
169.5ms
33.1%
BERT-Large-Man
61.9m
98.3ms
537.8ms
29.8%
LLaMA2-MoE
308.4ms
368.8ms
1987.7ms
34.2%
DeepSeek-V2-S
870.2ms
1247.8ms
5843.3ms
36.1% %
1.5 trillion through 15 MoE layers with 2048 experts each [9]. However, training MoE models on large-scale GPU clusters still suffers from serious scalability bottlenecks [10, 11, 12].
Recently, expert parallelism has been proposed to train MoE models in a distributed fashion by placing different experts on multiple workers since single worker (e.g., GPU) cannot hold complete MoE model [5]. In each iteration, input tokens need to be transferred to particular workers, which depends expert computing on different workers need to be collected via another A2A operation (combine) [13]. In addition, the parameters of the remaining parts of the model including the multi-head attention (MHA) layer and the gating function are replicated across all workers to perform data parallelism mode [14, 15, 16, 17, 18]. In particular, the parameters of the MHA layer and the gating function need to be synchronized among all workers using the all-reduce communication after backward propagation.
all-to-all (A2A) communication operation (dispatch), and the outputs from
Since the computing and communication tasks do not occupy the same hardware resources, pipelining
of computing and communication becomes one of the most efficient methods for accelerating expert parallelism. Existing works (e.g.. ScheMoE [10], Tutel [12], FasterMoE [11], Lina [20], PipeMoE [21]. Comet [22]) pipeline expert computing tasks and A2A communication tasks to hide the communication and reduce the training time of the MoE model by splitting the input data of the MoE layer into
However, they focus only overlapping tasks in the MoE layer and neglect MHA layer computing, gating, and all-reduce communication. We conducted experiments on 16-GPU cluster and the results are shown in Table (see Table 2 for model configurations). It can be observed that the MHA layer computing, gating, and all-reduce communication constitute 30%-40% of the iteration time. Therefore, designing pipeline scheduling method that considers all major tasks of the transformer block with MoE layer can maximize the overlap between computing and communication, thus improving the scaling efficiency of distributed MoE training.
In this paper, we propose FlowMoE, a scalable pipeline scheduling framework for multi-type tasks in distributed MoE training. We make the following main technical contributions: (1) We construct unified pipeline that consistently schedules MHA computing, gating, expert computing, and A2A communication. (2) We design priority-based scheduling mechanism for communication tasks using all-reduce tensor chunks to further overlap the all-reduce communication with computing tasks, and we leverage Bayesian optimization (BO) to automatically tune the partition size of all all-reduce chunks. (3) We implement the FlowMoR framework on PyTorch [23] and open-source the code. Extensive experiments on two GPU clusters using manually customized MoE layers and real-world consformer-based MoE models show that FlowMoE achieves better training performance than state-of-the-art MoE frameworks (including ScheMoE [10], FSMoE [24], Tutel [12] and FasterMoB [1I]) and vanilla expert parallelism [19]. Specifically, Flo outperforms ScheMooE by 26% average time efficiency in training 675 MoE layers with different configurations. In comparison with state-of-the--u frameworks in training real-world popular MoE models, FlowMoE achieves 1.13x-1.82 speedup, and reduces energy consumption by 10%-41% and memory usage by 7%%32% during each training iteration. The main differences between FlowMoE and the key literature is shown in Appendix A.
2 Background and Challenges
2.1 Transformer Block with MoE Layer
Fig. Ia illustrates a typical transformer structure with MoE layers, where the transformer block usually consists of an MHA layer and an MoE layer. For the input tensor I E 18 BxBxxxM the MHA

---

One Tronsfermer Backk
Expem
Mal: Haad
Geting
Exprrt
MJB Hood
Ozting
Mali-Hrued
Expert
Malisthaal
Loger
Expot22
Layer
Gatling
Aeplicane
E-3
3-
ExpestE-
cPL
200mrio Laget
Equet
Lager
(a) transformer structure with MoE layers.
(b) An illustration of expert parallelism.
Figure An example a transformer block with MoE layer and expert parallelism.
Feed-forward computing(FF)
Backward pwapppppppppiiooooomupipigggggr
Ali-to-all Communication
All-reduce Communication
FF without pipelining.
BP without pipelining.
<
ABT
(c) PF with R 2 in MoE layer.
(d) BP with I 2 in the MoE layer.
(e) FF with 2 in MHA and
(f) BP with R 2 in MHA and MoE layer priority based scheduling with all-reduce tensor chunks.
MoE layer.
Figure 2: An example of the execution timelines of the computing and communication tasks for model consisting of two transformer blocks with MoE layer.
layer utilizes the Query, Key, and Value matrices Wº, WK WV E MxM to calculate the attention data of each token and obtains the output tensor L' E IR Nxx through a linear transformation matrix p E Mad, where B denotes the number samples per GPU (or mini-batch size) in one iteration, N denotes the number of tokens per sample, and M denotes the embedding size of a token. The MoE layer includes a gating function and multiple experts. The gating function G is small learnable neural network followed by a softmax layer and is used to select the activated experts for the input tokens. Typically, only top-k experts are chosen to process a token [9]. The output Mep
tensor G(T') EXC of the gating function will be dispatched to the corresponding expert (each expert receives a tensor of shape x M), where denotes the total number of experts per MoE layer and C' denotes the maximum number of tokens assigned to one expert. C' can be calculated by k x x N/E, where f the capacity factor that determines the maximum number of tokens assigned each expert and is used to adjust C' [5]. An expert is usually small neural network with two structurally symmetric feed-forward layers (the first and the second layers are of sizes M H and H x M, respectively). where H denotes the hidden size of the feed-forward layer, and each expert is considered to have its own domain of expertise. After expert computing, the outputs from all the experts combined into tensor with shape of B x N x M as the input to the next transformer block.
2.2 Expert Parallelism
We define P as the number of workers (or GPUs) in the cluster, L as the number of transformer blocks in MoE model. Fig. 1b shows example of training MoE model on P workers by expert parallelism, where each worker has different expert parameters. Then, AT,(1) Fn!!) D. and (i()) respectively denote the rth subtask of MHA layer (including gating function), expert computing, dispatch A2A communication, and combining A2A communication of the ith transformer block. AR(!) denotes all-reduce communication task of the ith transformer block. Figs. 2a and 2b show the timelines of multiple computing and communication tasks in the feed-forward computing and backward propagation for one transformer block in vanilla expert parallelism [19]. is worth noting that the timeline for backward propagation is opposite to that of feed-forward computing, but the parameters of the MHA layer and the gating function need to synchronized through additional all-reduce communication operation.
2.3 Performance Bottlenecks in Distributed MoE Training
Most of existing pipeline scheduling works (e.g., ScheMoE [10], Tutel [12), PipeMoE [21]) partitions
the input token tensor of the MoE layer in the data dimension according to the pipelining degree R. where Figs. 2c and 2d illustrate pipelining example with R 2 in the MoE layer of
3

---

g
Stap1: Task breskdown
Step2: Pipeline scheduling
Pepals A3A
1111 Profiing
McE
Figure 3: Workflow of FlowMoE.
transformer block. In addition, FasterMoE [1I] splits the input tensor of the MoE layer based on the number of workers, enabling point-to-point communication between workers. Although these works effectively reduce the training time by pipelining the expert computing task and the A2A communication task, they neglect the MHA layer computing task, gating task and all-reduce communication task. Moreover, some MoE frameworks (e.g. FSMoE [24] and Lina [20]) also explore the pipeline of all-reduce communication tasks. Nevertheless, FSMoE focuses more on inter- and intra-node communication overlaps within the MoE layer, while Lina only optimizes MoF-layer-specife communication bottlenecks and not the entire transformer block. Therefore, we aim to pipeline all major computing and communication tasks across the transformer block to minimize the per-iteration time, which involves three main challenges:
Complex dependencies between multi-type tasks. In distributed MoE training, there exists highly complex dependencies between the computing and communication tasks (see Figs. 2a and 2b) [10]. In particular, the typical parallelization schemes (e.g.. pipeline parallelism [25, 26), tensor parallelism [27, 28, 29I are difficult to be overlaid directly on expert parallelism.
Coexistence of A2A communication and all-reduce communication. Although there has been extensive studies [30, 16, 31, 15, 32, 33, 34] providing efficient scheduling algorithms for data- parallel training of all-reduce communication tasks, they focus on traditional convolution- neural networks LLMs without MoE layers. These studies cannot be directly applied to distributed MoE training because the additional A2A communication tasks are not equivalent to all-reduce communication tasks (see Fig. 2b).
deep
Designing of an adaptive and generic pipeline scheduling framework. Currently, the performance of most scheduling frameworks depends on some hyperparameters setting [35, 36, 37], which introduces additional workload for tuning them before training. Ideally, an adaptive framework should automatically tune hyperparameters and be directly deployable by modifying only the model definition or dataset interface. Meanwhile, a general framework should also be designed to be compatible with different optimization frameworks and communication stacks.
In this paper, our main goal to address the above three challenges by proposing a scalable pipeline scheduling framework called FlowMoE for multi-type tasks in distributed MoE training. All frequently used notations are summarized in Table A. of the Appendix.
3 FlowMoE
3.1
Overview
Fig. 3 illustrates the workflow of FlowMoE. FlowMoE performs efficient pipeline training leveraging the given MoE structure and the dataset. First, multi-type computing and communication tasks are broken down according the pipelining degree and wait for pipeline scheduling. Second, FlowMoE. schedules all computing and communication tasks according to their dependencies and defined priority for distributed training the GPU cluster (see details in Secs. 3.2 and 3.3). Third, Flo automatically tunes the partition size of all-reduce tensor based on the first several iterations during training by BO profiling (see details in Sec. 4.
3.2 Pipeline for MHA and MoE Layer
Based on the pipeline of the MoE layer. FlowMoE first pipelines the MHA layer computing and gating. It is worth noting that for the Ith transformer block, the gating task only depends on the computing task of the MHA layer, and in the subsequent expression, we consider the computing task of both as a whole and denote it by AT (1 < < L). FlowMoE partitions the input tensor of each transformer block into R equal-sized parts, and each computing or communicating task (except all-reduce communicating task) in the MHA layer and the MoE layer will be partitioned into R independent subtasks. In particular, the tasks with the same type have the same execution time. The partitioned tasks can be represented as the set
T T= T= AT, D.D E.D 0. I snsn) n).
S
(I)

---

where AT and
E are computing tasks, and D chO) and AR are communication tasks.
Then, for 1 < L, during feed-forward computing, the scheduling order of computing tasks can be denoted as
AT AT (1) >>..->>ATT -> ->> >>.""..---> ~>AT (1+1)
(1+1)
2..
ERR
(2)
the scheduling order of A2A communication tasks can be denoted as
--. -"
11+1)
(3)
During backward propagation, the scheduling order of computing tasks can be denoted R ((+1) >I ->E -
R-1 2 AT
(4)
the scheduling order of A2A communication tasks can be denoted as
------------>0 - -D
(((+1) >>.->>D ((+1)) -
(1) R- R-
(5)
Figs. 2e and 2f show the example of execution timelines of the computing and communication
tasks when pipelining the MHA layer and the MoE layer with R = 2. Compared to pipelining the MoE layer only, the MHA layer computing task and gating task can be overlapped with the A2A communication tasks, thus shortening the per-iteration time.
3.3 Pipeline for All-reduce Communication
Existing state-of-the-an scheduling frameworks centrally executes all-reduce communication tasks at the end of backward propagation for each iteration (we call it centralized scheduling of all- reduce communication tasks). To reduce all-reduce communication time, we further consider the pipeline of all-reduce communication tasks. Intuitively, in the backward propagation, the all-reduce communication task of transformer block I can overlap with computing tasks of transformer block I since they both depend on the completion of MHA layer computing tasks of transformer block 1. However, in actual training, the all-reduce communication task of transformer block will conflict with the A2A communication tasks of transformer block 1. Therefore, we first mathematically model the timeline of backward propagation in one iteration to find the optimal scheduling of the two communication tasks that minimizes the training time. Considering resource competition in real training environments, we assume that only computing and communication tasks can be executed simultaneously on a GPU, while multiple computing or multiple communication tasks cannot be run simultaneously. Furthermore, there is preemption between tasks: once a task starts execution, it must run completion without interruption. We define (-) as the beginning execution timestamp of a task during the backward propagation, and (-()) as the elapsed time of task during the backward propagation. According to Fig. 2b, the objective function and the dependencies between tasks can be expressed as follows (1 < S R):
min
To = To (AR()) ( +1((AR)) T(((()
(6)
(C((-l)) 2 ( <II, T.
(6a)
Th (E())) 2 (C.))) +
(6b)
7 (D(")) 2 +t(E),)1((())
(6c)
7 (AT, 2 (
(6d)
(AR()) +
(6e)
We use the sign to denote the timeline for centralized scheduling of all-reduce communication tasks. Now, let T3 denote the backward propagation time when the all-reduce communication task of transformer block is inserted between any two A2A communication tasks. Also denote by T the backward propagation time under centralized scheduling of all-reduce communication tasks. With this notation, we have the following theorem:
Theorem 1. If the scheduling order satisfies Eqs. 4 and then we have To T Proof. Theorem proof is provided in Appendix B.
According to Theorem 1, we find that partitioning and inserting all-reduce communication tasks into the gaps between any of the A2A communication tasks can increase the overlap between all-reduce communication tasks and computing tasks. Thus, we design a priority scheduling mechanism of communication tasks based on the all-reduce tensor chunks in FlowMoF. Specifically, during the backward propagation, FlowMoE first slices the tensor for the all-reduce communication tasks of
5

---

Algorithm Training process in FlowMoR
Algorithm 2 Communication pool manage-
Input: Dataset D = ((11,:1),, (Zn, Vo)).4 L.
ment
Output:
Model Weights W = =( (W1.-,,W.
1: / Partition an all-reduce tensor and enqueue tensor chunks. +/
1: Initialize DataQueue for data transferred between tasks;
2: Initialize A2AQueue for A2A communication tasks;
2: Get Sg from BO;
Initialize ARQueue for all-reduce communication tasks;
3: procedure PARTITION( A RRTensor)
4: Split D for d1,,......... d2..... dg into DataQueue:
4: ARChunks ART ensor cettoorprattiin(Spp.
5: Communication pool management start()
5: ARQueue putt A (RCu(ARCCmmaa)
6: for = > dw//Fe/ddfooreeee forward Computing
f Communication task priority scheduling
7: for 1 R do
7: procedure COMMPOOL.ANNGEER
8: d, DetTTFFFDmmDDtaauuuue...c.o
while True de
9: A2AQueue.putcc->:
9:
if A2AQueue not empty then
10:
for r = + R do
10:
d= A2A.Comm( Com.mmAZAQuuue get();
11:
d.=EE.FFcommpDattaauuee
11:
DataQucce. (ut(d);
12:
AZAQuueu.pputdd>x
12:
else if AR AuQuee is not empty then
13:
for
T. do//Backward Propagation
13:
AR.Comm( CowmmmoRuuueeeeeeeeeeeeeeee get();;
14:
for = 1 do
15:
A2.AQucue.puu
16:
for = 1 >> R do
Semalo Eae Cuitatear lncered
17:
18:
A2AAQueuepput(d->
0.40
19:
for -=r=1 1. IR do
20:
10
(18)
21:
DataQueue.puttd-))
22:
Waiting until all-reduce communication is finished;
Figure 4: BO example: tuning Sp for training BERT-Large-MoE with a 16-GPU cluster.
23:
Update W:
each layer into tensor chunks with size Sp- Second, FlowMoF. maintains a communication task pool which includes all the all-reduce tensor chunks and A2A communication tasks, and sets the scheduling priority of all-reduce tensor chunks to be lower than that of A2A communication tasks. In other words, the communication tasks of the all-reduce tensor chunks will be executed immediately when there is no A2A communication tasks. Fig. 2f shows an example of execution timeline when using 2-degree pipelining in MHA layer and MoE layer as well as priority scheduling mechanism of communication tasks based on all-reduce tensor chunks. It can be observed that the all-reduce tensor chunks fully occupy the gap between the A2A communication tasks and thereby maximize the overlap between the computing tasks and the communication tasks.
4 Parameter Optimization, Algorithm Design and System Implementation
4.1 Auto-Tuning Partition Size by Bayesian Optimization
Ideally, if the communication of all-reduce tensor chunks does not introduce extra startup overhead, we have the following theorem:
When the scheduling order satisfies Eqs. 4 and 5, and using the priority scheduling mechanism of communication tasks based on all-reduce tensor chunks, the time per iteration will be minimized if Sp + 0 and there is startup overhead for communicating all-reduce tensor chunks,
Theorem 2.
Proof. Theorem 2 proof is provided in Appendix C.
In the actual training, the communication tasks of all-reduce tensor chunks will introduce extra startup overhead [15]. Therefore, the partition size of all-reduce tensor chunks, Sp, becomes the knob that determines the trade-off between optimal scheduling and system overhead. Since it is non-trivial to explicitly model the iteration time as a function of Sp. to guarantee the adaptivity of FoEE we adopt BO to automatically tune Sp during training, which attempts to find good parameters of an unknown objective function in as few number of trials as possible. BO's parameter settings, importance, and performance evaluation are described in detail in Appendix D. The goal of BO in FlowMoE is to minimize the per-iteration time. Specifically, BO fits an objective function by sampling different (Sp, per-iteration time) pairs and continuously suggests the next S, to obtain a new objective function value, where the per-iteration time corresponding to different S,, in all pairs is measured from the average of multiple iterations (e.g.. 10 iterations). In other words, BO can accurately predict oparmol Sp based on enough samples. Fig. 4 illustrates example of using BO to tune Sp- In the actual training, with only 8 samples, BO returns a near-optimal value at 2.5MB with good confidence.

---

Table 2: Benchmark models.
# Params
# Params
Configurations
MoE Model
Dataset
(MHA+Gating)
(Experts)
L
N
M
H
E/P
K
GPT2-Tiny-MoE
3.2M
50.4M
OpenWebText
12
256
256
512
2
1.0
BERT-Large-Mod
25.2M
806.5M
wikitext- 103
24
512
512 10242
1
1.0
LLaMA2-MoE
134.2M
4297.6M
wikitext-103
32
512
102440961
1.0
LLaMA2-MoE-L
268.4M
8595.2M
wikitext-103
64
512
1024
1
1.0
DeepSeek-V2-5
419.6M
2014.IM
OpenWebText
4
256
5120
15362
8
1.0
DeepSeek- V2-M
734.3M
3524.7M
Open WebText
7
4
256
5120 15362
1 1.0
Tutel Task Brokkereennnnna Maragri
FlowNoE Core
MPI
Poot
User
Pylorch API
Pipline Scretring Hanager
Pylorch Engine
NCCL
Scripts
medeaa Benatr
Ethernet/ InfiniBand
GO Tunter
Figure 5: Overview of Foo architecture (dark gray parts are new).
4.2 Algorithm Design
Algorithm formally provides the pipeline scheduling training process of R-degree computing and communication tasks using FlowMoE.. Moreover, Algorithm 2 provides the communication pool man- agement, including the splitting of all-reduce tensors and the priority scheduling of communication tasks. A detailed description of the two algorithms is provided in Appendix E.
Compared to the vanilla expert parallelism [19], the overhead of pipeline scheduling in FlowMoF mainly comes from two parts: the partition of all-reduce tensors and the objective function fitting during BO. First, when the structure of the gating function is a typical linear layer of M x E, the total number of parameters in the MHA layer and gating function is 4M² + M x E. According to Algorithm 2, the computational complexity of the scheduling algorithm of FlowMoE in each iteration is O(Lx x Mxx+ which is acceptable because it is less than I % one iteration time in practical experiments, and the time of all-reduce tensor partitioning can be further overlapped with the time of some computin/ccommunicttioo tasks (we discuss its scalability in Appendix E).
Sp
Second, BO utilizes the first several iterations of the training process to automatically tune Sp to a near-optimal value. In our experiments, BO samples 8 values of Sp and record the iteration time corresponding to each value by averaging 10 iterations. The computational overhead introduced by BO is negligible compared the total training time, quantified in Appendix D. Additionally, if the hardware environment changes, BO will be re-executed to tune Sp, which is discussed in Appendix K. These details are omitted from the presentation of Algorithm focus more on the pipeline scheduling process of Flow
4.3 System Implementation
We deploy FlowMoE in Py with its API, which usually supports class inheritance with flexible Python language and has a high generality to compatible with different optimization frameworks and communication stacks. In particular, MoE acceleration library that is deeply integrated into Py Torch and supports asynchronous execution of communication and computing tasks. Tutel has also been used default MoE training module by DeepSpeed [38]. Fig. shows the overview of FlowMoE architecture (the left side is closer to the user level), where its detailed description and technical implementations can be found in Appendix F.
implement FlowMoE atop Tutel [12], highly optimized
5 Evaluation
5.1 Experimental Settings
Testbed setup. We use two clusters. (1) Cluster consists of 2 nodes connected with 100Gb/s bandwidth. Each node is equipped with 8 NVIDIA RTX3090 GPUs (24 GB of memory per GPU) connected with PCIe3.0x16. The CPU is Intel Xeon(R) Gold 6248R. (2) Cluster 2 consists of nodes connected with 10Gb/s bandwidth. Each node is equipped with 2 NVIDIA RTX2080Ti GPUs (12 GB of memory per GPU) connected with PCIe switches. The CPU is Intel Xeon(R) Gold 5118.
Models with MoE Layers. (1) Customized MoE layers: We cover typical configurations of MoE layers by choosing combinations of input parameters, where B E (2,4,8), / E (1.0, 1.1, 1.2), N (512, 1024, 2048). M e {512, 1024, 2048, 4096, 8192} and H E (512, 1024, 2048, 4096, 8192).
7

---

Table 3: Comparison of average per-iteration time in milliseconds. S1, S2, S3, S4 and S5 are the speedups of FlowMoE over ScheMoE, FSMoE, Tutel, FasterMoE, and vanillaEP, respectively.
# of GPUs
Time (ms)
Model
Ss
S4
S3
S2
$1
vanillaEP FasterMoB
Tutel
FSMoE ScheMoE FlowMoE
GPT2-Tiny-MotE
104.2
90.1
85.6
82.8
80.8
66.1
1,58x1.36x1 1.
BERT-Large-MoL
373.1
300.9
314.6
278.1
273.6
239.5
1.66X126xxL3Ixx166xXX.
LLaMA2-MoE
1262.3
1187.5
1029.5
928.1
957.7
763.9
1.65x 11.25
DeepSeek-V22 V2-S
2789.7
2261.4
2276.5
2096.3
2166.1
1740.8
1.60x1.29xx.33xx122xx1...2
GPT2-Tiny-MP MoE
125.1
119.3
116.8
98.8
107.5
87.6
1,9xXLB66x1.3xxxIIxx1.222
BERT-Lange-MoE
428.8
354.4
377.5
345.1
331.2
283.2
1.11x1125x1.33xx222XX1177
LLaMA2-MoE
1563.3
1427.7
1187.9
1110.4
1089.9
906.0
1.73xL157x133xx122xx1.2
DeepSeek-V2-5 4037.8
3370.5
3351.7
2985.4
3138.3
2384.9
1.69x1.411x1.39xx125x1.....
GPT2-Tiny-Mod
169.5
135.3
129.3
114.8
116.4
95.6
1.77xL441x1.35xx12xx112222
BERT-Large-MoE
537.8
490.8
501.1
421.9
405.6
351.9
1.53x1.9Xx14DxxIIIXXI.B
16
1.IaMA2MoOE 1987.7
1759.1
1534.1
1292.6
1374.3
1124.0
DeepSeck-V24.
5843.3
4562.5
4481.4
3895.6
4093.7
3205.3
1.82x1.42x..9xx1.222
Table Average per-iteraune per-iteration time with different pipelining degrees on DeepSeek- V2-S. S1 and S2 are the speedups of FlowMoE over Tutel and ScheMoE, respectively.
Time (ms)
S2
S
Tutel
ScheMoE
FlowMoEE
2
4481.4
4093.7
3205.3
I.39x
1.28
4
4628.2
4164.0
3113.8
1.48x
1.33
8
4588.9
4308.7
3295.9
1.39x
1.30x
25
25
cases 85
20
1 20
10
10
Poo 1.05 1.10 1 1.20
1.25 1.30 1.35 1.40 1.45
foo 1.05 1.0 1:15 20 125 1.30 35 1.40 1.45 Speedup over ScheMoE
Speedup over ScheMoE
(a) Speedup Cluster 1.
(b) Specdup Cluster 2.
Figure 6: Statistic of the speedup over ScheMoE.
We set the number of experts equal to the number of GPUs (i.e., E= = P) and k: == 2. (2) Real-world and DeepScek-V2-S for the language modeling task on the OpenWetTex. dataset [39], and BERT-Large-MoE and LLaMA2-MoE for the text generation task on the wikitext-103 103 dataset [40]. We replace all feed-forward layers in GPT2-Tiny [41], BERT-Large [42], and LLaMA2 [43] with MoE layers to construct MoE models. Meanwhile, due to restricted memory, DeepSeek- V2-S only keeps four transformer blocks with MoE layers [44]. The detailed configuration is shown in Table 2. All parameters and gradients of MoE models stored as 32-bi single precision floating point numbers.
MoE models: We also choose four popular models, i.e., GPT2-Tiny-MM
Baselines. We compare FlowMoB with PyTorch-basee vanilla expert parallelism (vanillaFPP) [19] and existing state-of-the-ard MoF. frameworks, including ScheMoF. [10], FSMoE [24), Tutel [12] and FasterMoF, [11], with a pipelining degree of R unless specified otherwise. We use the official implementations of these frameworks. We focus on per-iteration training time, energy consumption and memory usage as performance metries. All the reported numbers are averaged over 1000 iterations.
5.2 Experimental Results
End-to-end Time on Real-world MoE Models. We measure the average iteration time for end-to- end training of the models listed in Table 2 on Cluster 1. The experimental results are shown in Table The results show that FlowMoE: obtains the best scalability and the shortest per-iteration time on different sizes of clusters and models. Specifically, FlowMoE is 14%-31% 13%-25%, 29%-42%.

---

Table 5: Per-iteration time of the MoE layer under different components. The reported speedup
values are based on vanillæEP as the baseline. "w/ Pipe-MoE'' indicates pipelining expert computing
and A2A communication. 'w/ Pipe-AT" indicates pipelining MHA layer computing and gating. 'w/
Pipe-AR' indicates pipelining all-reduce communication.
Name
w/ Pipe-MoE
w/ Pipe-AT
w/ Pipe-AR
Time (ms)
Speedup
vanillaEP
1630.8
1.0x
Tutel
X X
1115.2
1,46x
FlowMoE-AT
V
1012.6
1.61
FlowMoE-AR
X
/w/o BO)
971.5
1.68x
FlowMoE- (RRBOO
X
(w/ BO)
895.3
1.82x
FlowMoE
796.1
2.05x
26%-57% and 43%-82% faster than ScheMoF, FSMOE, Tutel, FasterMoB and vanillaEP, respectively. We also execute stress tests on different training frameworks across two scaled-up MoE. models (LLaMA2-MoE-LL and DeepSeek- -V--M) in Appendix G. Moreover, in most cases, Tutel faster than FasterMoE due to its dedicated optimization for MoE training. Both FlowMoE and ScheMoE are built atop Tutel and outperform Tutel. ScheMoE and FSMoE further reduce training time by optimizing scheduling within the MoE layer and pipelining intra- and inter-node communication tasks, and this strategy can also be integrated into MoE. contrast, our proposed Foo maximally overlaps communication with computing by pipelining all types of tasks. In bandwidth- limited scenarios, when the number of GPUs scales up. pipeline scheduling benefits less and longer communication time dominates the whole training process, thus the advantages of FlowMoE and other state-of-the-ee art MoE frameworks are not obvious. Moreover, Foo essentially only changes the scheduling order of different tasks and does not affect the model convergence. We verify it through theoretical analysis and practical convergence experiments in Appendix H. Furthermore, we discuss the performance lower bound of FlowMoE in Appendix and report the GPU SM utilization when training with FlowMoE in detail in Appendix J. We evaluate the sensitivity of FlowMoE to Sp selection and BO hyperparameters, while measure the computational overhead of BO in Appendix D. We also discuss the robustness of Flow MoE to heterogeneous clusters, dynamic hardware environments, and node dropouts in Appendix K.
Pipelining degree v.s. Training Speed. Table 4 presents the averaged iteration time with different R values on Cluster with 16 GPUs. FlowMoE consistently outperforms ScheMoE and Tutel. Prior work [21] has provided an effective method to select the optimal R by balancing overlapping degree and startup overhead of the tasks, which be directly applied to FlowMoE. Orthogonal to the methods in choosing the best R. our scheduling framework focuses on pipelining all computing and communication tasks efficiently for any given R 2.
Speedup on Customized MoE. Layers. We construct the MoE layer using a combination of input parameters from customized MoE layer. We collect 490 valid cases from Cluster I with 16 GPUs and 393 cases from Cluster 2 with 8 GPUs (excluding -ft-of-memory cases) successfully measured on ScheMoE and FlowMoE, respectively. FlowMoE is always faster than ScheMoE in all valid cases. The speedup statistic of FlowMoE over ScheMoE is shown in Fig. 6. On average, FlowMoE achieves 26% improvement over ScheMoE.
Ablation Study. We conduct the ablation study on a customized MoE layer with B = f= 1.2. N 512 M = 8192, and H = 8192. The per-iteration time on Cluster 1 with different components using 16 GPUs is shown Table 5. It is seen that -oo-ATE-AT improve the time performance by 10.3% over Tutel by pipelining the MHA layer computing and gating. Further, FlowMoE-AR utilizes the priority scheduling mechanism of communication tasks based on all-reduce tensor chunks to pipeline the all-reduce communication and improve the time performance by another 24.6% over Tutel. In particular, our designed BO finds the near-optimal S, and FlowMoE-AR((O)) obtains 8.3% improvement in training time compared to FlowMoE-AR (where we set Sp IMB). Putting our two optimizations together, FlowMoE runs 1.4x and 2.05 x faster than Tutel and vanillaEP, respectively.
Energy Consumption. We use the NVIDIA SMI tool to sample the real-time power of each GPU in the cluster every 5ms, then integrate these sampled data over time to calculate the energy consumption of each GPU during the training process. We sum the energy consumption of all GPUs to obtain the total energy consumption, then divide by the number of GPUs and the number of iterations to obtain the average per-worker energy consumption as shown in Table 6. The results indicate that

---

Table 6: Averaged per-worker energy consumption and memory usage in one iteration.
Model
vanillaFP
FasterMoE
Tutel
ScheMoE
FlowMoE
GPT2-Tiny-MoF:
1.7J/2.45GB
L5J//267GBB
1.3J/2.52GB
1,23/2.46GB
1.01/2.42GB
BERT-Large-Monn
5.93/4.97GB 1.5I/4.19GB
4.1J/4.16GB 1.1J/4.I3GB
3.71/3.89GB
LLaMA2-MoE
20.21/12.43BB
19.91/16.11GB
15.61/11.59BB
14.0J/11.81GB
12.1J/11.0100
DeepScck-V2-3
59.51/19.42GB 54.91/20.93GB
45.6//19.34GB 41.71/18.92GB 34.9U/17.57GB
higher overlapping degree of computing and communication tasks leads to higher computing and communication resource utilization, thus exhibiting lower energy consumption. Overall, FlowMoE saves 10%-16%% 22%-27% 33%-39% and 33%-41 % energy consumption compared to ScheMoE, Tutel, FasterMoE and vanillaEP, respectively. It is worth noting that the power reported by the NVIDIA SMI tool includes both communication and computing costs, According to the official documentation of NVIDIA SMI [45], the energy consumption measured by the NVIDIA SMI tool covers the entire GPU card's 's power consumption, where communication power includes (1) PCIe communication costs, (2) NVLink/NV costs (e.g., NCCL All- Reduce and A2A). Therefore, our mcasurement results effectively reflect the total cluster-wide energy usage, compatibly providing fair comparison with the baselines.
communication costs, and (3) collective communication
Memory Usage. Table 6 also demonstrates average memory usage per worker by monitoring all GPUs every Is Cluster I with 16 GPUs using the NVIDIA SMI tool. FlowMoE has the lowest memory usage because it performs all-reduce communication tasks in time to reduce gradient caching. Specifically, FlowMoE saves up to 7%, 9%, 32% and 11% of memory compared to ScheMoE, Tutel, FasterMoE and vanillaEP,, respectively. In addition, the memory usage of ScheMoE and Tutel similar to that of vanillaEP since there is optimization in the memory. FasterMoE needs to replicate expert parameters among workers in its load balancing mechanism and consumes more GPU memory.
6 Related Work
In addition to the frameworks mentioned in Section 2 that focus on caaputing--ommmniiciin pipelines in distributed training, we also discuss some other works. First, general pipeline schedulers such as PipeDream [26] and Gpipe [25] enable pipelining by splitting the model across multiple GPUs, requiring communication between GPUs to exchange activation values gradients across split layers. In contrast, Flo implements pipelining of computing-communiin
tasks, accelerating model training by minimizing communication overhead. Second, several orthogonal MoE optimization techniques can be combined with FlowMoF. (1) NetMoE [46] alleviates communication bottlenecks and load imbalance by reducing cross-node token routing and dynamically adjusting expert placement. (2) Lancet [47] determines task scheduling order by constructing a camputing-communication directed acyclic graph (DAG), without considering mierobatch-level partitioning of the same task. (3) Punniyamurthy et al.' work [48] focuses on fusing
operations and reducing kernel-level startup overhead. However, it does not consider scheduling relationships among overlap.
all major MoE tasks or how to maximize
7 Conclusion
In this paper, we propose a scalable and mathematically proven pipeline scheduling framework called MoE to accelerate the training of MoE models. FoE addresses the unified scheduling across all major MoE-related tasks and enables the optimal coexistence of heterogeneous communication tasks (A2A and all-reduce communications) in MoE training. This substantially advances distributed solutions beyond a simple extension of traditional pipeline methods. We implement the FlowMoE framework on Pytorch. We conduct extensive experiments with 675 typical MoE layers and four real- world NLP models across two GPU clusters. Experimental results show that FlowMoE outperforms state-of-the- of-tte-art MoE training frameworks including ScheMoE, FSMOE, Tutel, and FasterMoE by 13%- 57% in training time, 10%-39% in energy consumption, and 7%-32% in memory usage. Furthermore, Moo can be combined with many orthogonal optimization works, promoting distributed pipeline optimization to new stage.
8 Acknowledgements
This work was supported in part by the National Key Research and Development Project under Grant 2022YFB2201604 and in part by Zhejiang Provincial Natural Science Foundation of China under Grant LZ22F010008.
10

---

References
[1] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language Models are Few-Shot Learners. In Advances Neural Information Processing Systems, volume 33, pages 1877-1901, 2020.
[2] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The Llama Herd of Models. arXiv preprint arXiv:2407.21783, 2024.
[3] Daya Guo, Dejian Yang, Haowei Zhang, Junxiao Song, Ruoyu Zhang, Runxin Xu, Qihao Zhu, Shirong Ma, Peiyi Wang, Xiao Bi, et al. DeepSeek-RI: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning. arXiv preprint arXiv:2501. 2025.
[4] Robert A Jacobs, Michael Jordan, Steven J Nowlan, and Geoffrey E Hinton. Adaptive Mixtures of Local Experts. Neural Computation, 3(1):79 87, 1991.
[5] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. arXiv preprint arXiv:2006. 16668, 2020.
[6] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoe Le, Geoffrey Hinton, and Jeff Dean. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Expee Layer. arXiv preprint arXiv: 701.06538, 2017.
[7] Zixuan Ma. Jiaao He, Jiezhong Qiu, Huanqi Cao, Yuanwei Wang, Zhenbo Sun, Liyan Zheng, Haojie Wang, Shizhi Tang, Tianyu Zheng, et al. BaGuaLu: Targeting Brain Scale Pretrained Models with over 37 Million Cores. In Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, pages 192 192-204, 2022.
[8] Carlos Riquelme, Joan Puigcerver, Basil Mustafa, Maxim Neumann, Rodolphe Jenatton, André Susano Pinto, Daniel Keysers, and Neil Houlsby. Scaling Vision with Sparse Mixture of Experts. In Advances in Neural Information Processing Systems, volume 34, pages 8583-8595, 2021.
[9] William Fedus, Barret Zoph, and Noam Shazeer. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. Journal of Machine Learning Research, 23(120):1-39 2022.
[10] Shaohuai Shi, Xinglin Pan, Qiang Wang, Chengjian Liu, Xiaozhe Ren, Zhongzhe Hu, Yu Yang. Bo Li, and Xiaowen Chu. ScheMoE: An Extensible Mfxtureeoff-xpper Distributed Training System with Tasks Scheduling. In Proceedings of the 19th European Conference on Computer Systems, pages 236-249, 2024.
[11] Jiaao He, Jidong Zhai, Tiago Antunes, Haojic Wang. Fuwen Luo, Shangfeng Shi, and Qin Li. FasterMoF: Modeling and Optimizing Training of Large-Scale Dynamic Pre-Trained Models. In Proceedings of the 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, pages 120- 134-134,70 2022.
[12] Changho Hwang, Wei Cui, Yifan Xiong, Ziyue Yang, Ze Liu, Han Hu, Zilong Wang, Rafael Salas, Jithin Jose, Prabhat Ram, et al. Tutel: Adaptive Mixture-of-Experts at Scale. In Proceedings of Machine Learning and Systems, volume 5, pages 269-287. 2023.
[13] Jehoshua Bruck, Ching-Tien Ho, Shlomo Kipnis, and Derrick Weathersby. Efficient Algorithms for All-to-All Communications in Multi-Port Message-Passing Systems. In Proceedings of the 6th Anmual ACM Symposium on Parallel Algorithms and Architectures, pages 298 309, 1994.
[14] Jeffrey Dean, Greg Corrado, Rajat Monga, Kai Chen, Matthicu Devin, Mark Mao, Marc' aurelio Ranzato, Andrew Senior, Paul Tucker, Ke Yang, et al. Large Scale Distributed Deep Networks. In Advances in Neural Information Processing Systems, volume 25, 2012.
[15] Yunqi Gao, Bing Hu, Mahdi Boloursaz Mashhadi, A-Long Jin, Pei Xiao, and Chunming Wu. US-Byte: An Efficient Communication Framework for Scheduling Unequal-Size Tensor Blocks in Distributed Deep Learning. IEEE Transactions on Parallel and Distributed Systems, 35(1):123-13)) 2023.
11

---

[16] Shaohuai Shi, Xiaowen Chu, and Bo Li. MG-WFBP: Efficient Data Communication for Distributed Synchronous SGD Algorithms. In IEEE INFOCOM 2019-IEEE Conference on Computer Communications, pages 172-180, 2019.
[17] Xianyan Jia, Shutao Song. Wei He, Yangzihao Wang. Haidong Rong, Feihu Zhou, Liqiang Xie, Zhenyu Guo, Yuanzhou Yang. Liwei Yu, et al. Highly Scalable Deep Learning Training System with Mixed-Precision. Training ImageNet in Four Minutes. arXiv preprint 1807 12205, 2018.
(18) Yang You, Jing Li, Sashank J. Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh. Large Batch Optimization for Deep Learning: Training BERT in 76 minutes. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020, 2020.
[19] Jiaao He, Jiezhong Qiu, Aohan Zeng, Zhilin Yang, Jidong Zhai, and Jie Tang. FastMoF: A Fast Mixture-of-Expert Training System. arXiv preprint arXiv: 103. 13262, 2021.
[20] Jiamin Li, Yimin Jiang. Yibo Zhu, Cong Wang, and Hong Xu. Accelerating Distributed MoE Training and Inference with Lina. 2023 USENIX Annual Technical Conference (USENIX ATC 23), pages 945-959, 2023.
[21] Shaohuai Shi, Xinglin Pan, Xiaowen Chu, and Bo Li. PipeMoF: Accelerating Mixture-of- Experts through Adaptive Pipelining. In IEEE INFOCOM 2023-IEEE Conference Computer Communications, pages 1-10, 2023.
[22] Shulai Zhang, Ningxin Zheng, Haibin Lin, Ziheng Jiang, Wenlei Bao, Chengquan Jiang, Qi Hou, Weihao Cui, Size Zheng, and al. Chang, Li- Wen. Comet: Fine-grained computation- communication overlapping for mixture-of-expert.. arXiv preprint arXiv:2502. /98/ 2025.
[23] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. PyTorch: An Imperative Style, PirfoPPerformance Deep Learning Library. Advances Neural Information Processing Systems, 32, 2019.
[24] Xinglin Pan, Wenxiang Lin, Lin Zhang, Shaohuai Shi, Zhenheng Tang, Rui Wang, Bo Li, and Xiaowen Chu. FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Expeee Models. In Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1, page 524-539, 2025.
[25] Yanping Huang, Youlong Cheng, Ankur Bapna, Orhan Firat, Dehao Chen, Mia Chen, Hy- oukJoong Lee, Jiquan Ngiam, Quoe V Le, Yonghui Wu, et al. GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism. Advances in Neural Information Processing Systems, 32, 2019.
[26] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R Devanur, Gregory R Ganger, Phillip B Gibbons, and Matei Zaharia. PreaDream: Generalized Pipeline Parallelism for DNN Training. In Proceedings of the 27th ACM Symposium on Operating Systems Principles, pages 1-15, 2019.
[27] Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, Yanping Huang, Yida Wang, Yuanzhong Xu, Danyang Zhuo, Eric P Xing, et al. Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning. In 16th USENIX Symposiun Operating Systems Design and Implementation (OSDI 22), pages 559-578, 2022.
[28] Mohammad Shoeybi, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, and Bryan Catanzaro. Megatron-LMM Training Multi Milli--illion Parameter Language Models Using Model Parallelism. arXiv preprint arXiv: 000p:::9990000000 2019.
[29] Mingshu Zhai, Jiaao He, Zixuan Ma, Zan Zong, Runqing Zhang, and Jidong Zhai. SmartMoE: Efficiently Training Sparsely-Activanl Models through Combining Offline and Online Paral- lelization. In 2023 USENIX Annual Technical Conference (USENIX ATC 23), pages 961-975, 2023.
12

---

[30] Hao Zhang, Zeyu Zheng, Shizhen Xu, Wei Dai, Qirong Ho, Xiaodan Liang, Zhiting Hu, Jinliang Wei, Pengtao Xie, and Eric Xing. Poseidon: An Efficient Communication Architecture for Distributed Deep I Learning on GPU Clusters. In 2017 USENIX Annual Technical Conference (USENIX ATC 17), pages 181-193, 2017.
[31] Yanghua Peng, Yibo Zhu, Yangrui Chen, Yixin Bao, Bairen Yi, Chang Lan, Chuan Wu, and Chuanxiong Guo. A Generic Communication Scheduler for Distributed DNN Training Acceleration. In Proceedings of the 27th ACM Symposium on Operating Systems Principles, pages 16-29, 2019.
[32] Chang Chen, Xiuhong Li, Qianchao Zhu, Jiangfei Duan, Peng Sun, Xingcheng Zhang, and Chao Yang. Centauri: Enabling efficient scheduling for communication-compuuaaaa model training via communication partitioning. In Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3, pages 178-191, 2024.
overlap in large
[33] Abhinav Jangda, Jun Huang, Guodong Liu, Amir Hossein Nodchi Sabet, Saeed Maleki, Youshan Miao, Madanlal Musuvathi, Todd Mytkowicz, and Olli Saarikivi. Breaking the computation and communication abstraction barrier in distributed machine learning workloads. In Proceedings of the 27th ACM International Conference Architectural Support for Programming Languages and Operating Systems, pages 402-416, 2022.
[34] Shibo Wang, Jinliang Wei, Amit Sabne, Andy Davis, Berkin Ibbyii, Blake Hechtman, Dehao Chen, Karthik Srinivasa Murthy, Marcello Maggioni, Qiao Zhang, et al. Overlap communication with dependent computation via decomposition in large deep learning models. In Proceedings of the 28th ACM International Conference Architectural Support for Programming Languages and Operating Systems, Volume 1, pages 93-106, 2022.
[35] Yixin Bao, Yanghua Peng, Yangrui Chen, and Chuan Wu. Preemptive All-reduce Scheduling for Expediting Distributed DNN Training. In IEEE INFOCOM 2020-IEEE Conference on Computer Communications, pages 626-635, 2020.
[36] Shaohuai Shi. Xiaowen Chu, and Bo Li. Exploiting Simultaneous Communications to Accelerate Data Parallel Distributed Deep Learning. In IEEE INFOCOM 2021-IEEE Conference on Computer Communicationss, pages 1-10, 2021.
[37] Yungi Gao, Bing Hu, Mahdi Boloursaz Mashhadi, Wei Wang, Rahim Tafazolli, and Merouane Debbah. A Dynamic Sliding Window Based Tensor Communication Scheduling Framework for Distributed Deep Learning. IEEE Transactions on Network Science and Engineering, 12(2): 1080-1095, 2024.
[38] Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He. DeepSpeed-Mon Mooe: Advancing Mixture-of- Experts Inference and Training to Power Next-Generation Al Scale. In International Conference 2022.
on Machine Learning, pages
[39] Stephen Merity, Caiming Xiong, James Bradbury, Joshua Socher, Richard Peterson, Stephan Meylan, and David Bourgin. Openwebtext dataset. 2019.
[40] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer Sentinel Mixture Models. arXiv preprint arXiv:1609.0788 07843, 2016.
[41] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language Models are Unsupervised Multitask Learners. OpenAl blog, 1(8):9, 2019.
[42] Jacob Devlin, Ming-Wos Wei Chang, Kenton Lec, and Kristina Toutanova. BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long and Short Papers), pages 1171-4186, 2019.
[43] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open Foundation and Tuned Chat Models. arXiv preprint arXiv:2307.09288, 2023.
13

---

[44] Aixin Liu, Bei Feng, Bin Wang, Bingxuan Wang, Bo Liu, Chenggang Zhao, Chenggi Dengr, Chong Ruan, Damai Dai, Daya Guo, al. DeepSeck- -2: A Strong, Economical, and Efficient Mixture-of-Hxperts Language Model. arXiv preprint arXiv:2405.04434 2024.
[45] NVIDIA Corporation. NVIDIA System Management Interface (nvidia-smi) User Guide, 2025.
URL httpttss////eeeeopoor www.:opprr.iiiiiii.am nvidia. com/ www.swsttteeeeeeeeeooeeeoeeeooet system waw.a--aaaaaaenteaeamaaeaa inntttiintefccc..
[46] Xinyi Liu, Yujie Wang, Fangcheng Fu, Xupeng Miao, Shenhan Zhu, Xiaonan Nie, and Bin Cui. Netmoe: Accelerating moe training through dynamic sample placement. In The Thirteenth International Conference on Learning Representations, 2025.
[47] Chenyu Jiang, Ye Tian, Zhen Jia, Shuai Zheng. Chuan Wu, and Yida Wang. Lancet: Accelerating mixuure-of-experts training via whole graph computation-commuiicatie Proceedings of Machine Learning and Systems, volume 6, pages 74 86, 2024.
overlapping. In
[48] Kishore Punniyamurthy, Khaled Hamidouche, and Bradford M Beckmann. Optimizing dis- tributed ml communication with fused computation-collecte operations. SC24: Interna- tional Conference for High Performance Computing, Networking, Storage and Analysis, pages 1-17, 2024.
[49] Omid Alipourfard, Hongqiang Harry Liu, Jianshu Chen, Shivaram Venkataraman, Minlan Yu, and Ming Zhang. ChemyPick: Adaptively Unearthing the Best Cloud Configurations for Big Data Analytics. In 14th USENIX Symposium on Networked Systems Design and Implementation, NSDI, pages 462,482,20 2017.
[50] Lin Zhang, Shaohuai Shi, Xiaowen Chu, Wei Wang. Bo Li, and Chengjian Liu. DeAR: Accelerating Distributed Deep Learning with Fine-Grained All-Reduce Pipelining. In IEEE 43rd International Conference on Distributed Computing Systems (ICDCS), pages 142-153, 2023.
[51] Richard I. Graham, Timothy S Woodall, and Jeffrey M Squyres. Open MPI: A Flexible High Performance MPI. Parallel Processing and Applied Mathematies: 6th International Conference, PPAM 2005, Poznarí, Poland, September 11-14, 2005, Revised Selected Papers 6, pages 228-239, 2006.
[52] Min Si, Yutaka Ishikawa, and Masamichi Tatagi. Direct MPI Library for Intel Xeon Phi Co- Processors. In 2013 IEEE International Symposium Parallel & Distributed Processing, Workshops and Phd Forum, pages 816-824, 2013.
14

---

NeurIPS Paper Checklist
1. Claims
Question:
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We describe our scope and contributions in the abstract and introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We analyze the limitations of our approach based on the experimental results in Section 5.2.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
15

---

Justification: We describe the theory assumptions and proofs in Section 3.3, Section 4.1, Appendix and Appendix C.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: clearly and fully present the selected models and benchmarks Section and also describe some parameter settings in Appendix D.1.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution primarily new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution a new model (c.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
16

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification:
We have provided the URL of our code in abstract.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips) cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
//nips. httppsscccpppccccceeeeeeeeecsece /tt.ggeeeoooooeuiiiiieeeee /Cc////odiiissseeeeeeea wwsssonoolicy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer:
[Yes]
Justification: All experimental settings, benchmark and training details can be found in Section 5, and Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification:
All our experimental results come from the average 1000 training iterations.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for
example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
17

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
Experiments compute resources are reported in Section 5.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https / http///eurrpp..eeee wwwpurrps.ccpuuliieeeeee
Answer: [Yes]
Justification:
We fully comply with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- cration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: The proposed FlowMoE training framework for LLM, thus, has no negative societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
18

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We adhere to the usage restrictions of all assets and correctly cite their sources.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi twwode.hco....c. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
19

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: We do not release assets now, and the code will be open-sourced after the paper is accepted.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
We do not involve any crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We do not involve any crowdsourcing or research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
20

---

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs are used in any part of the core methodology scientific content of this work.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to LLM policy (https: //tpurpee..p///e/ee cc/Conferencces///200 /LLM) for what should or should not be described.
21

---

Appendix
FlowMoE: A Scalable Pipeline Scheduling Framework for Distributed Mixture-of-Experts Training
Table A.I: Frequently used notations
Name
Description
P
Number of workers (or GPUs) in the cluster.
L
Number of transformer blocks in MoE model.
B
Number of samples per GPU (or mini-batch size) in one iteration.
Number tokens per sample.
M
Embedding size of token.
Total number of experts per MoE layer.
H
Hidden size of the feed-forward layer in experts. Top-k experts should be selected for each token. Capability factor that determines the maximum number of tokens assigned to each expert.
k
f
R
Pipelining degree. Partition size of the all-reduce tensor block.
Sp
AT,
The rth MHA layer computing subtask of the Ith transformer block (including the rth gating subtask)
E
The rth expert computing subtask of the /th transformer block.
D
The rth dispatch subtask (A2A communication) of the 1th transformer block. The rth combining subtask (A2A communication) of the /th transformer block.
c
AR
The all-reduce communication task of the /th transformer block during the backward propagation.
(-())/()))
The beginning execution timestamp of a task during the feed-forward comput- ing/backward propagation.
(-())/t6-)
The elapsed time of a task during the feed-forward computing/backward propagation.
A Main differences between FlowMoE and the key literature
Table A.2: Main differences between FlowMoE and the key literature pipeline scheduling for distributed MoE training, where BO represents the Bayesian optimization. vanillaE[[19] [asterMoE[11] Tutel[12] ScheMoE[10] FlowMoE
Scheduling framework
A2A pipelining
Expert computing pipelining
MHA + gating pipelining
X
X
All-reduce pipelining
Auto-tuning
x
X
X
V (BO)
Robustness of dynamic
Weak
Weak
Weak
Weak
Strong
hardware environment
Speedup (Homogeneous
1.0x
1.28x
1.31x
1.45x
1.82x
16-GPU cluster)
Speedup (Heterogeneous
1.0x
1.31x
1.36x
1.44x
1.74x
16-GPU cluster)
B Proof of Theorem 1
According to Eqs. and 5, since the task scheduling timeline each transformer block is the same and the all-reduce communication time of each transformer block is also the same, only need to
22

---

prove that T3 when inserting AR((+)) between any two A2A communication tasks of transformer block I will be less than or equal to Tz
We discuss the relationship between T3 and 7 under the following four tentative scheduling orders 1 I, S < R):
Scheduling order 1: m.(C To (AR ro())))
Scheduling order 2: (C - (AR S ( (1)
Scheduling order mh (D (+11)= ro(A(((t))) s (D.)). Scheduling order 4: Th (D())) - (AA S (c
If Scheduling order 1 holds, according to Eq. 6b, have
Th (=(" max(na (C,") to() to (6(R((+)), To (o (r()))))))
(A.1)
And if AR(+1) is performed at the end of backward propagation, according to Eq. 6b, we have
Th (E,")) = max (T (C.)) + ta(C.)) To (E + to (((t)))))
(A.2)
Since Ti' (C.")) = (C. according to Eqs. A.1 and A.2, we can derive (o < T (E.)) + 4 (AR(+1))
(A.3)
Meanwhile, according Eqs. 4, and 6a-6e, we have
() + (ARR)) - To (F(D) + (o(AR(()))) = T (AR.)) + (((RR))) - - (A.4)
In other words, between task AR and task Ex" the timeline based on centralized scheduling of all-reduce communication tasks has one more (R((+)) than the timeline based on Scheduling order Then, adding the left and right terms of Eqs. A.3 and A.4, respectively, we have
To (AR())) + ()) S Th (A + to (AR()).
(A.5)
Moreover, since TO ( (c Ji, (C())) according Eq. 6, we have T. s TV. An easy-to-underrsand demonstration is presented in Fig. A.1.
AR 11-
time
time
(b)
E
time
time
(d)
8
et
Ct
time
time
J
(f)
Figure A.1: Demonstration of one example where (a). (c) and (e) represent the three timelines when using centralized scheduling of all-reduce communication tasks (the dotted lines indicate that the tasks before and after do not have direct dependency), and (b), (d) and (f) represent the three timelines when using Scheduling order 1. Comparing (a) with (b), (c) with (d) and (e) with (f), can be observed that T3 < Tg.
Similarly, we can obtain the same conclusion when using the tentative scheduling orders 2.3 and 4. Therefore, T, will be less than or equal to Th when inserting all-reduce communication task of one transformer block between any two A2A communication tasks, and the proof is complete.
23

---

C Proof of Theorem 2
When the scheduling order satisfies Eqs. 4 and 5 during backward propagation, the idle time of the communication resources (i.e., the gaps between the A2A communication tasks) is fixed if not considering the all-reduce communication tasks. Obviously, the approach to minimize the time of one iteration is to maximize the occupation of the idle time of the communication resources by the communication tasks of the all-reduce tensor chunks. Furthermore, if the communication of the all-reduce tensor chunks does not introduce extra startup overhead, the total time of the all-reduce communication tasks will remain unchanged for different Sp. When Sp is larger, A2A communication tasks may not be executed in time due to uncompleted communication tasks of the all-reduce tensor chunks, which will cause the timeline of the computing tasks to shift back according to Eqs. 6b and 6d and part of the communication tasks of the all-reduce tensor chunks do not occupy the free time of the communication resource (see Figs. A.1b and Af. On the contrary, when Sp + O, the communication tasks of the all-reduce tensor chunks will not affect the timeline of the A2A communication tasks and thus will not change the timeline of the computing tasks due to the higher priority of the A2A communication tasks. Meanwhile, will maximize the free time of the communication resources. Thereby, the proof is complete.
D Design and Performance Evaluation of BO
D.1 Parameter Settings and Advantage Analysis
We choose BO obtain a near-optimal Sp benefiting from the following three points:
First, BO is not limited by the expression of the objective function (we use F(Sp) as the objective function) and depends only on the sampling values obtained (i.e., F(S)), (($T)) We use Gaussian process regression with the Matern kernel to predict the value of the objective function as it is commonly used good surrogate model for BO [49]. A 95% confidence interval is associated with F(Sp), which is considered the most likely region for the value of F(S,).
Second, BO typically requires only limited number of trials to find high-quality solutions, resulting in low search overhead. To minimize the number of trials, it selects the next configuration Sp by maximizing an acquisition function [49]. In FlowMoE, we adopt the Expected Improvement (EI) acquisition function to choose the Sp that maximize training speedup improvement compared to the current best result.
Third, BO avoids the search process from falling into local optimum by tuning the hyperparameter Expected Improvement (EI). Specifically, a smaller El favors exploitation, i.e., collecting more points near the peak, while larger EI tends to exploration, i.e., collecting more scattered points in the range [49]. In FlowMoE, we choose El=0.1 I to prefer Sp exploration, which is a commonly adopted value [31, 50]. Meanwhile, the search space Sp in BO can be set as (OMB, the maximum value of the tensor to be communicated in each transformer block] and the required one initial sample value is randomly generated, which can effectively cover the optimal Sp for different hardware clusters and MoE models, e.g., from OMB- in Fig. 4.
D.2 Importance Analysis of BO
BO auto-tuning an indispensable part of FlowMoE and is absolutely necessary. The performance of FlowMoE depends on the introduced BO process. Specifically, an oversized all-reduce tensor chunk will affect the prioritization of A2A communication tasks and may prevent them from being started in time. In contrast, an undersized all-reduce tensor chunk will result in excessive startup overhead. Therefore, neither an oversized nor undersized all-reduce tensor chunk can achieve the shortest per-iteration time, and there must exist a unique optimal solution that maximizes training speed. The BO optimizer balances these two competing effects by sampling and learning from actual iterations, efficiently finding an all-reduce tensor chunk size that maximizes overlap without incurring excessive overhead. This tuning is crucial because the optimal trade-off point varies across models, GPU interconnect topologies, and MoE configurations.
24

---

D.3
Performance Evaluation of BO
Table A.3: Comparison of average per-iteration time milliseconds when using different approaches
tune the partitioning size Sp-
Time (ms)
Model
B0
Grid Search
Random Number Generation
GPT2-Tiny-Mal
95.6
101.3
109.3
BERT-Large-Mal
351.9
373.80
388.96
LLaMA2-MoE
1124.0
1208.23
1250.09
DeepSeck-v2-S 3205.3
3498.8
3902.75
We compare the performance of BO with the other methods including grid search and random number generation on tuning Sp. In experiments, we set the other two methods to use the same search space as BO. For grid search, the search space divided into equal parts construct discrete sampling points. Similarly, the grid search utilized the first 80 iterations of the training process to determine Sp. Specifically, each iteration time corresponding to each sample point is obtained by averaging 10 iterations, and the sample point corresponding to the smallest per-iteration time is used as the value Sp for subsequent training. For random number generation, we randomly selected number from the search space the value of Sp in each iteration. Table A.3 shows the comparison of average iteration times when using different methods of tuning Sp in FlowMoE on Cluster 1 with 16 GPUs. It can be observed that using BO to tune Sp obtained the shortest per-iteration time among the three methods. Although grid search obtains better performance than random number generation, its limited number of sampling points makes it difficult cover the optimal solution especially when the search space large, and over-increasing the number of sampling points brings higher search overhead and longer search process. On the contrary, BO can obtain near-optimal Sp with very little overhead, and therefore we choose it.
Table A.4: Comparison of average per-iteration time in milliseconds when using FlowMoE with BO auto-tuning different fixed partition sizes.
Time (ms)
Model
BO
S, 0.5MB
S,=IMB
5,=2MB
5,=4MB
S,=8MB
GPT2-Tiny-MoE
95.6
130.1
115.9
104.7
109.7
122.3
BERT-Large-Moie
351.9
388.9
378.6
362.2
386.3
395.1
I.AMM2-MoEE
1124.0
1213.8
1167.9
1185.9
1211.8
1240.5
DecpSeck-v2-S 3205.3
4438.5
3948.9
3654.7
3493.9
3740.2
BO Auto-tuning v.s. Different Fixed Partition Sizes. In addition, we count the average per-iteration time when using FlowMoE with BO auto-tus auto-tuning different fixed partition sizes for training four MoE models (see Table 2 for their detailed configurations) in Cluster with GPUs to validate the sensitivity of FlowMoE Sy. As shown in Table A.4, different partition sizes greatly affect the training efficiency and BO auto-tuning is essential to maximize the performance of FlowMoF.
Table A.5: Comparison of average per-iteration time when training BERT- BERT-Large-MoE with differ- BO parameter configurations on Cluster where the surrogate model uses Gaussian Process Regression (GPR) with different kernel functions.
BO Hyperparameter
Time (ms)
Acquisition Function
Surrogate Model
Expected Improvement (EI=0.1)
GPR + Matern
351.9
Expected Improvement (EI=0.05)
GPR Matern
358.9
Expected Improvement (EI=0.2)
GPR + Matern
354.2
Probability of Improvement
GPR Matern
355.1
Lower Confidence Bound
GPR Matern
355,4
Expected Improvement (EI=0.1)
GPR + RBF
357.2
Expected Improvement (EI=0.1)
GPR + Rational Quadratic
360,2
25

---

Sensitivity of BO hyperparameters. In FlowMoE, the objective function for per-iteration time regarding the all-reduce tensor chunk size is single-peaked, smooth objective function with fixed search space. For this objectíve function, the process of using BO to find the optimal solution is insensitive to the two hyperparameters (acquisition function and surrogate model) of BO, and BO always converges and approaches the optimal solution. Table A.5 shows the comparison results of average per-iteration time when training BERT-Large-MoL with different BO parameter configura- tions on Cluster 1. The results indicate that although the performance of FlowMoE is sensitive to the BO process, BO hyperparameters have a minor impact on it, and different BO configurations lead to similar iteration time.
Table A.6: Percentage of the computational overhead of BO to the training time of first 1000 iterations. Model GPI2-Tiny-MoE BERI-Large-Mof LLaMA2-MoE DeepSeek-v2-S
Overhead 3.22%
1.38%
0.43%
0.16%
Computational Overhead of BO. We also measure the percentage of the computational overhead of BO to the training time of first 1000 iterations when training four MoE models. The experimental results are illustrated in Table A.6. It can be observed that this overhead is negligible compared to the gains brought by BO (see Table 3), and it will be smaller when training the model convergence, which involves tens of thousands of iterations. Thus, the auto-tuning process of BO is lightweight and practical.
E Algorithm Description and Scalability Analysis
Algorithm 1 describes the pipeline of R-degree computing and communication tasks during the training process. Line 1-4 initializes the necessary parameters and queues. Line 6-12 performs the feed-forward computing of each iteration according to Eqs. 2 and 3. Meanwhile Line 13-21 performs the backward propagation of each iteration according to Eqs. 4 and 5. Algorithm 2 shows the management in the communication pool. The PARTITION procedure is responsible for splitting the tensor of the MHA layer and gating function during the backward propagation process and placing it into the queue of the all-reduce communication tasks. The COMMPOOLMANAGER procedure performs the two types of communication tasks according to the defined priorities, where Line 9- 11 prioritizes the execution of the A2A communication tasks and queues the data obtained after the communication, while Line 12-13 executes the communication tasks of the all-reduce tensor chunks A2A communication tasks.
when there
To analyze the scalability of FlowMoE's scheduling algorithm, we examine the computational complexity of each iteration of the MoE model in distributed training. We assume that each GPU processes the same number of tokens and do not consider the non-overlapped time of A2A communi- cation and all-reduce communication in each iteration, then the time complexity of one iteration is O (L x M + B + BNME + 2BNMH ), where O M + B² B²N'M M) the linear mapping of the Q.K.V and linear transformation matrix in the MHA layer and the attention score computation, { and O (2.N H) denotes the expert computing on each GPU. Then, when the complexity of the MoE model scales up. the increase in the complexity of FlowMoE'' scheduling algorithm is significantly smaller than the increase in the complexity of one iteration time. Meanwhile, when the number of GPUs P in the cluster increases, the complexity of the scheduling algorithm of FlowMoE. not affected, but the complexity of one iteration time may increase due to the longer time of the communication tasks. In summary, the ratio of the time overhead of FlowMoE's scheduling algorithm one iteration time will decrease when the model complexity and the number of GPUs increase. Therefore, FlowMoE's scheduling algorithm has good scalability.
denotes
denotes the gate function computing with the structure of M E,
F
System Description
From the closest to user level to the lowest level, distributed training frameworks and communication stacks usually include: User Script level, PyTorch frontend with high-level APIs, PyTorch Engine, message-level communication library. To ensure the generality of FlowMoF, we can not modify user scripts and framework engines heavily. However, due to the diversity of MPI (c.g.. OpenMPI [51],
26

---

Intel MPI [52]) and network interfaces (e.g.. Ethernet, Infiniband) in the communication library level, it is almost impossible to make one piece of code work in all communication libraries. Therefore, FlowMoF is deployed at the PyTorch API level, where its four modules is elaborated as follow:
Task Breakdown Manager is responsible for splitting the dataset of each batch size according to the degree R. and requesting different communicattonn/oompuiii subtasks based on the tensor transferred between different computing and communication tasks. Task Breakdown Manager has been improved based on Tutel.
BO autotuner is responsible for searching the near-optimal partition size and guiding the partitioning of the all-reduce tensor.
Communication Task Pool maintains queues of A2A communication tasks and all-reduce chunk communication tasks. On the one hand, it receives A2A communication tasks requested from the Task Breakdown Manager and queues them according to the request order. On the other hand, it splits the all-reduce tensor generated in the backward propagation and queues the all-reduce chunks.
Pipeline Scheduling Manager submits the requested multiple computing tasks and commu- nication tasks in the queue to the PyTorch Engine and the communication library according the defined scheduling order and priority.
In addition, the backward propagation of each transformer block is not intuitively visible in Py Torch, and the gradients of the model parameters can only be accessed after the backward propagation is fully completed. To obtain the gradient tensors of the MHA layer and gating function for each transformer block in time during backward propagation, we use register_ull conster_flllbackwweee backward _ook to access the gradients and split the gradient tensors into queues based on the partition size Sp. Meanwhile, we achieve the overlapping of computing and communication tasks by using multi-threadien selli-threading. Specifically, the main thread includes Task Breakdown Manager, BO autotuner and Pipeline Scheduling Manager and is responsible for scheduling all tasks. We also create one sub-thread in Communication Task Pool schedule the communication tasks in the two queues based on the defined priorities, and threading Lock is used to maintain security for all threads.
G Stress Tests in two Scaled-up MoE Models
Table A.7: Comparison of average per-iteration time when training two scaled-up MoE models. S1, S2, and S3 are the speedups of FlowMoE over ScheMoE, Tutel and vanillaEP, respectively.
# of
Time (ms)
Model
S₂
S₁
GPUs
vanillEEP Tutel
ScheMoE FlowMoE
LLaMA2-MoE-L
2405.1
1927.0
1806.1
1493.8
1.61 xl..xx1xxxxx x1.21X
4
DeepSeck-V22 V2-M
535.3
468.4
432.2
352.2
1.52xI.33x1.223
L.aMA2-MoF-L
2989.1
2493.9
2297.9
1833.8
1.6x1126x
8
DeepSeck-V2-M.
944.6
773.4
723.6
552.4
40x.31
LLaMA2-MoE-L
OOM
OOM
OOM
OOM
/
1
16
DeepSeek-V24 V2-M
1254.6
956.9
893.4
708.8
x.77 x1.226
To stress-test FlowMoE's performance, we measured the training performance of different training frameworks on two scaled-up MoE models [LLaMA2-MoE-4) and DeepSeek-V2-M, V2-MM. both which close the memory upper limit of Cluster 1. Table A.7 shows the average per-iteration time using different frameworks on these two MoE models (FasterMoE is OOM in any cases). The results indicate that FlowMoE still achieved the best training performance among all baselines larger models.
H Convergence Analysis and Experiments
We denote the samples processed in each pipeline as one microbatch, and the number of microhatches is equal to the pipelining degree R.
27

---

Specifically, during the backpropagation of the lth transformer block in one iteration, the gradients from all AT; and E, (1 S I R) after backpropagation are accumulated and summed. Once the gradient from ATŞ is also accumulated, the All-Reduce chunk communication task for the MHA and gate function in the /th transformer block is started. Similarly, when the gradient of E is accumulated, the expert parameters are updated. This effectively prevents the parameters of MHA, the gate function, and the expert from being updated early, thereby avoiding gradient staleness. Additionally, to ensure that the accumulated gradients are equívalent with and without pipelining, we scale the loss calculated for each microbatch by R, i.e., loss!") (1
R), where loss is the loss calculated using the samples from the rth microbatch, and the detailed theoretical derivation is as follows:
R
The loss and gradient when performing backpropagation using the entire mini-batch are expressed as follows:
Σ loss B
(x,, V = V
Ex w1)
(A.6)
B
which is currently used by all mainstream MoE training frameworks [12, 38, 28].
When dividing mini-batee into R microbatches, the number of samples in each microbatch is = R The loss for each microbatch is:
loss(r) (r)
Σ ((FrssBrs)
(A.7)
Then, the scaled loss for each microbatch is:
loss (r)
loss(r) R 2 (Trri Dr,s)
1
_
(arsis
(A.8)
B
The cumulative loss of all microbatches is:
R
loss Σ
((xr, Ur.))
Σ
( Br.))
1
Σ
((x
(A.9)
According to Eqs. A.6 and A.8, obtain:
loss (r)
1 B
V
Σ
V
E(1,, 81)
VL
(A.10)
This analysis shows that the gradient update strategy is numerically equivalent with and without pipelining. The only difference between them is the optimization of task scheduling order during execution, which does not compromise model convergence during training.
10.5
FasterMoE
5.6
FasterMoE
4.8
ScheMoE
24.0
ScheMoE
6.0
FlowMoE
932
FlowMoE
2.4
4.5
1.6
10
20
30
40
50
1
7
Time(h)
Time(h)
(a) GPT2-Timy-MoE
(b) BERT-Large-MoF
Figure A.2: Loss with time for training GPT2-Tiny-Moo and BERT-Large-MoE.
Furthermore, we experimentally verify the convergence of FlowMoE by training GPT2-Tiny-Moft and LarRT-Largee-M on Cluster with 16 GPUs As shown in Fig. A.2, FlowMoF. reaches the same loss as the baselines, and it takes much less time due to reduced per-iteration time.
28

---

I Performance Lower Bound Analysis
We analyze the performance lower bound of FlowMoE from the following three cases:
(I) The communication task time is much longer than the computing task time. When the communication time very long, the time of the All-to-All (A2A) communication task alone completely covers the computing task time. In this case, the All-Reduce (AR) chunk based priority scheduling mechanism will fail because the AR chunk cannot be inserted into the A2A communication task. The performance of FlowMoE will be the same ScheMoE, Tutel, and FasterMoE, but better than vanillaFP due to the hidden computing task time.
(2) The computing task time is much longer than the communication task time. When the computing time very long, all the communication task time can be covered by the computing task time. In this case, FlowMoE will outperform ScheMoE, Tutel, and FasterMoB due to the hidden AR task time, and better than vanillaEP because it further hides the A2A task time.
(3) The communication task time is comparable to the computing task time. In this case, Flow- MoE outperforms all baselines because maximizes the overlap of multi-typ typliittypeeeeee tasks.
In summary, in all cases, the performance FlowMoE is greater than or equal to that of ScheMoE, Tutel, and FasterMoB, and is always better than that of vanillaEP.
J GPU SM Utilization
Table A.8: Average GPU SM utilization with different microbatch sizes.
Name
Model
R
Average GPU SM Utilization
FlowMoE
GPT2-Tiny-Mot:
72.63%
FlowMoE
GPT2-Tiny-Moo
48.43%
vanillaFP
GPT2-Tiny-MoE:
87.09%
FlowMoE
BERT-Large-Mott
87.84%
Flo
BERT-Largc-MoE
78. 16%
vanillaEP
BERT-Large-Moz
88.90%
FlowMoE
LLaMA2-MoF
89.16%
FlowMoE
LLaMA2-MoE
88. 88.19%
vanillaEP
LLaMA2-MoE
89.49%
FlowMoE
DeepSeck-V22 V2-S
89.27%
FlowMoE
DeepSeek-V24. V2-S
88.85%
vanillaEP
DeepSeek V2-S
90.77%
First, we use the CUPTI tool to measure GPU SM utilization on Cluster 1 with 16 GPUs under different microbatch sizes. Specifically, we adjust the microbatch size by changing the pipelining degree R (a larger R results in smaller microbatch). VanillaEP represents the original mini-batch without pipelining. The results are shown in Table A.8. We observe that smaller microbatch sizes may result in lower GPU SM utilization (e.g.. when training GPT2- GPT2-Tiny-MoE: with f using FlowMoB), but in most cases, GPU SM utilization is nearly identical that without pipelining (e.g., when training BERT-L BERT-Large-Me MooE L.LMMA2MMo. and DeepSeek-V22 V2-S using FlowMoF)). This is because the actual tensors involved in GPU computation remain sufficiently large for larger MoE models, and GPU SM resources still utilized efficiently. In other words, the microbatches introduced by FlowMoEE's pipelining do not waste GPU SM resources for large MoE models.
Second, we measure GPU SM utilization with different batch sizes when training different MoE models using FlowMoE on Cluster with 16 GPUs. As illustrated in Table A.9, smaller batch sizes are more likely to lead to lower GPU SM utilization. e.g., when training EPT2-Tiny-Mor or BERT-Large-Mon with a batch size of 2. Moreover, GPU SM utilization remained nearly unchanged when training LLAMA2-MoE or DeepSeek-V22 V2-S. This is because these two MoE models have larger number of parameters, which keeps the tensor dimensions involved in GPU computation sufficiently large, cnsuring that SM resources still efficiently utilized even with lower batch size.
29

---

Table A.9: Average GPU SM utilization with different batch sizes when using FlowMoE.
Model
Batch Size
Average GPU SM Utilization
GPT2-Tiny-MoE
4
72.63%
GPT2-Tiny-Mob
2
36.62%
BERT-Large-MoE
4
87.84%
BERT-Large-MoE
2
61.48%
LLaMA2-MoE
4
89. 99.16%
I.LMM2-MoE
2
88.45%
DeepSeck- -
4
89.27%
DeepSeek V2-S
2
89.06%
Table A.10: Parameter Configurations for BERT-Large-Moofs
Params
Params
Configurations
MoE Model
Dataset
(MHA+Gating) (Experts)
B
M H
R/P k
BERT-Large-MoBww
25,2M
3325.9M
wikitext- 103
24
4
512
512
10248
Table A.11: The maximum and minimum GPU SM utilization for large number of experts with different numbers activated experts when using FlowMoE Cluster with 16 GPUs..
Model
f
Maximum GPU SM Utilization Minimum GPU SM Utilization
BERT-Large-Moud MoE-w
1.0
89.20%
87.81%
HERT--Lage-MoE..
4.0
89.72%
50.65%
BERT-Large-MoE.
8.0
90.30%
31.60%
BFRT--Large-MoFw
16.0
90.68%
19.41%
Third, we also evaluate the impact of the number of activated experts GPU utilization when using FlowMoE on Cluster with 16 GPUs for a large number of experts. Specifically, we construct BERT-Large-MMo-w by increasing the number of experts per GPU from 2 to 8 and adjusting the number of activated experts by modifying the capacity factor f. where a larger f indicates more uneven token routing by the gate function and fewer activated experts since popular experts. The detailed configurations of BERT-Large-Me --w are shown in Table A. We the CUPTI tool and report the maximum and minimum GPU SM utilization in Table A. The results demonstrate that, with a large number of experts, the fewer the number of activated experts, the more unbalanced the computation load on the GPU and the greater the difference SM utilization between GPUs.
tokens are routed to
K Robustness Analysis
We discuss the robustness of FlowMoE to heterogeneous clusters, dynamic hardware environments, and node dropouts.
K.1 Robustness for Heterogeneous Clusters
Table A.12: Comparison average iteration time in milliseconds when the GPUs have different computing power. S1, S2, S3 and S4 are the speedups of FlowMoE over vanillaEP, FasterMoE, Tutel and ScheMoE, respectively.
Time (ms)
Model
S₁
S3
S2
S
vanillaEP FasterMoE
Tutel
ScheMoE
FlowMoE
GPT2-Tiny-Mof
235.8
201.6
189.1
178.2
153.3
1.54x1.32x
1.16x
BERT-Large-Mod
657.7
608.7
590.8
500.6
449.2
1.3x 1.1IX
LLaMA2-MoE
2439.1
2152.2
1849.2
1707.4
1468.3
1.66x
1.66 1.15x
DeepSeek-V2-S
7233.7
5495.1
5323.0
4958.3
4142.4
1.3
30

---

We conduct detailed analysis on how FlowMoE adapts to heterogeneous GPU clusters. (1) For clusters with mixed GPU types, since all-reduce and A2A tasks can only begin once the slowest GPU completes its corresponding computing task (other GPUs will be idle in waiting time after com- pleting their computing tasks), the task timeline for each iteration of FlowMoF and other baselines determined by the slowest GPU. In this case, FlowMoE can maximize the computin/ccommuncattio overlap of the slowest GPU, and its performance still outperforms the baselines. To verify this conclusion, construct a heterogeneous cluster with 16 GPUs Cluster I with different GPU types, We add delay to each iteration on GPUs in one node of Cluster to simulate the computing power reduction. Specifically, we cess the forward computing and backward propagation of each layer. Meanwhile, we synchronize the tensor of each layer with torch. cadh..uaaaaaaa synchhhycceeeeee after it completes the forward computing or gradient computing. Then, we obatin the forward computing or backward propagation time of each layer according to the interval between two tensor synchronizations and inject equivalent delays during both passes, which effectively simulates halving the computing power of the GPUs. Table A.12 demonstrates that FlowMoE still obtains the shortest end-to-end training time when the GPUs have different computing power. (2) For clusters with bandwidth asymmetry, since AR and A2A tasks are collective communications, all GPUs start and finish communication tasks simultaneously regardless of whether their bandwidths are the same. Therefore, the task scheduling timeline is the same for all GPUs In this case, FlowMoE still achieves better performance than the baseline due to its efficient pipeline solution.
forward hook and register _hook in PyTorch to ac-
K.2 Robustness for Dynamic Hardware Environments
In real large-scale training clusters, the hardware environments including network bandwidth and GPU computing power may change dynamically. To address it, we design the re-Bayesian tuning mechanism for FlowMoE and BO will be automatically re-executed when the hardware enviromment changes. Specifically, we define a re-execution threshold ò. We denote the iteration time correspond- ing to the near-optimal Sp predicted by the last BO as F(.hhest. If the change of the current iteration time T' compared to F(Sheat) is more than i.e..
T F(Sbeat) F(Speat)
3 6,
(A.11)
the BO will be re-executed to find new Sp.
K.3 Robustness for Node Dropouts
In FlowMoE, to improve the fault tolerance of the system to handle node dropouts, we adopt the following approach:
First, a replica of each expert parameter is stored separately on two different nodes to ensure that when a node fails, its corresponding expert parameter can still continue to be served by the backup replica. To ensure the parameter consistency between the replicas, we set synchronize the updating of expert parameters every fixed number of iteration steps (e.g., every 1000 steps) during the training process.
Second, FlowMoE periodically calls torch.distributed. ) during the training process with a reasonable timeout (e.g., 5 minutes). When synchronization operation times out and throws exception, we determine that a node has failed, and obtain the rank of the failed node from the exception information. Then, FlowMoE will execute the following recovery process:
(1) Based on the node replica information, the gating function routing table is updated on all surviving nodes, and the requests originally assigned to the faulty node are remapped to its corresponding backup node;
(2) The remaining surviving nodes adapted to the new of nodes and subsequently trained by destroying the current communication group (dist.destroy_prog Sotostroyyproes___s_ (roop(() and reinitializing a new distributed communication group (dist.nitt Bocosiini_pproes______ pwtpcosssgrouppii)
Therefore, FlowMoE is able to effectively handle the risks from node dropouts in large-scale dis- tributed training, and improves the stability and reliability of the overall training.
31


Reviews:
[{"review_id": "nIv0QMaM3w", "reviewer": "Reviewer_hxBa", "summary": "The paper proposes FlowMoE, a scalable pipeline scheduling framework to improve the training efficiency of distributed Mixture-of-Experts (MoE) models. Unlike prior MoE frameworks that focus primarily on pipelining within the MoE layer, FlowMoE constructs a unified pipeline that includes Multi-Head Attention (MHA) computing, gating, expert computing, and All-to-All (A2A) communication, and additionally pipelines all-reduce communication using a tensor chunk-based priority scheduling mechanism. It leverages Bayesian Optimization (BO) to automatically tune the all-reduce partition size for optimal overlap, aiming to reduce iteration time, energy consumption, and memory usage during training. Extensive experiments on two GPU clusters with real-world MoE models (e.g., GPT2-Tiny-MoE, LLaMA2-MoE, BERT-Large-MoE, DeepSeek-V2-S) and 675 customized MoE layers demonstrate 14%-57% speedup, 10%-39% energy reduction, and 7%-32% memory savings over existing state-of-the-art MoE frameworks such as ScheMoE, Tutel, and FasterMoE.", "questions": "* Have you observed any cases where FlowMoE's aggressive overlap leads to slower per-iteration time on specific settings?\n* Could you elaborate on how FlowMoE adapts to heterogeneous GPU clusters (e.g., mixed GPU types or bandwidth asymmetry) and whether BO retuning is required for every hardware change? How sensitive is FlowMoE’s performance to the BO process?", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "# Strengths:\n\n* The proposed FlowMoE addresses a practical bottleneck in training large-scale MoE models, which is highly relevant as LLM scaling continues.\n* The paper carefully formalizes pipeline scheduling and priority mechanisms, and integrates BO for adaptive optimization, reducing manual tuning.\n* The experiments are comprehensive, with detailed comparisons under diverse settings.\n\n# Weaknesses:\n\n* While the paper claims novelty in pipelining MHA computing, gating, expert computing, A2A, and all-reduce communication together, it is unclear how fundamentally novel this scheduling is compared to prior works [1-3].\n* While BO tuning of the all-reduce chunk size improves performance, the paper does not provide insights into why the BO-tuned configurations lead to specific speedups compared to baselines or how sensitive the system is to these hyperparameters beyond empirical gains.\n\n[1] Chen, Chang, et al. \"Centauri: Enabling efficient scheduling for communication-computation overlap in large model training via communication partitioning.\" Proceedings of the 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 3. 2024.\n\n[2] Jangda, Abhinav, et al. \"Breaking the computation and communication abstraction barrier in distributed machine learning workloads.\" Proceedings of the 27th ACM International Conference on Architectural Support for Programming Languages and Operating Systems. 2022.\n\n[3] Wang, Shibo, et al. \"Overlap communication with dependent computation via decomposition in large deep learning models.\" Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1. 2022.", "was_revised": true, "final_justification": "The authors have address my concerns and I would like to maintain my score.", "rebuttal": "We thank Reviewer hxBa for the thoughtful feedback and constructive comments. We have responded and addressed every weakness and question in detail. Meanwhile, for weaknesses and questions that have been addressed with additional explanations or experiments in the rebuttal, we will incorporate the corresponding discussions or results into the revised version of the paper. Please see the one-by-one responses below.\n\n>### **[Weakness 1]**\n\n**[Answer 1]** Prior works [1–3] focus on non-MoE models and primarily consider pipelining between computing and All-Reduce (AR) communication. In contrast, FlowMoE is specifically designed for MoE models, where All-to-All (A2A) communication is introduced due to the sparsely activated expert layers. Compared to works [1-3], FlowMoE has two fundamental novelties: \n(1) the co-existence and scheduling of both AR and A2A communication tasks, and\n(2) the pipelining of computing tasks (MHA, gating and expert) with two heterogeneous communication tasks (AR and A2A). \nIn particular, to the best of our knowledge, FlowMoE is the first work to pipeline these five tasks together.\n\nIn addition, as shown in Fig. 2b, in the backpropagation of the MoE model, the expert computing of the $l$th transformer block depends on the combining A2A of the $l$th transformer block, the dispatch A2A of the $l$th transformer block depends on the expert computing of the $l$th transformer block, and the MHA computing of the $(l+1)$th transformer block depends on the dispatch A2A of the $l$th transformer block. However, the AR communication of the $l$th transformer block only depends on the $l$th transformer block MHA computing and gating. Therefore, the A2A communication is not equivalent to the AR communication. Works [1-3] do not consider these complex dependencies and cannot be directly used to solve the pipeline problem of distributed MoE training. The scheduling mechanism of FlowMoE is essentially an efficient solution for scheduling multi-type tasks in distributed MoE training.\n\n>### **[Weakness 2]**\n\n**[Answer 2]** BO auto-tuning is an indispensable part of FlowMoE and is absolutely necessary. While FlowMoE's efficient performance depends on the introduced BO process (Table 10 in Appendix D.2 demonstrates this conclusion), but it is less sensitive to the hyperparameters used in BO, which are of lower importance and can be manually adjusted. \n\nFirstly, regarding your point on providing insights, we clarify that: An oversized All-Reduce (AR) chunk size will affect the prioritization of All-to-All (A2A) tasks and may prevent them from being started in time. In contrast, an undersized AR chunk size will result in excessive startup overhead. Therefore, neither an oversized nor undersized AR chunk size can achieve the shortest per-iteration time, and there must exist a unique optimal solution that maximizes training speed. The BO optimizer balances these two competing effects by sampling and learning from actual iteration times, efficiently finding an AR chunk size​ that maximizes overlap without incurring excessive overhead. This tuning is crucial because the optimal trade-off point varies across models, GPU interconnect topologies, and MoE configurations. We elaborate on these in Appendix D and Theorem 2, where we show theoretically (and empirically in Table 9/10 and Fig. 4) that BO-tuned configurations consistently reduce iteration time compared to fixed or grid-based alternatives. We will clarify this explanation further in the main text if accepted.\n\nThen, in FlowMoE, the objective function for per-iteration time regarding the AR chunk size is a single-peaked, smooth objective function with a fixed search space. For this objective function, the process of using BO to find the optimal solution is insensitive to the two hyperparameters (acquisition function and surrogate model) of BO, and BO always converges and approaches the optimal solution. The following table shows the comparison results of average per-iteration time when training BERT-Large-MoE with different BO parameter configurations on Cluster 1. The results indicate that BO hyperparameters have a minor impact on FlowMoE performance, and different BO configurations lead to similar iteration time.\n\n|Acquisition Function|Surrogate Model|Average Per-iteration Time(ms)|\n|-|-|-|\n|EI (0.1)|GPR + Matern|351.9|\n|EI (0.05)|GPR + Matern|358.9|\n|EI (0.2)|GPR + Matern|354.2|\n|Probability of Improvement|GPR + Matern|355.1|\n|Lower Confidence Bound|GPR + Matern|355.4|\n|EI (0.1)|GPR + RBF|357.2|\n|EI (0.1)|GPR + Rational Quadratic|360.2|\n\n\\*EI: Expected Improvement  \n\\*Surrogate Model use Gaussian Process Regression (GPR) with different kernel functions.\n\n>### **[Question 1]**\n\n**[Answer 3]** No, FlowMoE will not obtain slower per-iteration time in any case. We analyzed three cases:\n\n(1) **The communication task time is much longer than the computing task time.** When the communication time is very long, the time of the All-to-All (A2A) communication task alone completely covers the computing task time. In this case, the All-Reduce (AR) chunk based priority scheduling mechanism will fail because the AR chunk cannot be inserted into the A2A communication task. The performance of FlowMoE will be the same as ScheMoE, Tutel, and FasterMoE, but better than vanillaEP due to the hidden computing task time.\n\n(2) **The computing task time is much longer than the communication task time.** When the computing time is very long, all the communication task time can be covered by the computing task time. In this case, FlowMoE will outperform ScheMoE, Tutel, and FasterMoE due to the hidden AR task time, and better than vanillaEP because it further hides the A2A task time.\n\n(3) **The communication task time is comparable to the computing task time.** In this case, FlowMoE outperforms all baselines because it maximizes the overlap of multi-type tasks.\n\nIn summary, in all cases, the performance of FlowMoE is greater than or equal to that of ScheMoE, Tutel, and FasterMoE, and is always better than that of vanillaEP.\n\n>### **[Question 2]**\n\n**[Answer 4]** ●Firstly, we conducted a detailed analysis on how FlowMoE adapts to heterogeneous GPU clusters. (1) For clusters with mixed GPU types, since All-Reduce (AR) and All-to-All (A2A) tasks can only begin once the slowest GPU completes its corresponding computing task (other GPUs will be idle in waiting time after completing their computing tasks), the task timeline for each iteration of FlowMoE and other baselines is determined by the slowest GPU. In this case, FlowMoE can maximize the computing/communication overlap of the slowest GPU, and its performance still outperforms the baselines. Table 12 in Appendix G.1 also demonstrates the same conclusion when comparing different frameworks on GPU clusters with heterogeneous computing power. (2) For clusters with bandwidth asymmetry, since AR and A2A tasks are collective communications, all GPUs start and finish communication tasks simultaneously regardless of whether their bandwidths are the same. Therefore, the task scheduling timeline is the same for all GPUs. In this case, FlowMoE still achieves better performance than the baseline due to its efficient pipeline solution.\n\n●Secondly, it is not necessary to conduct the BO retuning every time there is a hardware change. We have designed a re-Bayesian tuning mechanism for FlowMoE in Appendix G.2. Specifically, we define a re-execution threshold $\\delta$. We denote the iteration time corresponding to the near-optimal AR chunk size $S_p$ predicted by the last BO as $\\hat{\\mathcal{F}}(S_p^{best})$. If the change of the current iteration time $T$ compared to $\\hat{\\mathcal{F}}(S_p^{best})$ is more than $\\delta$, i.e.,\n$$\n\\frac{\\left | T-\\hat{\\mathcal{F}}(S_p^{best}) \\right |}{\\hat{\\mathcal{F}}(S_p^{best})}> \\delta,\n$$\nthe BO will be re-executed to find a new $S_p$. This is based on the understanding that there is only one optimal AR chunk size to minimize the iteration time for the same hardware cluster (see **[Answer 2]** for detailed analysis). Therefore, when $T$ and $\\hat{\\mathcal{F}}(S_p^{best})$ are close, it indicates that the hardware conditions have not changed or have only slightly fluctuated, and there is no need to readjust the AR chunk size. However, when $T$ and $\\hat{\\mathcal{F}}(S_p^{best})$ differ significantly, it indicates that the hardware conditions have changed significantly, and the optimal AR chunk size has also changed. At this point, the BO needs to be re-executed to update the AR chunk size.\n\n●Thirdly, FlowMoE's performance is sensitive to the BO process. Table 10 in Appendix D.2 compares the per-iteration time of FlowMoE using BO auto-tuning and different fixed AR chunk sizes. The results show that randomly selected fixed AR chunk sizes are likely to severely limit FlowMoE's performance. Therefore, BO auto-tuning is an indispensable part of FlowMoE and can maximize its performance. Although FlowMoE's performance is sensitive to the BO process, it is insensitive to the hyperparameters used in BO and they can be manually adjusted (see **[Answer 2]** for detailed analysis)."}, {"review_id": "6SPCiWKWZT", "reviewer": "Reviewer_gKp7", "summary": "This paper proposed FlowMoE for distributed Mixture-of-Experts (MoE) training that unifies pipeline scheduling across multi-head attention (MHA), gating, expert computation, and both all-to-all (A2A) and all-reduce (AR) communications. By partitioning tensors into micro-batches and introducing a priority-based scheduling mechanism for AR tensor chunks, tuned via Bayesian optimization, FlowMoE maximizes overlap between compute and communication.", "questions": "1. Figure 7 shows loss over time but no convergence stability analysis. We understand from other papers that pipelining can sometimes introduce gradient staleness. This is unaddressed in this paper. Is gradient staleness avoided? If there is gradient staleness that can happen, there is a lack of discussion on whether the model performance is affected or not.\n\n2. Energy savings are measured per worker but lack total cluster-wide analysis.  Are their communication costs not covered by the per worker analysis?", "limitations": "None.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: \n1. The paper addresses a critical bottleneck in distributed MoE training by pipelining all major tasks (MHA, gating, experts, A2A, all-reduce), unlike prior work that only optimizes MoE-layer tasks. .\n2. The design of BO-autotuned tensor partitioning is novel.\n3. The proposed design achieve a significant speedup.\n4. The approach is backed up with meaningful theory\n5. The approach addresses real world issues such as dynamic hardware environments and node dropouts\n\nWeaknesses\n1. Experiments use small clusters (max 16 GPUs) and tiny models (e.g., GPT2-Tiny-MoE with 50M expert params). Claims of \"trillion-parameter scalability\" are unsupported. Also largest model (DeepSeek-V2-S) has only 4 transformer blocks, failing to stress-test the framework. Validate scalability on more devices and larger models will be an added strength.", "was_revised": true, "final_justification": "Thank you for addressing all the issues we raised during the review process.  We believe the paper is a well written and a would be a good contribution to NeurIPS. We will thus keep our rating of (accept) as is.", "rebuttal": "We thank Reviewer gKp7 for the thoughtful feedback and constructive comments. We have responded and addressed every weakness and question in detail except for providing evaluation results for larger-scale GPU clusters. We sincerely apologize for this, because it is very difficult and unrealistic for us to obtain larger-scale hardware resources and measure results within one week. But we will continue to seek larger-scale hardware resource in the future to demonstrate the scalability of FlowMoE. In addition, for weaknesses and questions that have been addressed with additional explanations or experiments in the rebuttal, we will incorporate the corresponding discussions or results into the revised version of the paper. Please see the one-by-one responses below.\n\n>### **[Weakness 1]**\n\n**[Answer 1]** ●Firstly, to stress-test FlowMoE's performance, we measured the training performance of different training frameworks on two scaled-up MoE models (LLaMA2-MoE-L and DeepSeek-V2-M), both of which are close to the memory upper limit of Cluster 1. Specifically, we constructed LLaMA2-MoE-L by increasing the number of transformer blocks in LLaMA2-MoE from 32 to 64. We constructed DeepSeek-V2-M by reducing k to 1 while keeping 7 transformer blocks. The detailed configurations of the two scaled-up MoE models are as follows:\n| MoE Model | # Params (MHA+Gating) | # Params (Experts) | Dataset | L | B | N | M | H | E/P | k | f | \n|-----------------|----------------------------------|----------------------------|-------------|----|----|----|----|------|---|---|---|   \n| LLaMA2-MoE-L | 268.4M | 8595.2M | wikitext-103 | 64 | 4 | 512 | 1024 | 4096 | 1 | 1 | 1.0 | \n| DeepSeek-V2-M | 734.3M | 3524.7M | OpenWebText | 7 | 4 | 256 | 5120 | 1536 | 2 | 1 | 1.0 | \n\nThen, the average per-iteration time using different frameworks on these two MoE models is shown as follows (FasterMoE is OOM in any cases):\n| # of GPUs | Model | VanillaEP | Tutel | ScheMoE | FlowMoE | $S_3$ | $S_2$ | $S_1$ | \n|----------------|---------|---------------|--------|---------------|--------------|----------|----------|-----------|\n| 4 | LLaMA2-MoE-L | 2405.1ms | 1927.0ms | 1806.1ms | 1493.8ms | 1.61$\\times$ | 1.29$\\times$ | 1.21$\\times$ | \n| 4 | DeepSeek-V2-M | 535.3ms | 468.4ms | 432.2ms | 352.2ms | 1.52$\\times$ | 1.33$\\times$ | 1.23$\\times$ | \n| 8 | LLaMA2-MoE-L | 2989.1ms | 2493.9ms | 2297.9ms | 1833.8ms | 1.63$\\times$ | 1.36$\\times$ | 1.25$\\times$ | \n| 8 | DeepSeek-V2-M | 944.6ms | 773.4ms | 723.6ms | 552.4ms | 1.71$\\times$ | 1.40$\\times$ | 1.31$\\times$ | \n| 16 | LLaMA2-MoE-L | OOM | OOM | OOM | OOM | / | / | / | \n| 16 | DeepSeek-V2-M | 1254.6ms | 956.9ms | 893.4ms | 708.8ms | 1.77$\\times$ | 1.35$\\times$ | 1.26$\\times$ |\n\n\\* OOM: out-of-memory  \n\\* $S_1$, $S_2$ and $S_3$ are the speedups of FlowMoE over ScheMoE, Tutel and vanillaEP, respectively.\n\nThe results above show that FlowMoE still achieved the best training performance among all baselines on larger models.\n\n●Secondly, verifying the scalability of FlowMoE on a larger-scale GPU cluster would indeed be an added strength. However, we sincerely apologize that obtaining a larger-scale hardware resource and evaluating the results within a week is very difficult and unrealistic for us. Nevertheless, our current clusters and models are already sufficient to demonstrate the effectiveness of FlowMoE's pipeline mechanism. We will continue to seek larger-scale hardware resource in the future to demonstrate the scalability of FlowMoE.\n\n●Thirdly, we agree with the reviewer that any claim of “trillion-parameter scalability” is indeed inappropriate in our paper since we lack experiments on trillion-parameter models. The sentence: \"The parameter size of modern large language models (LLMs) can be scaled up to the trillion-level via the sparsely-activated Mixture-of-Experts (MoE) technique to avoid excessive increase of the computational costs\", we had in the abstract may have caused misunderstanding. We apologize for that. Thereby, we will modify this sentence to: \"The parameter size of modern large language models (LLMs) can be scaled up to a massive scale using the sparsely-activated Mixture-of-Experts (MoE) technique to avoid excessive increases in computational costs.\" to rectify this issue.\n\n>### **[Question 1]**\n\n**[Answer 2]** Yes, FlowMoE avoids the stale gradient updating problem, and FlowMoE's scheduling mechanism only changes the execution order of computing and communication tasks without affecting the convergence of the model.\n\nWe denote the samples processed in each pipeline as one microbatch, and the number of microbatches is equal to the pipelining degree $R$.\n\nSpecifically, during the backpropagation of the $l$th transformer block in one iteration, the gradients from all $AT_r^l$ and $E_r^l$ ($1 \\le r \\le R$) after backpropagation are accumulated and summed. Once the gradient from $AT_1^l$ is also accumulated, the All-Reduce chunk communication task for the MHA and gate function in the $l$th transformer block is started. Similarly, when the gradient of $E_1^l$ is accumulated, the expert parameters are updated. This effectively prevents the parameters of MHA, the gate function, and the expert from being updated early, thereby avoiding gradient staleness. Additionally, to ensure that the accumulated gradients are equivalent with and without pipelining, we scale the loss calculated for each microbatch by $R$, i.e., $\\frac{loss^{(r)}}{R}$ ($1 \\le r \\le R$), where $loss^{(r)}$ is the loss calculated using the samples from the $r$th microbatch, and the detailed theoretical derivation is as follows:\n\n●The loss and gradient when performing backpropagation using the entire mini-batch are expressed as follows:\n$$\nloss=\\frac{1}{B}{\\textstyle \\sum_{i=1}^{B} \\ell (x_i,y_i)}, \\nabla L_{full}=\\nabla \\left (\\frac{1}{B}{\\textstyle \\sum_{i=1}^{B}\\ell (x_i,y_i)}\\right ),\\tag{1}\n$$  \nwhich is currently used by all mainstream MoE training frameworks [1,2,3].  \n●When dividing mini-batch into $R$ microbatches, the number of samples in each microbatch is $b=\\frac{B}{R}$. The loss for each microbatch is:\n$$\nloss^{(r)}=\\frac{1}{b}{\\textstyle \\sum_{i=1}^{b}\\ell (x_{r,i},y_{r,i}) }.\\tag{2}\n$$  \nThen, the scaled loss for each microbatch is:\n$$\n\\tilde{loss}^{(r)}=\\frac{1}{R}\\cdot loss^{(r)}=\\frac{1}{R}\\cdot \\frac{1}{b}{\\textstyle \\sum_{i=1}^{b}\\ell (x_{r,i},y_{r,i}) }=\\frac{1}{B} {\\textstyle \\sum_{i=1}^{b}\\ell (x_{r,i},y_{r,i})}.\\tag{3}\n$$  \nThe cumulative loss of all microbatches is:\n$$\n {\\textstyle \\sum_{r=1}^{R}\\tilde{loss}^{(r)}} =  {\\textstyle\\sum_{r=1}^{R}\\frac{1}{B} {\\textstyle \\sum_{i=1}^{b}\\ell (x_{r,i},y_{r,i})}}=\\frac{1}{B}\\textstyle\\sum_{r=1}^{R}\\textstyle \\sum_{i=1}^{b}\\ell (x_{r,i},y_{r,i})= \\frac{1}{B}\\textstyle \\sum_{i=1}^{B}\\ell (x_{i},y_{i}).\\tag{4}\n$$  \nAccording to Eqs. 1 and 4, we can obtain:\n$$\n\\nabla \\left (  {\\textstyle \\sum_{r=1}^{R}}\\tilde{loss}^{(r)} \\right )= \\nabla \\left (\\frac{1}{B} {\\textstyle \\sum_{i=1}^{B}}\\ell (x_i,y_i)\\right )=\\nabla L_{full}.\\tag{5}\n$$  \nThis analysis shows that the gradient update strategy is numerically equivalent with and without pipelining, thereby FlowMoE does not introduce gradient staleness. The only difference between them is the optimization of task scheduling order during execution, which does not compromise model convergence during training. We will add this discussion in the Appendix if accepted.\n\n>### **[Question 2]**\n\n**[Answer 3]** Please let us clarify that our measurement process actually does cover the entire cluster-wide energy usage, including both the communication cost and computing cost.\n\nWe use the NVIDIA SMI tool to sample the real-time power of each GPU in the cluster every 5ms, then integrate these sampled data over time to calculate the energy consumption of each GPU during the training process. We sum the energy consumption of all GPUs to obtain the total energy consumption, then divide it by the number of GPUs and the number of iterations to obtain the average per-worker (per-GPU) energy consumption as shown in Table 6. It is worth noting that the power reported by the NVIDIA SMI tool includes both communication and computing costs. According to the official documentation of NVIDIA SMI [4], the energy consumption measured by the NVIDIA SMI tool covers the entire GPU card's power consumption, where communication power includes (1) PCIe communication costs, (2) NVLink/NVSwitch communication costs, and (3) collective communication costs (e.g., NCCL All-Reduce and All-to-All). In addition, some representative works [5,6] also used the NVIDIA SMI tool to monitor GPU communication cost. Therefore, our measurement results effectively reflect the total cluster-wide energy usage, compatibly providing a fair comparison with the literature.\n\n**References:**  \n[1] Tutel: Adaptive Mixture-of-Experts at Scale. *MLSys*, 2023.  \n[2] DeepSpeed-MoE: Advancing Mixture-of-Experts Inference and Training to Power Next-Generation AI Scale. *ICML*, 2022.  \n[3] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism. *arXiv:1909.08053*, 2019.  \n[4] NVIDIA Corporation, *NVIDIA System Management Interface (nvidia-smi) User Guide*.  \n[5] Evaluating Modern GPU Interconnect: PCIe, NVLink, NV-SLI, NVSwitch and GPUDirect. *IEEE TPDS*, 2019.  \n[6] AUGEM: Automatically Generate High Performance Dense Linear Algebra Kernels on x86 CPUs. *SC*, 2013."}, {"review_id": "9CoWBt2lOE", "reviewer": "Reviewer_YDn2", "summary": "The paper proposes a scalable framework for scheduling multi-type task pipelines in the context of distributed Mixture-of-Experts training.", "questions": "- How does FlowMoE compare with other potential optimization strategies, such as different scheduling algorithms or alternative communication patterns?\n- Have the authors considered or compared against other automatic tuning methods beyond Bayesian optimization?\n- How sensitive is the performance to the hyperparameters, particularly the BO search space and acquisition function choice?\n- Can the authors provide deeper theoretical analysis addressing the convergence properties of the proposed scheduling approach?", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strength:\n- The paper addresses a well-motivated and practically important problem in distributed MoE training. The authors provide concrete evidence (Table 1) showing that MHA computing, gating, and all-reduce communication constitute 30-40% of iteration time, which existing frameworks neglect.\n- The paper provides mathematical formulation of the optimization problem and includes theorems (Theorems 1 and 2) that support the design decisions, particularly for the priority scheduling mechanism.\n\nWeaknesses:\n- The core contribution is primarily an engineering optimization that extends existing pipelining techniques. While the unified approach appears valuable, the fundamental algorithmic innovations seem incremental. The engineering contribution, while substantial, is primarily an optimization of existing techniques rather than a fundamental advancement. The impact is thus limited to the specific domain of MoE training optimization.\n- While the paper provides mathematical formulation, the theoretical analysis could be deeper. The convergence properties of the proposed scheduling approach are not analyzed, and the Bayesian optimization component appears basic and potentially fragile across various configurations.\n- The paper primarily compares against existing MoE frameworks (ScheMoE, Tutel, FasterMoE) and vanilla expert parallelism, but lacks comparison with other recent advances in distributed training optimization.", "was_revised": true, "final_justification": "The authors have provided a satisfactory rebuttal. As mentioned in my response below, I now consider the work holds research value more than a simple \"engineering\" tweak. As such, I have upgraded the ratings accordingly. \n\n- Rating: 3 (Borderline Reject) → Revised Score: 4 (Borderline Accept)\n- Quality: 2 → 3 (Fair → Good)", "rebuttal": "We thank Reviewer YDn2 for the thoughtful feedback and constructive comments. We have responded and addressed every weakness and question in detail. Meanwhile, for weaknesses and questions that have been addressed with additional explanations or experiments in the rebuttal, we will incorporate the corresponding discussions or results into the revised version of the paper. Please see the one-by-one responses below.\n\n>### **[Weakness 1]**\n\n**[Answer 1]** We appreciate the opportunity to clarify the innovation and impact of our work.\n\n●Firstly, the MoE architecture has become one of the dominant architectures for modern LLMs, including LLaMA4 (Meta, 2025), DeepSeek-V3 (2024), Grok-1 (xAI, 2024), and DBRX (Databricks, 2024). As MoE models become more prevalent,  optimizing their distributed training is both critical and urgent. We believe FlowMoE will make a meaningful contribution to the development of this field.\n\n●Secondly, existing pipelining techniques focus on overlapping of All-Reduce (AR) and computing, or All-to-All (A2A) and expert computing, but leave two key issues unresolved:  (1) optimal coexistence of AR and A2A, and (2) a unified pipeline scheduling scheme for all major MoE tasks. FlowMoE explicitly addresses these two issues. Therefore, FlowMoE is not an optimization of existing pipeline technologies but proposes a new pipeline acceleration scheme, representing a fundamental innovation.\n\n●Thirdly, we introduce a formal task-dependency pipelining model and prove that our scheduling strategy strictly reduces or matches per-iteration time compared to sequential or partially pipelined scheduling (see Theorems 1 and 2). This mathematical foundation distinguishes our method from heuristic-driven pipeline strategies, going beyond standard engineering optimization.\n\n●Fourly, FlowMoE not only focuses on system-level implementation but also emphasizes algorithm design for the pipelining scheme and AR chunk based priority scheduling mechanism (e.g., Algs. 1 and 2 in Sec. 4.2). These contributions clearly extend the impact beyond the specific domain of MoE training optimization.\n\n●Finally, our submission is highly consistent with the NeurIPS infrastructure track (e.g., libraries, improved implementations and scalability, distributed solutions). FlowMoE provides an efficient distributed solution that improves the implementation of MoE training systems. Additionally, the AI community has been extensively focusing on research to accelerate distributed training through pipelining techniques [1-7], and we believe that FlowMoE's contributions represent a fundamental advancement in this important and evolving field.\n\n>### **[Weakness 2, Question 3, Question 4]**\n\n**[Answer 2]** To deepen the theoretical analysis, we analyze (1) the convergence of our scheduling method and (2) FlowMoE's sensitivity to Bayesian optimization (BO) configurations. Since the above two theoretical analyses also address Questions 3 and 4, we provide a unified response here.\n\n(1) **Convergence of the proposed scheduling method.**  \nFlowMoE only changes the execution order of computing and communication tasks without affecting model convergence. We denote the samples processed in each pipeline as one microbatch, and the number of microbatches is equal to the pipelining degree $R$.\n\nSpecifically, during the backpropagation of the $l$th transformer block in one iteration, the gradients from all $AT_r^l$ and $E_r^l$ ($1 \\le r \\le R$) are accumulated and summed. Once the gradient from $AT_1^l$ is also accumulated, the All-Reduce chunk communication task for the MHA and gate is started. Similarly, when the gradient of $E_1^l$ is accumulated, the expert parameters are updated. This effectively prevents the impact of gradient staleness on convergence from early parameter updates. Additionally, to ensure that the accumulated gradients are equivalent with and without pipelining, the loss of each microbatch is scaled by $\\frac{1}{R}$, i.e., $\\frac{loss^{(r)}}{R}$ ($1 \\le r \\le R$), and the detailed derivation follows:\n\n●The loss and gradient when performing backpropagation using the entire mini-batch are expressed as follows:\n$$\nloss=\\frac{1}{B}{\\textstyle\\sum_{i=1}^{B}\\ell (x_i,y_i)}, \\nabla L_{full}=\\nabla\\left (\\frac{1}{B}{\\textstyle\\sum_{i=1}^{B}\\ell (x_i,y_i)}\\right ),\\tag{1}\n$$  \nwhich is currently used by all mainstream MoE training frameworks, e.g., Tutel, DeepSpeed-MoE and Megatron-LM.  \n●When dividing mini-batch into $R$ microbatches, the number of samples in each microbatch is $b=\\frac{B}{R}$. The loss for each microbatch is:\n$$\nloss^{(r)}=\\frac{1}{b}{\\textstyle\\sum_{i=1}^{b}\\ell (x_{r,i},y_{r,i})}.\\tag{2}\n$$  \nThen, the scaled loss for each microbatch is:\n$$\n\\tilde{loss}^{(r)}=\\frac{1}{R}\\cdot loss^{(r)}=\\frac{1}{R}\\cdot \\frac{1}{b}{\\textstyle\\sum_{i=1}^{b}\\ell (x_{r,i},y_{r,i})}=\\frac{1}{B} {\\textstyle\\sum_{i=1}^{b}\\ell (x_{r,i},y_{r,i})}.\\tag{3}\n$$  \nThe cumulative loss of all microbatches is:\n$$\n{\\textstyle\\sum_{r=1}^{R}\\tilde{loss}^{(r)}}={\\textstyle\\sum_{r=1}^{R}\\frac{1}{B}{\\textstyle\\sum_{i=1}^{b}\\ell (x_{r,i},y_{r,i})}}=\\frac{1}{B}\\textstyle\\sum_{r=1}^{R}\\textstyle\\sum_{i=1}^{b}\\ell (x_{r,i},y_{r,i})=\\frac{1}{B}\\textstyle \\sum_{i=1}^{B}\\ell (x_{i},y_{i}).\\tag{4}\n$$  \nAccording to Eqs. 1 and 4, we can obtain:\n$$\n\\nabla\\left ({\\textstyle\\sum_{r=1}^{R}}\\tilde{loss}^{(r)} \\right )=\\nabla\\left (\\frac{1}{B}{\\textstyle\\sum_{i=1}^{B}}\\ell (x_i,y_i)\\right )=\\nabla L_{full}.\\tag{5}\n$$  \nThis analysis shows that the gradient update strategy is numerically equivalent with and without pipelining, and FlowMoE does not introduce gradient staleness. The only difference between them is the optimization of task scheduling order during execution, which does not compromise model convergence during training.\n\n(2) **Sensitivity analysis of FlowMoE to BO configurations.**  \nFlowMoE is insensitive to BO hyperparameters, which can be manually adjusted.\n\nFirstly, the BO search space should be fixed as (0MB, max tensor size to be communicated per transformer block] to ensure that the optimal AR chunk size is included. Thus, the BO configuration only involves two hyperparameters: the acquisition function and surrogate model. Then, an oversized AR chunk size will affect the prioritization of A2A tasks and may prevent them from being started in time, while an undersized AR chunk size will cause excessive startup overhead. Therefore, there must exist a unique optimal solution that maximizes training speed. For this single-peaked, smooth objective function with a fixed search space, BO consistently converges near the optimum, largely unaffected by hyperparameter variations. The table below shows that different BO settings have minor effect on FlowMoE performance when training BERT-Large-MoE on Cluster 1.\n\n|Acquisition Function|Surrogate Model|Average Per-iteration Time(ms)|\n|-|-|-|\n|EI (0.1)|GPR + Matern|351.9|\n|EI (0.05)|GPR + Matern|358.9|\n|EI (0.2)|GPR + Matern|354.2|\n|Probability of Improvement|GPR + Matern|355.1|\n|Lower Confidence Bound|GPR + Matern|355.4|\n|EI (0.1)|GPR + RBF|357.2|\n|EI (0.1)|GPR + Rational Quadratic|360.2|\n\n\\*EI: Expected Improvement  \n\\*Surrogate Model use Gaussian Process Regression (GPR) with different kernel functions.\n\nWe will add the above theoretical analysis in the Appendix if accepted.\n\n>### **[Weakness 3, Question 1]**\n\n**[Answer 3]** ●Firstly, we added a new recent baseline, FSMoE[8]. FSMoE focuses on pipelining for expert computing and communication tasks within the MoE layer, and it maximizes computation-communication overlap within the MoE layer by reasonably scheduling inter-node  and intra-node communication tasks. The table below shows the comparison of average per-iteration time of FlowMoE and FSMoE on Cluster 1.\n\n|Model|# of GPUs|FSMoE|FlowMoE|$S_1$|\n|-|-|-|-|-|\n|GPT2-Tiny-MoE|4|82.8ms|66.1ms|1.25$\\times$|\n||8|98.8ms|87.6ms|1.13$\\times$|\n||16|114.8ms|95.6ms|1.20$\\times$|\n|BERT-Large-MoE|4|278.1ms|239.5ms|1.16$\\times$|\n||8|345.1ms|283.2ms|1.22$\\times$|\n||16|421.9ms|351.9ms|1.19$\\times$|\n|LLaMA2-MoE|4|928.1ms|763.9ms|1.21$\\times$|\n||8|1110.4ms|906.0ms|1.23$\\times$|\n||16|1292.6ms|1124.0ms|1.15$\\times$|\n|DeepSeek-V2-S|4|2096.3ms|1740.8ms|1.20$\\times$|\n||8|2985.4ms|2384.9ms|1.25$\\times$|\n||16|3895.62ms|3205.3ms|1.22$\\times$|\n\n\\* $S_1$ is the speedups of FlowMoE over FSMoE.\n\nThe results show that FlowMoE still achieves better performance gain compared to FSMoE. This is because FlowMoE implements pipeline scheduling for all major MoE-related tasks and introduces BO to adaptively adjust the AR chunk size to maximize computation-communication overlap across the entire transformer block. It is worth noting that the scheduling strategies of FSMoE and FlowMoE are orthogonal and can be combined to further accelerate MoE training.\n\n●Secondly, Table 9 in Appendix D.2 also shows the superiority of BO over other auto-tuning optimization strategies in FlowMoE.\n\n>### **[Question 2]**\n\n**[Answer 4]** Yes, we have compared BO, Grid Search, and Random Number Generation in Appendix D.2. The results from Table 9 show that using BO to adjust the AR chunk size achieves the shortest per-iteration time among the three methods.\n\n**References:**  \n[1] Pipeline Parallelism with Controllable Memory. *NeurIPS*, 2024.  \n[2] Rethinking Memory and Communication Costs for Efficient Data Parallel Training of Large Language Models. *NeurIPS*, 2024.  \n[3] SAPipe: Staleness-Aware Pipeline for Data Parallel DNN Training. *NeurIPS*, 2022.  \n[4] Theoretical Limits of Pipeline Parallel Optimization and Application to Distributed Deep Learning. *NeurIPS*, 2019.  \n[5] Pipe-SGD: A Decentralized Pipelined SGD Framework for Distributed Deep Net Training. *NeurIPS*, 2018.  \n[6] PipeOffload: Improving Scalability of Pipeline Parallelism with Memory Optimization. *ICML*, 2025.  \n[7] Pfeife: Automatic Pipeline Parallelism for PyTorch. *ICML*, 2025.  \n[8] FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Experts Models. *ASPLOS*, 2025."}, {"review_id": "IbCnnoE8xt", "reviewer": "Reviewer_GGaM", "summary": "This paper presents a pipeline scheduling framework named FlowMoE for Mixture-of-Expert (MoE) model training. Instead of only pipelinining within the MoE layer, it proposes that 1) other layers including MHA/gating/allreduce/expert layers can also get pipelined in order to achieve higher utilization, and 2) allreduce(AR) ops can be chunked by tensors so as to filling in the gaps between A2A. This paper also presents a BO-based method to find best allreduce chunks and a priority scheduling algorithm for A2A and AR.The evaluations were conducted on a variety of model architectures and the results show convincing speedup over exisiting works.", "questions": "1. Could you comment on the GPU utilization when using 1) a smaller batch size, and 2) large number of experts with small number of activated experts? \n2. (Section 4.2) Could you clarify if bayesian optimization for AR/A2A scheduling is on the critical path? Can it be made async with respect to training loop in order to further hide the overhead?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n* This paper addresses some less explored areas in MoE training and proposes effective solutions: overlapping for non-MoE layers and priority scheduling of collective communications. The unified pipeline scheduling framework could inspire future work on auto-parallelism for MoE models. \n* The authors augment their claims with strong evidence empirically and theoretically, showing effectiveness of their approach. \n* This paper is well written and easy to follow. \n\nWeakness:\n* It is unclear if microbatching on MHA in forward pass and sharded ARs in backward pass will reduce device utilization. It is my understanding that they will reduce the arithmatic intensity, but the GPU SM utilization (or occupancy) is missing from the evaluation.", "was_revised": true, "final_justification": "The authors have address my concerns and I would like to maintain my score.", "rebuttal": "We thank Reviewer GGaM for the thoughtful feedback and constructive comments. We have responded and addressed every weakness and question in detail. Meanwhile, for weaknesses and questions that have been addressed with additional explanations or experiments in the rebuttal, we will incorporate the corresponding discussions or results into the revised version of the paper. Please see the one-by-one responses below.\n\n>### **[Weakness 1]**\n\n**[Answer 1]** We conducted experiments on Cluster 1 with 16 GPUs to measure GPU utilization under different microbatch sizes. Specifically, we adjusted the microbatch size by changing the pipelining degree $R$ (a larger $R$ results in a smaller microbatch). VanillaEP represents the original mini-batch without pipelining. We used the NVIDIA SMI tool to sample the utilization of all GPUs every 5ms and averaged the values over 1000 iterations. Additionally, we reported the standard deviation of all GPU utilization values to illustrate the degree of data variability. The experimental results are as follows: \n\n|  | Model | $R$ | Average GPU Utilization | Standard Deviation |\n|--------|--------|--------|--------|--------|\n| FlowMoE | GPT2-Tiny-MoE | 2 | 94.96% | 13.35% | \n| FlowMoE | GPT2-Tiny-MoE | 4 | 69.34% | 21.04% | \n| vanillaEP | GPT2-Tiny-MoE | /  | 91.16% | 40.40% | \n| FlowMoE | BERT-Large-MoE | 2 | 96.94% | 3.28% | \n| FlowMoE | BERT-Large-MoE | 4 | 96.14% | 5.91% | \n| vanillaEP | BERT-Large-MoE | /  | 98.67% | 12.59% | \n| FlowMoE | LLaMA2-MoE | 2 | 99.49% | 2.98% |\n| FlowMoE | LLaMA2-MoE | 4 | 98.98% | 2.67% |\n| vanillaEP | LLaMA2-MoE | /  | 99.24% | 15.60% |\n| FlowMoE | DeepSeek-V2-S | 2 | 98.26% | 12.08% |\n| FlowMoE | DeepSeek-V2-S | 4 | 94.93% | 11.10% |\n| vanillaEP | DeepSeek-V2-S | /  | 98.46% | 23.29% |\n\nBased on the above results, we have the following conclusions:  \n\n(1) According to the comparison of average GPU utilization, smaller microbatch sizes may result in lower GPU utilization (e.g., when training GPT2-Tiny-MoE with $R=4$ using FlowMoE), but in most cases, GPU utilization is nearly identical to that without pipelining (e.g., when training BERT-Large-MoE, LLaMA2-MoE and DeepSeek-V2-S using FlowMoE). This is because the actual tensors involved in GPU computation remain sufficiently large for larger MoE models, and GPU SM resources are still utilized efficiently. In other words, the microbatches introduced by FlowMoE's pipelining do not waste GPU SM resources for large MoE models;  \n\n(2) According to the standard deviation of all GPU utilization values, compared to vanillaEP, FlowMoE maintains relatively stable GPU utilization during training (e.g., the standard deviation is even less than 6% when training BERT-Large-MoE and LLaMA2-MoE). This is because FlowMoE overlaps computing and communication tasks, and the GPU rarely enters idle states waiting for communication. In contrast, vanillaEP's communication tasks must be executed separately, resulting in low GPU utilization during communication task execution.\n\nWe will add the above discussion and results on the GPU utilization in the Appendix if the paper is accepted.\n\n>### **[Question 1]**\n\n**[Answer 2]** ●Firstly, smaller batch sizes are more likely to result in lower GPU utilization. We measured GPU utilization with different batch sizes when training different MoE models using FlowMoE on Cluster 1 with 16 GPUs. Similarly, we used the NVIDIA SMI tool to sample the utilization of all GPUs every 5ms and averaged the values over 1000 iterations. The measurement results are as follows:\n\n| Model | Batch Size | Average GPU Utilization |\n|--------|--------|--------|\n| GPT2-Tiny-MoE | 4 | 94.96% | \n| GPT2-Tiny-MoE | 2 | 65.68% | \n| BERT-Large-MoE | 4 | 96.94% | \n| BERT-Large-MoE | 2 | 84.23% | \n| LLaMA2-MoE | 4 | 99.49% | \n| LLaMA2-MoE | 2 | 98.52% | \n| DeepSeek-V2-S | 4 | 98.26% | \n| DeepSeek-V2-S | 2 | 96.28% | \n\nThe results demonstrate that smaller batch sizes are more likely to lead to lower GPU utilization, e.g., when training GPT2-Tiny-MoE or BERT-Large-MoE with a batch size of 2. Moreover, GPU utilization remained nearly unchanged when training LLaMA2-MoE or DeepSeek-V2-S. This is because these two MoE models have a larger number of parameters, which keeps the tensor dimensions involved in GPU computation sufficiently large, ensuring that SM resources are still efficiently utilized even with a lower batch size.\n\n●Secondly, when using a large number of experts with a small number of activated experts, the utilization of GPUs will vary significantly. GPUs hosting frequently activated experts have higher utilization, while those with rarely activated experts experience underutilization.\n\nWe also evaluated the impact of the number of activated experts on GPU utilization when using FlowMoE on Cluster 1 with 16 GPUs for a large number of experts. Specifically, we constructed BERT-Large-MoE-w by increasing the number of experts per GPU from 2 to 8 and adjusting the number of activated experts by modifying the capacity factor $f$, where a larger $f$ indicates more uneven token routing by the gate function and fewer activated experts since most tokens are routed to popular experts. The detailed configurations of BERT-Large-MoE-w are as follows:\n\n| MoE Model | # Params (MHA+Gating) | # Params (Experts) | Dataset | L | B | N | M | H | E/P | k |\n|-----------------|----------------------------------|----------------------------|-------------|----|----|----|----|------|---|---| \n| BERT-Large-MoE-w | 25.2M | 3325.9M | wikitext-103 | 24 | 4 | 512 | 512 | 1024 | 8 | 1 |\n\nWe used the NVIDIA SMI tool to sample the utilization of each GPU every 5ms and averaged the utilization of each GPU over 1000 iterations. We reported the maximum and minimum GPU utilization in the table below:\n\n| Model | $f$ | Maximum GPU Utilization | Minimum GPU Utilization |\n|--------|--------|--------|--------|\n| BERT-Large-MoE-w | 1.0 | 97.63% | 96.14% | \n| BERT-Large-MoE-w | 4.0 | 98.37% | 56.58% |  \n| BERT-Large-MoE-w | 8.0 | 98.92% | 38.16% |  \n| BERT-Large-MoE-w | 16.0 | 99.23% | 26.92% |  \n\nThe above results show that, with a large number of experts, the fewer the number of activated experts, the more unbalanced the computation load on the GPU and the greater the difference in utilization between GPUs.\n\nWe will add the above discussion and results on the GPU utilization in the Appendix if the paper is accepted.\n\n>### **[Question 2]**\n\n**[Answer 3]** The Bayesian optimization (BO) process used to adjust the AR chunk size ($S_p$) is not a critical path in the training loop and can be made asynchronous with respect to the training loop. However, BO is crucial to FlowMoE. Meanwhile, FlowMoE's performance is not sensitive to BO hyperparameters, and users do not need to spend time carefully designing these parameters.\n\nSpecifically, BO is only executed during the initial profiling phase (typically the first ~80 iterations), during which it continuously suggests and samples a small number of ($S_p$, per-iteration time) pairs (8 by default). Then, BO predicts an approximate optimal $S_p$ based on these sampling pairs and uses it for subsequent training. All training iterations during the initial profiling phase use the existing FlowMoE pipeline and run normally. Therefore, we agree with the reviewer that BO can be redesigned to piggyback on standard iterations (e.g., using a background thread or dedicating evaluation to a shadow process, with BO and the main training loop process exchanging only sampled values and updated $S_p$) without blocking training, thereby further reducing even this minimal overhead. We view this as a valuable direction for future research."}]

Author rebuttal:

Meta-review: Overall, the paper will be a good contribution to NeurIPS which is also a consensus among the reviewers.

(a) Summary of Claims

The paper proposes FlowMoE, a pipeline scheduling framework for distributed MoE training. It unifies scheduling of different components such as MHA, gating, expert computation, A2A, and AR communication, and introduces a tensor chunk-based priority mechanism tuned via Bayesian Optimization (BO). The authors show empirical speedup, energy savings, and memory reduction compared to different MoE frameworks.

(b) Strengths

- Tackles a practical bottleneck in MoE training with a unified and adaptive scheduling framework.
- Strong empirical results.
- Theoretical backing (Theorems 1 & 2) and robustness to BO hyperparameters.

(c) Weaknesses

- Novelty is limited; some reviewers view it as an engineering extension of prior work.
- Experiments are limited to 16 GPUs and sub-trillion parameter models, despite scalability claims.
- GPU utilization initially measured with less precise tools (later addressed).

(d) Justification for Recommendation

Despite concerns about novelty, the paper presents an important system contribution. In particular, the unified scheduling of heterogeneous tasks and integration of BO tuning are valuable. Therefore, this is a good contribution to NeurIPS.

(e) Rebuttal Summary and Discussion

Reviewer hxBa maintained a borderline accept, citing limited novelty even with satisfactory responses on BO tuning.
Reviewer gKp7 appreciated the added experiments and theoretical clarifications, maintaining an accept.
Reviewer YDn2 upgraded to borderline accept after recognizing the co-design of scheduling policies as a meaningful contribution.
Reviewer GGaM accepted the paper and appreciated the authors’ follow-up with more accurate GPU utilization metrics.
Decision: Accept (poster)