FlowMoE: Pipeline Scheduler for Distributed MoE

• FlowMoE is a framework for distributed MoE training that pipelines the entire transformer block, integrating MHA, gating, expert compute, and communications into a unified schedule.
• It employs a tensor chunk-based priority scheduling mechanism that overlaps all-reduce operations with computation, reducing training time by 13%-57% and energy by 10%-39%.
• The adaptive design leverages Bayesian optimization and a unified task graph to efficiently manage compute and communication, optimizing resources on PyTorch without backend modifications.

FlowMoE most directly denotes a distributed training framework for Mixture-of-Experts models that treats the transformer block, rather than only the MoE sublayer, as the scheduling unit. Its defining idea is to place MHA computing, gating, expert computing, all-to-all communication, and all-reduce communication into a single pipeline, then use a tensor chunk-based priority scheduling mechanism so that all-reduce can be overlapped with ongoing compute and communication. Implemented as an adaptive and generic framework atop PyTorch, FlowMoE was reported to reduce training time by 13%-57%, energy consumption by 10%-39%, and memory usage by 7%-32% relative to state-of-the-art MoE training frameworks (Gao et al., 30 Sep 2025).

1. Problem setting and design objective

FlowMoE is formulated for standard expert parallelism in distributed MoE training. In this setting, the MHA layer and gating function are replicated across workers and trained in data parallel style, while the experts are distributed across workers, so each worker holds only a subset of experts. For input tensor $I \in \mathbb{R}^{B \times N \times M}$, MHA produces $I' \in \mathbb{R}^{B \times N \times M}$, and the gating output is described as $G(I') \in \mathbb{R}^{E \times C \times M}$, where each expert can receive up to

$C = f \times k \times B \times N / E.$

Here $L$ is the number of transformer blocks, $B$ the samples per GPU per iteration, $N$ the number of tokens per sample, $M$ the token embedding size, $H$ the expert FFN hidden size, $E$ the number of experts, $I' \in \mathbb{R}^{B \times N \times M}$0 the top-$I' \in \mathbb{R}^{B \times N \times M}$1 routed experts per token, and $I' \in \mathbb{R}^{B \times N \times M}$2 the capacity factor (Gao et al., 30 Sep 2025).

The framework is motivated by a systems diagnosis: prior MoE schedulers largely optimize only the MoE layer, especially overlap between expert computation and dispatch/combine A2A, while leaving MHA, gating, and all-reduce outside the scheduling model. On a 16-GPU RTX3090 cluster, the sum of MHA+gating and all-reduce contributes 33.1% for GPT2-Tiny-MoE, 29.8% for BERT-Large-MoE, 34.2% for LLaMA2-MoE, and 36.1% for DeepSeek-V2-S. FlowMoE’s central claim is therefore that MoE-layer-only scheduling is structurally incomplete for sparse LLM training.

2. Unified pipeline across heterogeneous task types

FlowMoE constructs a unified pipeline by partitioning the transformer-block input into $I' \in \mathbb{R}^{B \times N \times M}$3 equal parts and expressing nearly all work as a common task graph. The task set is

$I' \in \mathbb{R}^{B \times N \times M}$4

where $I' \in \mathbb{R}^{B \times N \times M}$5 denotes the MHA + gating subtask, $I' \in \mathbb{R}^{B \times N \times M}$6 the dispatch A2A subtask, $I' \in \mathbb{R}^{B \times N \times M}$7 the expert compute subtask, $I' \in \mathbb{R}^{B \times N \times M}$8 the combine A2A subtask, and $I' \in \mathbb{R}^{B \times N \times M}$9 the all-reduce task for block $G(I') \in \mathbb{R}^{E \times C \times M}$0 (Gao et al., 30 Sep 2025).

Symbol	Meaning
$G(I') \in \mathbb{R}^{E \times C \times M}$1	MHA + gating subtask
$G(I') \in \mathbb{R}^{E \times C \times M}$2	dispatch A2A subtask
$G(I') \in \mathbb{R}^{E \times C \times M}$3	expert compute subtask
$G(I') \in \mathbb{R}^{E \times C \times M}$4	combine A2A subtask
$G(I') \in \mathbb{R}^{E \times C \times M}$5	all-reduce task

The feed-forward compute order is staged as

$G(I') \in \mathbb{R}^{E \times C \times M}$6

while A2A tasks follow

$G(I') \in \mathbb{R}^{E \times C \times M}$7

Backward compute and backward A2A follow the reverse staged orders defined in the paper. The resulting execution “flow” is the source of the framework’s name: chunks of MHA/gating, A2A, and expert work advance through the block in lockstep.

The scheduling abstraction assumes that one compute task and one communication task can overlap, that two compute tasks cannot overlap and two communication tasks cannot overlap, and that there is no preemption. Under these assumptions, FlowMoE formulates backward execution as a dependency-constrained optimization with objective

$G(I') \in \mathbb{R}^{E \times C \times M}$8

This formalization matters because it expands the scheduling domain from “expert kernels plus A2A” to the entire transformer-block critical path.

3. All-reduce chunking and priority scheduling

FlowMoE’s second major contribution is the treatment of all-reduce as a low-priority, gap-filling communication workload rather than a centralized tail at the end of backward. The paper proves that if the scheduling order satisfies its dependency constraints, then the backward time $G(I') \in \mathbb{R}^{E \times C \times M}$9 is no worse than centralized all-reduce scheduling, written as $C = f \times k \times B \times N / E.$0 (Gao et al., 30 Sep 2025).

The mechanism is tensor chunk-based priority scheduling. Each all-reduce tensor is partitioned into chunks of size $C = f \times k \times B \times N / E.$1, and these chunks are inserted into an ARQueue. Communication is then scheduled from two queues: an A2AQueue and an ARQueue, with A2A tasks having higher priority. The communication manager executes an A2A task whenever one is ready, and executes an all-reduce chunk only when the A2A queue is empty. This creates a non-preemptive priority scheduler in which all-reduce occupies communication gaps without delaying latency-critical A2A.

The paper further proves an idealized result: if the scheduling order satisfies the same constraints, then the time per iteration is minimized when $C = f \times k \times B \times N / E.$2 and there is no startup overhead for communicating all-reduce tensor chunks. In practice, that limit is unattainable because smaller chunks increase communication startup overhead. FlowMoE therefore treats $C = f \times k \times B \times N / E.$3 as a tuning knob and uses Bayesian optimization to find a near-optimal value.

This design changes the role of all-reduce in distributed MoE training. Instead of being a monolithic synchronization step, it becomes a fragmented background workload that is opportunistically inserted between higher-priority communication events. A plausible implication is that FlowMoE’s performance advantage depends not only on expert sparsity, but also on the amount of exploitable communication slack in a given cluster and model configuration.

4. System architecture, implementation, and computational model

FlowMoE is implemented at the PyTorch API level, rather than by deeply modifying the backend engine or communication library. The framework is described as adaptive because it automatically tunes $C = f \times k \times B \times N / E.$4 using Bayesian optimization, and generic because it sits atop PyTorch and existing distributed stacks (Gao et al., 30 Sep 2025).

The implementation is organized around four modules: the Task Breakdown Manager, the BO autotuner, the Communication Task Pool, and the Pipeline Scheduling Manager. Because per-block backward readiness is not directly visible in PyTorch, FlowMoE uses register_full_backward_hook to access gradients of MHA and gating parameters as soon as each transformer block finishes backward; these gradients are then chunked and enqueued for all-reduce. The runtime uses multithreading, with the main thread handling task breakdown, BO autotuning, and pipeline scheduling, and a sub-thread managing the communication pool. A threading.Lock is used for thread safety.

For a gating function implemented as a linear layer of size $C = f \times k \times B \times N / E.$5, the paper states that MHA + gating parameters per block have size

$C = f \times k \times B \times N / E.$6

and gives the scheduling overhead per iteration as

$C = f \times k \times B \times N / E.$7

That overhead is reported to remain below 1% of iteration time in experiments. The MoE model iteration complexity is given as

$C = f \times k \times B \times N / E.$8

with the terms corresponding to MHA projections and attention score computation, gating, and expert FFN on each GPU.

The framework also includes a semantic equivalence argument. By scaling microbatch losses by $C = f \times k \times B \times N / E.$9, the accumulated scaled loss matches the full-batch loss, and the accumulated gradient equals the full-batch gradient. This means FlowMoE changes the execution order of distributed MoE training, not the mathematical training objective itself.

5. Empirical evaluation and observed behavior

The evaluation covers 675 typical/customized MoE layers and four real-world MoE models across two GPU clusters. The customized benchmark spans $L$0, $L$1, $L$2, $L$3, $L$4, $L$5, and $L$6. The real-world models are GPT2-Tiny-MoE, BERT-Large-MoE, LLaMA2-MoE, and DeepSeek-V2-S; larger stress-test models LLaMA2-MoE-L and DeepSeek-V2-M are also included. Cluster 1 consists of 2 nodes, 100 Gb/s interconnect, and 8 × NVIDIA RTX3090 per node; Cluster 2 consists of 4 nodes, 10 Gb/s interconnect, and 2 × RTX2080Ti per node. Baselines are vanillaEP, FasterMoE, Tutel, FSMoE, and ScheMoE, and results are averaged over 1000 iterations (Gao et al., 30 Sep 2025).

Across end-to-end experiments, FlowMoE is reported as 14%–31% faster than ScheMoE, 13%–25% faster than FSMoE, 29%–42% faster than Tutel, 26%–57% faster than FasterMoE, and 43%–82% faster than vanillaEP. On the 675-layer benchmark, it is faster than ScheMoE in all valid cases, with an average speedup over ScheMoE of 26%. The paper also reports per-worker energy savings of 10%–16% versus ScheMoE, 22%–27% versus Tutel, 33%–39% versus FasterMoE, and 33%–41% versus vanillaEP, while memory savings reach up to 7% versus ScheMoE, 9% versus Tutel, 32% versus FasterMoE, and 11% versus vanillaEP.

The ablation study isolates the two main design components. On a customized MoE layer with $L$7, $L$8, $L$9, $B$0, $B$1, and 16 GPUs, the reported iteration times are: 1630.8 ms for vanillaEP, 1115.2 ms for Tutel, 1012.6 ms for FlowMoE-AT, 971.5 ms for FlowMoE-AR with fixed $B$2MB, 895.3 ms for FlowMoE-AR(BO), and 796.1 ms for full FlowMoE. The paper interprets this as showing that pipelining MHA+gating already contributes a meaningful gain beyond MoE-layer-only scheduling, while all-reduce chunk scheduling adds a larger gain, and BO tuning is necessary because performance varies substantially with fixed chunk sizes.

The sensitivity study over pipeline degree $B$3 shows that the best value is model-dependent. For DeepSeek-V2-S on 16 GPUs, FlowMoE reports 3205.3 ms at $B$4, 3113.8 ms at $B$5, and 3295.9 ms at $B$6. The paper characterizes $B$7-selection as orthogonal to the framework and notes that methods such as PipeMoE’s can be applied.

6. Scope, limitations, and related usages of the name

FlowMoE assumes a specific overlap model: one compute task and one communication task can run concurrently, but multiple compute tasks or multiple communication tasks cannot. It also assumes no preemption and uniform subtask execution times for partitioned tasks of the same type. The paper states that benefits are less pronounced in bandwidth-limited scenarios where communication dominates so heavily that pipeline overlap helps less as cluster size grows. It also emphasizes that chunk-size tuning is crucial, because poor $B$8 choices can substantially reduce performance, and the implementation depends on timely access to gradients through PyTorch hooks (Gao et al., 30 Sep 2025).

The name “FlowMoE” also sits near several distinct research threads, which can create terminological confusion. FoMoE is a federated MoE training system for weakly connected multi-site pretraining that breaks the “full replica everywhere” assumption through partial expert replication; its novelty is synchronization-payload sparsification rather than transformer-block pipeline scheduling (Sani et al., 17 Jun 2026). MoE-FM, instantiated in YAN, uses mixture-of-experts flow matching for non-autoregressive language modeling; here the “flow” refers to continuous-time latent transport rather than distributed scheduling (Li, 16 Apr 2026). FP8-Flow-MoE is an FP8-centric MoE training recipe that introduces a scaling-aware transpose and fused FP8 operators to eliminate most explicit casts; it addresses quantization-consistent FP8 dataflow, not pipeline scheduling (Wang et al., 4 Nov 2025). FloE is an on-the-fly MoE inference system for memory-constrained GPUs based on hybrid expert compression, routing-aware prefetch, and PCIe traffic reduction; it concerns serving rather than distributed training (Zhou et al., 9 May 2025).

This suggests that “FlowMoE” is best understood in the narrow technical sense of the framework introduced in (Gao et al., 30 Sep 2025): a whole-transformer-block pipeline scheduler for distributed MoE training. The broader naming landscape, however, reflects a recurrent systems theme across MoE research: performance bottlenecks often emerge not from sparse computation alone, but from the organization of dataflow, whether that dataflow is across GPUs, datacenters, quantization domains, or the CPU–GPU memory hierarchy.