最少负载专家并行：面向非均衡混合专家模型的负载均衡

动态重分配过载设备的token和专家权重至空闲设备，实现5倍加速

前置知识

混合专家模型（Mixture-of-Experts, MoE）

MoE是一种稀疏激活的神经网络架构，由一个路由器（Router）和N个专家网络（通常是FFN层）组成。对于每个输入token，路由器通过计算softmax概率选择top-K个专家来处理该token，最终输出是被选中专家输出的加权和。形式化地，给定token的隐藏表示 $u \in \mathbb{R}^D$，路由器计算 $s_i = \text{softmax}_i(u^T W_r)$ 选择得分最高的K个专家，输出 $h = \sum_{i} g_i \text{FFN}_i(u)$。这种设计使得模型可以在保持每个token的计算量（激活参数）恒定的同时扩展模型总参数量。

MoE是本文研究的核心对象，理解其路由机制和专家选择方式是理解负载不均衡问题的前提

专家并行（Expert Parallelism, EP）

EP是专门为MoE模型设计的分布式并行策略，将N个专家分布在P个GPU设备上，每个设备仅存储M=N/P个本地专家。计算时，所有设备先通过All-to-All通信将token发送到对应专家所在的设备，每个设备用本地专家处理收到的token，再通过第二次All-to-All将结果返回原设备。这种策略比张量并行或流水线并行更适合MoE模型，因为它能更高效地利用内存和通信带宽。

EP是本文要改进的基础方法，理解其工作原理和假设（负载均衡）是理解LLEP动机的关键

All-to-All通信

All-to-All是一种集合通信操作，每个设备向所有其他设备发送不同的数据，同时也从所有其他设备接收数据。在EP中，第一次All-to-All用于dispatch阶段（将token路由到专家所在设备），第二次用于combine阶段（将专家输出返回原设备）。通信量取决于路由到每个专家的token数量，在负载不均衡时可能导致严重的通信瓶颈。

All-to-All是EP的核心通信原语，其开销直接影响LLEP中token重分配的代价计算

Grouped-GEMM

Grouped General Matrix Multiplication是指对多个不同大小的矩阵乘法批处理执行。在MoE中，每个专家的GEMM操作 $B_i W_i$ 的输入大小 $B_i$（token数量）各不相同。可以用cuBLAS循环执行多个独立GEMM，或用Triton编写融合的Grouped-GEMM核函数。论文实验表明，尽管cuBLAS需要多次kernel启动，但由于其高度硬件优化，仍比Triton融合版本更快。执行少量大GEMM比执行大量小GEMM效率更高。

理解GEMM执行效率是理解为什么负载均衡能带来加速的计算基础

路由不均衡（Routing Imbalance）

指MoE模型中token被不均匀地路由到少数专家的现象。即使经过良好预训练的MoE模型也会表现出这种特性，因为部分专家会专业化处理特定领域知识。例如论文观察到gpt-oss-20b中E11专家在数学数据上接收高达20%的token负载（均衡情况下应为~3%）。这种不均衡在领域特定的后训练或推理场景中尤为严重，因为只有相关领域的专家被频繁激活。

路由不均衡是本文要解决的核心问题，理解其自然性和必然性是理解LLEP价值的基础

研究动机

混合专家模型（MoE）在预训练时通常使用辅助损失或随机偏置项来强制实现均衡路由，但即便如此，训练良好的MoE模型仍然表现出显著的路由不均衡。论文通过实验观察到，在gpt-oss-20b模型（32个专家，8路EP）上处理数学对话数据时，E11专家接收的token负载高达20%，而均衡情况下仅为~3%；GPU-0的负载达到30-35%，而均衡基准仅为~12.5%。更重要的是，这种不均衡是自然且可取的——经过大规模预训练后，部分专家专业化处理特定领域知识，在领域特定的后训练或推理时，这些专家自然被更频繁地激活。标准EP假设负载均衡，在极端不均衡场景下（95%的token集中到1个专家），延迟可能增加4.6倍，峰值内存使用可能增长4倍，甚至导致GPU内存溢出（OOM）崩溃。现有的缓解方法如降低batch size会严重降低吞吐量，使用冗余专家副本会增加内存消耗且仅适用于推理场景。

本文的目标是本文的目标是在系统层面而非模型层面解决MoE的负载不均衡问题，设计一种新的专家并行算法，能够在不改变模型数学行为的前提下，动态地将过载设备上的token和对应专家权重转移到空闲设备上，使得所有设备在最小的集体延迟内完成工作负载，同时保持每GPU的最小峰值内存使用。具体目标包括：在极端不均衡场景下实现显著加速（目标5倍以上）、稳定峰值内存使用避免OOM、支持反向传播以应用于训练场景、以及在实际预训练模型上验证端到端吞吐量提升。

与已有工作不同的是，现有方法存在明显局限：辅助损失和路由偏置会改变模型行为，破坏已学习的专家专业化模式；降低batch size直接牺牲吞吐量；链式梯度检查点处理效率低下且受硬内存上限约束；DeepSeek-V3的EP Load Balancer（EPLB）基于延迟路由统计数据复制高负载专家，增加内存开销且不适用于微调，在极端情况下仍会OOM。本文的独特切入角度是：不试图纠正路由不均衡（这是模型的自然属性），而是在系统层面拥抱这种不均衡，通过动态重分配计算任务来利用设备间的空闲资源。LLEP是一种精确的MoE计算算法，不改变模型行为，同时支持前向和反向传播，这使其既适用于推理也适用于训练。

核心方法

LLEP的核心直觉是：当某个GPU上的专家接收到过多token时，与其让这个设备独自处理导致高延迟甚至OOM，不如将多余的token连同对应的专家权重一起转移到有空闲计算能力的设备上。技术路线分为三个阶段：首先，收集全局路由信息，计算每个专家的负载分布；其次，通过最少负载分配算法（LLA）制定任务分配计划，决定每个GPU处理哪些专家的哪些部分token；最后，按照分配计划执行dispatch-compute-combine操作，包括必要的权重P2P传输。整个过程是动态的，每个batch都可能有不同的分配方案。关键的优化点包括：自适应阈值λ判断是否需要启用LLEP（均衡情况下直接用标准EP避免额外开销）、容量因子α控制每GPU的最大token容量、最小GEMM大小m避免过小的计算任务。

LLEP的核心创新在于它同时转移token和专家权重，而非仅仅转移token。这与现有方法有本质区别：标准EP只做token的All-to-All路由，专家权重固定在原设备上；LLEP允许任何设备临时加载其他设备的专家权重来处理溢出的token。这种设计的关键优势是它保持了精确的MoE数学计算——不改变任何专家的权重，不修改路由决策，不引入近似，因此模型行为完全不变。另一个重要创新是负载感知的溢出决策：只有当转移token的代价低于本地处理的代价时才触发溢出，这通过考虑计算延迟、内存约束和通信开销的综合模型来实现。LLEP还支持反向传播，溢出专家的梯度会返回到其原始设备并正确累积。

方法步骤详情

LLEP的执行分为以下步骤：第一步，全局负载收集——每个设备通过All-reduce汇总所有专家的全局负载向量 $l \in \mathbb{R}^N$，计算负载不均衡比率 $\max(l)/\mean(l)$。第二步，自适应判断——如果负载比率低于阈值λ（如1.3），则认为路由基本均衡，直接调用标准EP算法（算法1），避免额外开销。第三步，最少负载分配（LLA，算法2）——将专家按负载降序排列，从最大负载专家开始，依次确定每个专家的token分配给哪些GPU。对于每个专家，首先检查其原生GPU能否处理全部token（不超过容量 $m_\alpha = \alpha \times \frac{1}{P}\sum_{i=0}^{N-1} \hat{l}_i$）。如果能，全部分配给原生GPU；如果不能，原生GPU取尽量多的token，剩余部分通过最少负载分配溢出（LLAS，算法3）循环分配给其他负载最低的GPU，直到所有token分配完毕。第四步，权重传输计划——根据token分配结果，构建权重P2P传输计划W，如果专家i的token被溢出到GPU q，则需要从原生GPU p传输权重 $W_i$ 到GPU q。第五步，执行dispatch——按照分配计划将token通过All-to-All发送到对应的计算设备，同时通过P2P传输专家权重。第六步，Grouped-GEMM计算——每个设备不仅计算本地专家的GEMM，还计算被分配的外部专家的GEMM。第七步，combine——通过反向All-to-All将计算结果返回原设备，反向排序和重组。第八步，反向传播支持——在训练时，溢出专家权重的梯度返回其原始设备并累积。

技术新颖性

LLEP的技术新颖性体现在多个层面。首先，它是第一个同时进行token和权重动态重分配的EP算法，区别于仅做token路由的标准EP和仅复制权重的EPLB。其次，LLEP是一种精确计算算法（exact MoE computation algorithm），不改变模型行为，这与引入辅助损失或修改路由偏置的方法有本质区别。第三，LLEP的负载感知溢出决策是成本模型驱动的——考虑了GEMM计算效率（大batch比小batch更高效）、通信开销（权重传输vs计算节省）、内存约束（避免OOM）的综合权衡。第四，自适应阈值λ机制使LLEP在均衡场景下自动退化为标准EP，避免不必要的开销。第五，LLEP支持完整的反向传播，使其适用于训练场景，而不仅是推理。第六，理论分析揭示了GEMM效率与专家数量的非线性关系——执行少量大GEMM比大量小GEMM效率高得多，这为LLEP的设计提供了理论基础。

**Figure 2**: Comparison of standard Expert Parallelism and LLEP

**Figure 8**: Grouped-GEMM benchmark: execution time vs. number of experts

实验结果

论文在多个层面验证了LLEP的有效性。在受控实验中，使用8个H200 GPU测试三种MoE架构（gpt-oss-120b: 128专家，DeepSeek-V3: 256专家，Kimi-K2: 384专家），在极端不均衡场景（95%的token集中到1个专家）下，LLEP实现了高达6.11倍的加速和5倍的峰值内存节省。在均衡场景下，由于自适应阈值λ的存在，LLEP与标准EP性能相当。在端到端全模型测试中，使用真实的gpt-oss-20b和gpt-oss-120b模型处理Megatron-Math数据集，LLEP分别实现了2.2倍和1.88倍的吞吐量提升，且GPU数量越多相对加速越大。在训练场景中，使用gpt-oss-20b进行全参数SFT训练（Zero-3+CPU offloading），LLEP实现了1.25倍的训练加速，同时在AIME'25上达到与标准EP相当的准确率。消融研究揭示了关键的超参数影响：batch size越大加速比越高（大batch饱和GPU容量使均衡分配的收益超过通信开销）；α因子越小加速越大（更早触发溢出）；隐藏维度越大加速越大（GEMM效率提升超过通信开销）；专家数量越多加速越大。

**Figure 1**: LLEP vs. standard expert parallelism (EP)

**Figure 4**: Performance comparison of LLEP vs. standard EP across three MoE architectures

**Figure 5**: Performance (accuracy on AIME'25) vs. wall-time for EP and LLEP

**Figure 6**: Speedup of LLEP over standard EP with 4 imbalanced experts

**Figure 9**: Speedup of LLEP over standard EP as a function of number of experts

查看结构化数据

任务	指标	本文	基线	提升
MoE层前向传播（gpt-oss-120b，95%不均衡到1专家）	加速比	6.11×	标准EP (1×)	6.11倍加速
MoE层峰值内存（Kimi-K2，95%不均衡）	峰值内存使用	~20GB（稳定）	标准EP ~100GB	5倍内存节省
gpt-oss-20b全模型吞吐量（8 GPU）	吞吐量加速	2.2×	标准EP (1×)	2.2倍加速
gpt-oss-120b全模型吞吐量（32 GPU）	吞吐量加速	1.88×	标准EP (1×)	1.88倍加速
gpt-oss-20b SFT训练（AIME'25准确率）	训练加速	1.25×	标准EP (1×)	1.25倍加速，准确率相当
cuBLAS vs Triton Grouped-GEMM（128专家）	执行时间	cuBLAS ~14ms	Triton+TMA ~17ms	cuBLAS更快，说明少量大GEMM比融合小GEMM更高效

局限与改进

论文的局限性主要包括：第一，当前实现使用Python实现LLA算法和标准Torch NCCL进行通信，存在优化空间，实际性能可能受限于Python开销和通用通信库效率。第二，超参数α、m、λ需要针对每个使用场景调优，论文建议针对N、P、Bp、K、D、H、模型大小和物理系统配置进行调参，这增加了使用复杂度。第三，权重P2P传输在极端不均衡时可能成为瓶颈，特别是当多个专家需要同时传输时。第四，论文主要在单节点8 GPU环境下验证，多节点场景下的跨节点通信开销可能显著影响性能，论文承认需要修改LLEP优先溢出到节点内设备。第五，论文未充分讨论在动态工作负载（如在线服务）场景下的适用性，LLEP的batch级别重分配可能不适合延迟敏感的实时推理。第六，LLEP的理论分析主要关注前向传播，反向传播的额外开销（梯度回传）未被充分量化。

独立分析的弱点

第一，Python实现的LLA算法存在性能瓶颈——排序、循环、贪心分配等操作在Python中效率低下，特别是当专家数量达到256或384时，LLA的执行时间可能不可忽略。改进方向是将LLA用C++/Triton重写，或设计并行化的分配算法。第二，标准NCCL All-to-All在token重分布不规则时效率不高——论文建议可以使用修改版DeepEP或直接在未排序张量上执行All-to-All来避免内存密集的index select操作。第三，权重传输未与计算重叠——当前实现是先传输权重再计算，理想情况下应将P2P权重传输隐藏在Grouped-GEMM计算背后。第四，多节点优化缺失——当前未区分节点内和节点间通信成本，应优先将溢出分配给节点内设备以减少跨节点通信。第五，自适应阈值λ的粒度较粗——目前是全局二值判断（均衡/不均衡），可以考虑逐层或逐专家的更细粒度判断。

未来方向

作者提出的未来方向包括：使用低级C++/Triton内核优化通信操作、采用修改版DeepEP实现更高效的All-to-All、将通信与计算重叠或隐藏在Grouped-GEMM背后、以及针对多节点场景修改LLEP优先溢出到节点内设备。基于本文成果可延伸的方向包括：第一，将LLEP与其他并行策略（张量并行、数据并行、流水线并行）结合，设计混合并行方案；第二，将LLEP应用于MoE模型的长序列推理场景，其中token分布可能更加不均衡；第三，设计硬件感知的自动调优框架，根据GPU型号、互联拓扑和模型配置自动选择最优α、m、λ参数；第四，探索LLEP在训练中的动态batch size策略——在不均衡程度高时自动减小batch size以减少溢出开销；第五，将LLEP的思想扩展到其他稀疏激活架构，如Switch Transformer或Mixture of Attention。

复现评估

论文提供了完整的开源代码（github.com/SalesforceAIResearch/LeastLoadedEP），这大大降低了复现门槛。核心算法（LLA、LLAS、LLEP dispatch-combine）的伪代码在论文中清晰给出（算法1-4），实现细节充分。实验使用公开可获得的H200 GPU，MoE配置参数明确（专家数量、隐藏维度、batch size等）。复现所需的主要资源是8张H200 GPU（或同等性能的GPU）和对应的MoE模型权重。论文的消融研究覆盖了关键超参数的影响，为调参提供了指导。主要的复现挑战在于：需要正确配置MoE模型的路由机制、实现精确的负载不均衡模拟、以及处理All-to-All通信的正确性。总体而言，论文的可复现性较高。

最少负载专家并行：面向非均衡混合专家模型的负载均衡 Least-Loaded Expert Parallelism: Load Balancing An Imbalanced Mixture-of-Experts