学术研究报告：LAER-MOE——用于高效混合专家模型训练的负载自适应专家重布局系统
一、 主要作者、机构与发表信息
本研究报告的论文《LAER-MOE: Load-Adaptive Expert Re-layout for Efficient Mixture-of-Experts Training》由以下作者共同完成：刘欣怡（Xinyi Liu，北京大学）、王宇杰（Yujie Wang，北京大学）、付方诚（Fangcheng Fu，上海交通大学）、肖雪峰（Xuefeng Xiao，字节跳动Seed）、李会霞（Huixia Li，字节跳动Seed）、李佳时（Jiashi Li，字节跳动Seed）以及崔斌（Bin Cui，北京大学，通讯作者）。该研究于2026年3月22日至26日，在美国匹兹堡举行的第31届ACM编程语言与操作系统架构支持国际会议（ASPLOS ‘26）上发表。
二、 研究背景与目标
本研究隶属于人工智能与高性能计算交叉领域，具体聚焦于大规模语言模型（Large Language Models, LLMs）的高效分布式训练。随着模型参数规模的指数级增长，计算与内存需求成为训练万亿参数级模型的重大挑战。混合专家（Mixture-of-Experts， MoE）架构通过为每个输入令牌（token）动态激活少数专家网络，能在保持计算量相对恒定的情况下显著增加模型总参数量，已成为扩展模型规模的关键技术。
在分布式训练Moe模型时，专家并行（Expert Parallelism， EP）是一种主流技术，它将不同的专家放置在不同的计算设备（如GPU）上。然而，MoE层中动态路由机制导致令牌被不均匀地分配给各个专家，从而引发严重的负载不平衡问题：少数“热门”专家（接收大量令牌）的计算时间远长于其他“冷门”专家，形成训练瓶颈（尾部延迟），极大降低了整体训练效率。图1(a)展示了训练Mixtral 8x7B模型时，专家负载在迭代过程中持续且剧烈波动的现象，图1(b)显示由此导致的通信延迟占比从不足10%激增至40%以上。
现有解决负载平衡的系统级方法主要分为专家复制（如FasterMoE, Prophet）和专家重定位（如SmartMoE）两类，或二者结合（如FlexMoE）。然而，这些方法存在根本性局限：专家复制会引入额外的梯度同步通信开销，且复制本身可能不均衡；专家重定位则涉及迁移专家参数和优化器状态，通信量巨大（通常达专家参数大小的6倍）。这些“重布局”操作本身带来了显著的额外开销，迫使现有方法必须在负载平衡效果与重布局开销之间进行权衡，例如限制重布局频率或惩罚高开销策略，从而无法对快速变化的负载做出及时、最优的调整。
因此，本研究旨在从系统层面出发，提出一种新颖的解决方案，其核心目标是：在不牺牲模型训练质量的前提下，通过一种灵活且开销极低的专家重布局机制，实现对动态路由负载的快速自适应平衡，从而显著提升大规模MoE模型的训练效率。
三、 研究设计与详细工作流程
本研究提出并实现了一个名为LAER-MOE的全新高效MoE训练框架。其核心创新在于一个名为全分片专家并行（Fully Sharded Expert Parallelism， FSEP）的新型并行范式，以及一个与之协同工作的智能负载平衡规划器。整个系统的工作流程可以概括为：规划器实时监控训练过程中的专家负载，动态生成优化的专家重布局策略和令牌路由方案；执行器则基于FSEP范式，在训练迭代中无缝实施这些策略，同时通过精细的通信调度隐藏重布局开销。
1. 核心创新一：全分片专家并行（FSEP）范式
FSEP的提出灵感来源于将全分片数据并行（Fully Sharded Data Parallelism， FSDP）的思想引入专家并行（EP）。与传统EP将完整专家固定分配到不同设备不同，FSEP对每个专家参数进行完全分片。假设有N个设备和E个专家，每个设备不存储任何完整的专家，而是存储所有E个专家的各一个分片（每个分片大小为原专家参数的1/N）。在训练的前向和后向传播过程中，每个设备根据当前迭代的布局策略，通过All-to-All通信从其他设备收集所需完整专家（数量C ≤ E）的参数块，临时重构出这些专家的完整参数进行计算。计算完成后，梯度通过All-to-All通信进行归约和重新分片。
FSEP的关键操作流程： * 分片（Shard）：在模型初始化阶段，将所有专家参数展平并拼接，然后均匀划分为N个块，每个设备保存其中一块。这类似于FSDP，但分片单位是跨所有专家的聚合参数。 * 解分片（Unshard）：在前向/后向计算前，根据规划器提供的布局（指定本设备需要计算哪C个专家），通过All-to-All通信，从其他设备获取构成这C个完整专家所需的所有参数分片。这是一个定制化的All-to-All操作，而非FSDP中的All-Gather。 * 再分片（Reshard）：在后向计算后，对完整的专家梯度进行分片，并通过All-to-All通信将分片发送到对应设备进行归约，同步更新各设备持有的参数分片。
FSEP的核心优势在于它将专家重布局过程与FSDP固有的参数预取和梯度归约-分散通信深度融合。具体来说，为下一层预取专家参数时，可以根据新的布局策略来获取，这使得重布局的通信成本被完全隐藏在计算时间内（见图5的通信优化调度）。因此，LAER-MOE能够在每个训练迭代中无额外开销地动态调整专家布局，为实现极致的负载平衡提供了前所未有的灵活性。
2. 核心创新二：智能负载平衡规划器
规划器的目标是在每个训练迭代中，联合优化两个决策变量：专家重布局策略（A，一个N×E的0/1矩阵，A_ij=1表示在设备i上恢复专家j）和令牌路由策略（S，一个N×E×N的张量，S_ijk表示设备i上路由到专家j且需要发送到设备k的令牌数）。其优化目标是最小化总时间T，包括通信时间T_comm和计算时间T_comp，并受限于每个设备只能恢复C个专家，且所有令牌必须被正确路由。
由于这是一个复杂的非线性整数规划问题，直接求解耗时且无法满足训练实时性要求。因此，研究团队设计了一个高效的启发式贪心算法，将问题分解为两个异步-同步协作的组件： * 异步专家布局调谐器：运行在CPU上，基于历史路由信息，为下一个训练迭代生成专家重布局策略。它首先通过一个优先级队列方法或均匀分配方法，确定每个专家应有多少个副本；然后，采用一个拓扑感知的贪心算法（算法1），在考虑节点内副本均衡和设备负载最小化的原则下，为每个专家副本分配合适的设备位置。为了找到更优解，算法会探索多种副本分配方案（包括比例分配、均匀分配及随机扰动变体），并评估每种方案的成本，最终选择最优布局。 * 同步令牌分发器：运行在GPU上，在当前训练迭代中，根据实时路由结果和上述调谐器提供的布局策略，快速决定每个令牌应被发送到目标专家的哪个具体副本上进行计算。其设计的核心原则是：1) 拓扑感知：优先将令牌路由到同一节点内的专家副本，以减少跨节点通信；2) 轻量级：无需全局协调，仅基于全局布局信息在本地进行均匀分配决策，确保极低的运行时开销（表3显示其耗时仅占总时间的不到0.1%）。
3. 系统整体工作流程与实现
如图7所示，LAER-MOE的工作流程是流水线化的：当计算当前MoE层时，其路由信息被发送给CPU侧的布局调谐器，用于计算下一迭代该层的布局策略；同时，GPU侧的令牌分发器根据当前布局快速路由令牌以进行All-to-All分发通信。在计算当前层专家MLP的同时，系统会根据已计算好的下一层布局策略，预取下一层专家的参数，实现了计算与重布局通信的重叠。
研究团队在PyTorch上高效实现了LAER-MOE系统，并进行了多项深度优化： * 异构并行策略与细粒度重计算：支持对MoE层（使用FSEP）和非MoE层（如注意力层，可使用FSDP或张量并行）采用不同的并行策略。并允许对注意力层和MoE层的专家计算部分进行细粒度的激活重计算，避免引入额外的通信开销。 * 定制的All-to-All内核：开发了自定义CUDA通信内核，直接在分片和未分片的参数之间进行All-to-All操作，避免了PyTorch原生接口带来的额外内存缓冲区和内存重排开销。 * 主机端优化：将主机-设备间的数据转移改为异步操作，并使用专用CUDA流进行管理，避免了CPU阻塞，提高了GPU利用率。同时，使用Triton内核实现令牌重排，避免了可能造成阻塞的PyTorch操作。
四、 主要实验结果与分析
研究团队在由4个节点（每节点8块NVIDIA A100 80GB GPU）组成的集群上进行了全面的实验评估。测试模型包括Mixtral-8x7B、Mixtral-8x22B以及结构改动的Qwen-8x7B，并在两种配置（E8K2：8专家Top-2；E16K4：16专家Top-4）下进行测试。对比基线系统包括：1) 最先进的分布式训练框架Megatron（支持异构专家并行）；2) 基于PyTorch FSDP扩展的FSDP+EP基线；3) 复现的当前最先进的负载平衡策略FlexMoE（与FSEP结合进行比较）。
1. 端到端性能 如图8所示，LAER-MOE在所有测试场景下均 consistently outperformed 其他方法。与Megatron相比，最高取得了1.69倍的加速；与FSDP+EP相比，最高取得了1.50倍的加速；与FlexMoE相比，最高取得了1.39倍的加速（平均提升1.20倍）。实验结果分析表明： * FSDP+EP在E8K2系列模型上表现更好，因为它通过完全分片节省了内存，允许使用更大的微批次。Megatron在E16K4系列模型上更优，因为参数较少使其可以使用更高效的张量并行配置。但两者都受困于专家负载不平衡导致的尾部延迟。 * FlexMoE在专家数量少（E8K2）时表现良好，但在专家空间更大（E16K4）时，其迭代搜索难以找到最优解，且其策略会惩罚布局变更，限制了灵活性。 * LAER-MOE通过FSEP实现了每迭代无开销的动态布局调整，其规划器能提供更灵活、更及时的负载平衡方案，因此在所有配置下均取得最佳性能。
2. 收敛性验证 如图9所示，在Mixtral-8x7B E8K2模型上的收敛实验表明，当使用相同的辅助损失权重（1e-4）时，LAER-MOE与Megatron的收敛曲线几乎完全一致（相对误差 < 1e-3），证明了FSEP范式不会损失训练精度。同时，LAER-MOE能够在低辅助损失权重下实现快速训练，获得了最佳的收敛速度（综合考虑迭代速度和所需步数）。
3. 案例分析 对Mixtral-8x7B的深入案例分析（图10）揭示了性能提升的来源： * 时间分解（图10a）：LAER-MOE的加速主要来自于All-to-All通信时间的大幅减少。FSDP+EP由于负载不平衡，通信时间占比高达40%。LAER-MOE通过负载平衡将通信占比降至20%以下，通信速度相比基线提升了最高2.68倍，而专家计算和其他操作时间相近。 * 负载平衡效果（图10b）：通过测量每层每设备处理的最大令牌数（灰色虚线为理想平衡线），LAER-MOE在所有情况下都最接近完美平衡。FlexMoE在E16K4配置下因调整策略受限而出现次优平衡，而LAER-MOE得益于每个迭代的全局优化和更多的每设备专家数（C=4），实现了近乎完美的负载均衡。
4. 规划器性能与消融实验 * 规划器效率：令牌分发器的耗时可忽略不计（<0.1%）。专家布局求解器的时间复杂度为O(|ε|N²C)，即使在模拟的1024 GPU规模下，其求解时间也低于单个Transformer层的平均计算时间，且可通过多进程并行进一步加速，证明其不会成为系统瓶颈（图11）。 * 消融研究（图12）：验证了关键组件的有效性。仅使用单一专家副本分配方案（“pq”或“even”）无法在所有路由场景下都有效，而结合多种方案的求解器能提供稳健性能。此外，移除通信优化（图5所示）会导致计算-通信重叠效果变差，影响端到端效率。
五、 研究结论与价值
本研究成功提出了LAER-MOE，一个用于高效分布式MoE训练的创新系统。其核心贡献在于： 1. 提出了全新的全分片专家并行（FSEP）范式：通过将每个专家参数完全分片，并将重布局通信与训练固有通信（参数预取、梯度同步）深度融合，首次实现了在训练过程中灵活、无额外开销的专家动态重布局。 2. 设计了一个智能负载平衡规划器：通过异步布局调谐器和同步令牌分发器的协同，能够动态、实时地联合优化专家布局和令牌路由，有效应对快速变化的负载分布。 3. 实现了显著的性能提升：实验证明，LAER-MOE能带来高达1.69倍的端到端训练加速，且不损害模型收敛性。
该研究的科学价值在于为分布式MoE训练中的负载不平衡这一核心挑战提供了全新的系统级解决思路，打破了重布局开销与平衡效果之间的传统权衡，推动了高效稀疏模型训练系统设计理论的发展。其应用价值则体现在为训练超大规模MoE语言模型提供了切实可行的高效工具，使研究人员能够更自由地探索低辅助损失（可能带来更好模型质量）的算法，而无需担心系统效率的下降，从而促进更强大AI模型的开发。
六、 研究亮点
七、 其他有价值的讨论
论文还就一些潜在问题进行了深入讨论： * 平衡场景下的性能：在负载完全平衡的理想情况下，LAER-MOE的通信量与FSDP+EP相当，性能相近。这强调了其设计初衷是解决不平衡问题，将系统效率从算法约束中解耦。 * 全局峰值内存：通过分析各层最大令牌数，指出由于不同层的“热门”专家不同，FSEP对全局峰值内存的影响并不显著，不影响其优化目标（平衡每层计算负载）。 * 大规模集群的可扩展性：尽管受资源所限未进行超大规模实测，但通过理论分析和基于真实路由轨迹的模拟表明，随着集群规模扩大，更大的微批次可以补偿带宽下降，维持计算-通信重叠；同时，LAER-MOE可与流水线并行等混合并行策略兼容，以应对跨机柜带宽限制。模拟显示从8GPU扩展到128GPU，其重布局算法带来的加速比保持稳定。