Title: FlowMoE: A Scalable Pipeline Scheduling Framework for Distributed Mixture-of-Experts Training
Venue: NeurIPS 2025
Accepted: yes
Keywords: Distributed Deep Learning

Large Language Model

Mixture-of-Experts

Pipeline Scheduling
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a scalable framework for scheduling multi-type task pipelines in the context of distributed Mixture-of-Experts training.\n\nSTRENGTHS AND WEAKNESSES:\nStrength:\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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: \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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nNone.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n# 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.\n\nQUESTIONS:\n* 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?\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-05"}