Least-Loaded Expert Parallelism: Load Balancing an Imbalanced Mixture-of-Experts

标题：最低负载专家并行：对不平衡的混合专家模型进行负载均衡
作者/机构：Xuan-Phi Nguyen, Shrey Pandit, Austin Xu, Caiming Xiong, Shafiq Joty | Salesforce AI Research

A1 主要贡献

本文研究的核心问题是，在混合专家（Mixture-of-Experts, MoE）模型中普遍存在的专家路由不平衡现象。即使是经过良好预训练的MoE模型，在后训练或推理阶段也表现出显著的路由不平衡，即一小部分专家处理了绝大多数的token。这种现象本身是自然且有益的，因为它允许模型将领域特定的知识集中在部分专家中。然而，现有的专家并行（Expert Parallelism, EP）方法假设路由是均衡的，导致在不平衡场景下，少数设备会因接收过多token而过载，引发高计算延迟、内存峰值过高，甚至内存不足（OOM）的故障。特别是在后训练和推理阶段，通常不适用通过辅助损失等方式来强制实现负载均衡。

为解决此问题，本文提出了最低负载专家并行（Least-Loaded Expert Parallelism, LLEP），这是一种新颖的专家并行算法，其核心创新点在于：
1. 动态负载均衡：LLEP能够动态地将过载设备上多余的token及其对应的专家权重，重新路由到利用率不足的设备上。当某个GPU上的专家组负载超过预设阈值时，LLEP仅将token分配至该GPU的容量上限，然后将剩余的token和相应的专家权重转移到负载最低的设备，让它们分担超出的工作量。
2. 系统级解决方案：与在模型层面强制均衡路由的方法不同，LLEP在系统层面接纳并高效处理不平衡现象，它不改变模型的路由行为，保证了MoE计算的精确性，从而保留了模型预训练学到的专业化知识。
3. 全面的成本感知设计：LLEP的负载均衡算法不仅考虑计算负载，还同时兼顾了每个GPU的内存分配和跨GPU的通信开销。只有当转移token的成本低于在本地处理它们的成本时，才会触发转移。
4. 支持训练与推理：LLEP提供了反向传播支持，使其能够在训练循环的每次迭代中动态应用，同样也适用于推理过程。

根据论文的实验结果，LLEP相较于标准EP展现了显著的性能优势。在MoE层级别，极端不平衡场景下，LLEP实现了高达5倍的加速和4倍的峰值内存使用量降低，同时保持了内存使用的稳定性。在完整的模型（如gpt-oss-120b）上进行端到端测试时，LLEP也实现了高达1.9倍的吞吐量提升。

A3 背景知识

2.1 混合专家模型（MIXTURE-OF-EXPERTS）

MoE作为横向扩展LLM的标准。自Lepikhin等人【【6】GShard: Scaling giant models with conditional computation and automatic sharding，2020，arXiv】的开创性工作以来，MoE模型已成为横向扩展大语言模型（LLM）的实际标准，著名例子包括DeepSeek-V3【【7】Deepseek-v3 technical report，2024，arXiv】、gpt-oss【【1】gpt-oss-120b & gpt-oss-20b model card，2025，arXiv】和Kimi-K2【【17】Kimi k2: Open agentic intelligence，2025，arXiv】。MoE前馈层使模型能够跨越不同的专家权重矩阵学习广泛的知识，同时允许单个token由一个稀疏的专家子集处理，从而提高效率和可扩展性。尽管具体实现各不相同，MoE架构依赖于一个核心思想：一个路由器层（router layer）选择top-K个专家来路由每个token。token随后由每个被激活的专家处理，这些专家通常构建为前馈（FFN）层。形式上，给定一个token x的隐藏表示$u \in R^D$，MoE模块有N个专家${FFN_0, FFN_1, \ldots, FFN_{N-1}}$和一个路由器层Router，它选择top-K个专家来路由x。为简洁起见，我们定义$FFN_i(u) = u^T W_i$，其中$W_i \in R^{D \times H}$是专家i的权重矩阵，$W_r \in R^{D \times N}$是路由器层的权重矩阵。那么，MoE的输出h为：

$$\begin{aligned} \boldsymbol{h}=\sum_{i=0}^{N-1} g_{i} \mathrm{FFN}_{i}(\boldsymbol{u}), \quad \text { where } \quad g_{i}= \begin{cases}s_{i}, & \text { if } s_{i} \in \operatorname{top-} K\left(\left\{s_{j} \mid 0 \leq j \leq N-1\right\}, K\right) \\ 0, & \text { otherwise }\end{cases} \end{aligned}$$ $$s_{i}=\operatorname{softmax}_{i}\left(\boldsymbol{u}^{T} \boldsymbol{W}_{r}\right).$$

公式解释。在上述公式中，$g_i$是专家i的门控亲和度分数，只有得分最高的top-k个专家被选中对输出做出贡献。

实现方式。处理一个token所需的通用矩阵乘法（GEMM）数量与激活的专家数量成正比。朴素的实现（例如，每个token按顺序遍历每个专家）可能极其低效。一种提高性能的简单方法是通过重新索引来高效地构建每个专家的token批次。假设对于一批tokens $B \in R^{B \times D}$，有$G \leq N$个专家被激活，即接收到路由的tokens。不失一般性，假设这G个被激活的专家是专家$0, \ldots, G-1$。然后，我们可以形成一个重新索引的$B' = [B_0, B_1, \ldots, B_{G-1}]$，其中$B_i \in R^{B_i \times D}$包含所有路由到专家i的tokens。举个例子，如果B包含四个tokens $B = [a, b, c, d]$，分别路由到专家$[2, 6, 2, 1]$，那么我们可以形成$B' = [d, a, c, b]$，其中$B_1 = [d]$，$B_2 = [a, c]$，$B_6 = [b]$。在这种重新索引的方法下，每个MoE层计算G个GEMM操作$B_i W_i$，对应专家$i=0, 1, \ldots, G-1$。

2.2 专家并行（EXPERT PARALLELISM）

EP的原理与优势。为了在多个GPU上训练大型MoE模型，专家并行（EP）通常优于张量或流水线并行【【14】Megatron-lm: Training multi-billion parameter language models using model parallelism，2019，arXiv】，因为它能更有效地利用内存和通信带宽。在EP中，专家被分布在多个GPU上，每个设备只托管一部分本地专家。计算过程中，输入tokens首先由一个路由器层处理，以产生全局路由索引和相应的路由权重（亲和度分数）。然后，产生的（输入，索引，权重）元组被排序和重新索引，并随后分派到托管所选专家的GPU上，这是由路由索引决定的。

dispatch-combine流程。路由过程通常使用dispatch-combine程序进行。算法1形式化了该程序的高效单设备实现，而图2a则直观地展示了EP在不平衡路由场景下的情况。具体来说，在分派（dispatch）阶段，每个设备将其本地输入tokens发送到其他专家所在的设备，并从其他设备接收分配给其本地专家的tokens。这种数据交换模式被称为All-to-All通信操作【【11】Bruck algorithm performance analysis for multi-gpu all-to-all communication，2024，HPCA】，【【9】Optimizing distributed ml communication with fused computation-collective operations，2024，SC24】。在专家FFN计算完成后，输出在合并（combine）阶段被聚合。此时，所有专家的输出通过另一次All-to-All操作被送回其原始设备。除了标准的基于NCCL的集合通信，EP还可以使用专门的库在核函数（kernel）层面更高效地实现，例如DeepEP【【7】Deepseek-v3 technical report，2024，arXiv】和Triton-Distributed【【21】Triton-distributed: Programming overlapping kernels on distributed ai systems with the triton compiler，2025，arXiv】。

A3 关键观察与设计原则

3.1 不平衡路由分析

不平衡路由的实证观察。即便是大型且高性能的MoE模型【【7】Deepseek-v3 technical report，2024，arXiv】，【【17】Kimi k2: Open agentic intelligence，2025，arXiv】也被证明存在专家路由不平衡现象，即token被路由到一小部分专家。我们通过在8路EP（每个GPU托管4个专家）下，用多批次数据运行gpt-oss-20b【【1】gpt-oss-120b & gpt-oss-20b model card，2025，arXiv】来研究token路由的具体动态。为了保持数据分布的熟悉性，我们用对话数据喂给模型，其中问题来自DeepScaleR【【8】Deepscaler: Surpassing o1-preview with a 1.5b model by scaling rl，2025，Notion Blog】，回答的思维链由gpt-oss-20b自身生成。在图3a中，我们观察到token持续地被路由到特定的专家位置，其中位置E11占主导。这意味着在24个MoE层中，至少有一个第11个专家在不同数据批次中持续接收主要负载。尽管如此，E11所在的GPU（gpu-2）并非负载最高的；托管专家E0-E3的gpu-0在所有GPU中负载最高。这表明在极端不平衡的路由下，某些GPU设备可能处理了压倒性数量的token。有趣的是，虽然E11通常承载最多的token，但某些批次会导致更多token路由到其他专家。也就是说，不平衡的程度是按批次变化的。

不平衡路由的合理性。人们可能倾向于将不平衡的MoE路由归因于预训练的缺陷，例如数据倾斜或设计不佳的负载均衡损失【【22】Mixture-of-experts with expert choice routing，2022，NeurIPS】，【【13】Outrageously large neural networks: The sparsely-gated mixture-of-experts layer，2017，arXiv】。确实，如果训练不当，MoE模型可能会出现“专家坍塌”，即只有一小部分专家被激活，导致模型性能次优或较弱。然而，正如我们之前所示，即使是顶尖的MoE模型也表现出一定程度的不平衡路由，尽管比极端的专家坍塌要温和得多。我们认为轻微的不平衡是一个训练良好的MoE模型的自然属性，而不是追求完美的平衡路由。经过大规模预训练后，专家子集通常会在特定知识领域、任务或能力上形成专业化【【10】Demons in the detail: On implementing load balancing loss for training specialized mixture-of-expert models，2025，arXiv】，【【4】Communication-efficient moe fine-tuning with locality-aware expert placement，2025，ICDCS】，【【16】Mixture of weight-shared heterogeneous group attention experts for dynamic token-wise kv optimization，2025，arXiv】。因此，当MoE模型在特定领域（如数学）上进行微调或评估时，为该领域专业化的专家会被更频繁地激活，从而导致路由不平衡。同时，一些专家可能演变为广泛适用的、领域无关的“共享”专家，持续处理通用的语言模式（如语法）。这种现象在先前的工作中也已被观察并接受【【7】Deepseek-v3 technical report，2024，arXiv】。从这个角度看，强制执行平衡路由，例如通过辅助负载均衡损失【【3】Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity，2022，JMLR】或移动平均路由偏差【【7】Deepseek-v3 technical report，2024，arXiv】来改变模型行为，可能会破坏特定专家内部学到的这些专业化模式。我们不从模型层面纠正不平衡，而是接受它，并提出一个系统级的机制，用于在不平衡路由下最大化训练和推理的吞吐量。这既尊重了专家间的内在专业化，又减轻了不平衡带来的低效问题。

现有缓解不平衡的方法及其局限性。已有几种方法被提出来缓解EP下的路由不平衡。一个朴素的解决方案是减小批量大小，但这会严重降低吞吐量。另一个策略是使用链式梯度检查点来分小块处理token；然而，这种方法效率低下，并且仍然受限于硬性的内存上限。对于推理，Liu等人【【7】Deepseek-v3 technical report，2024，arXiv】提出了一种EP负载均衡器（EPLB），它根据时间延迟的路由统计数据在设备间复制重载专家。虽然在某些情况下有效，但该方法会产生额外的内存开销，且不适用于微调；此外，在极端路由不平衡下仍可能导致内存溢出（OOM）失败。相比之下，Huang等人【【5】Toward efficient inference for mixture of experts，2024，NeurIPS】建议为超额专家预留额外内存，这同样会产生CPU和GPU内存开销。

3.2 分布式延迟与内存分析

本地计算延迟分析。为了更深入地了解EP下MoE层的最坏情况成本模型，我们从整体上分析了延迟和峰值内存使用。我们首先考虑单个GPU设备上的本地计算。给定一个批次被路由到G个本地专家，MoE层执行G次GEMM操作，总延迟可以近似为：

$$T_{\mathrm{local}}=\sum_{i=0}^{G-1}(T_{\mathrm{overhead}}+B_{i}\times T_{B_{i},D,H})$$

其中$T_{overhead}$表示核函数启动延迟，$T_{B_i,D,H}$是每个token的计算时间，它取决于token数量$B_i$和模型维度D和H（在§2.1中定义）。$T_{B_i,D,H}$的效率直接受到GEMM核函数的实现、优化以及针对不同输入输出尺寸和配置的调优影响。通常，随着$B_i, D$和H的增加，GEMM变得更高效。例如，当D和H固定时，若$B_1 > B_2$，则$T_{B_1,D,H} < T_{B_2,D,H}$。因此，在固定的FLOPs下，执行少量大型GEMM比执行大量小型GEMM要高效得多。EP利用了这一特性，通过跨设备聚合token，从而减少了本地专家的数量G，并增加了每个专家的有效批量大小$B_i$。

GEMM实现方式的影响。多个$T_{overhead}$可以通过使用一个融合的grouped-GEMM核函数减少到一个，但这并不总是更快，因为单个针对硬件优化的GEMM核函数（如cuBLAS）在D和H较大时效率更高。图8显示，即使我们计算完全相同的FLOPs，随着专家数量的增加，执行时间也会增加，并且启动许多小的cuBLAS GEMM仍然比一个融合的Triton grouped-GEMM核函数要快。MoE层的峰值内存使用量近似定义为：

$$M_{\text{local}} = \sum_{i=0}^{G-1} (B_i \times D + D \times H + B_i \times H)$$

标准EP在不平衡场景下的问题。在标准专家并行下，$B_i$是所有EP设备路由到专家i的token总数。在最坏的情况下，$B_i$可能接近全局批量大小，导致所有token集中在一个设备上，而其他设备则处于空闲状态。这会导致过载设备的延迟和内存使用急剧增加，甚至导致内存溢出崩溃。图1a和1b展示了标准EP设置在不同不平衡场景下的减速和峰值内存使用情况。如图1a所示，当95%的token路由到单个专家时，EP可能比平衡基线慢4.6倍。至于峰值内存使用，EP的每GPU峰值内存使用可能增长高达4倍，这可能导致OOM错误。

A2 方法细节

4 最低负载专家并行（LEAST-LOADED EXPERT PARALLELISM, LLEP）

LLEP核心概念。我们详细解释我们提出的LLEP是如何工作的。从概念上讲，我们的方法会根据每个专家的负载提前检测全局路由的不平衡程度。如果不平衡度低于一个阈值λ，那么我们认为路由是平衡的，并继续执行标准的EP程序。否则，我们将执行最低负载分配算法（算法2），为每个GPU设备确定它需要为哪些专家计算GEMM，以及处理这些专家全局token的多大比例。如果该GPU本身没有驻留指定的专家，它将从托管该专家的GPU导入。这个分配过程会考虑权重和数据传输的开销成本，并与仅处理本地专家token的延迟和内存成本进行比较。算法2到4详细地形式化描述了我们的方法。

约束条件。我们的方法通过制定受某些约束的路由决策来工作。首先，算法2中的因子α决定了一个GPU可以处理的最大token容量，我们定义为$m_{\alpha} = \alpha \sum_{i=0}^{N-1} \hat{l}_i / P$个token。$m_{\alpha}$不一定是物理内存限制，而是一个GPU被视为过载的阈值。如果一个本地专家的负载超过$m_{\alpha}$，它会将超出的负载溢出到其他GPU。其次，m是GEMM要达到高效所需的最小token数。如果一个本地专家的负载超过了本地GPU的已占用容量，但超出的部分小于m，我们认为溢出是不值得的，而是强制本地GPU计算它，尽管超出了容量（见§3.2）。第三，不平衡比率阈值λ用于确定全局负载是否相对平衡，在这种情况下我们会切换回标准EP。原因是我们的方法采用了一种贪婪的最低负载分配（LLA）算法（算法2），它在平衡情况下会产生与标准EP相同的路由计划，但会带来微小的时间开销。如果不跳过这个不平衡比率检查，我们的方法在完全平衡的场景下会比标准EP稍慢。α、m和λ的最优值取决于N、P、Bp、K、D、H、整体模型大小以及物理系统配置。因此，我们建议为每个用例调整这些值。

算法详述。最低负载分配（LLA）算法（算法2）为每个专家确定哪些GPU处理其全局负载的哪些部分。首先，它按降序对专家负载进行排序。然后，它从负载最大的专家到负载最小的专家依次确定GPU的分配。对于每个专家，它首先确定其原生GPU（托管该专家权重的GPU）是否能处理该专家的所有token。如果可以，它将所有token分配给原生GPU。如果不能，它会将超出的token溢出到负载最低的可用GPU，直到达到容量阈值。如果仍有剩余的超额token，它将继续这个溢出循环（LLAS，算法3），直到所有token都被分配。一旦token路由计划最终确定，它也会相应地构建权重转移计划。例如，如果原生于GPU p的专家i的超额负载被溢出到GPU q，那么权重转移计划将包括一个从p到q的$W_i$权重转移操作。LLA算法确保每个GPU优先计算其本地专家的大部分（如果不是全部）负载，然后再接受外部专家的负载。这是为了最小化所需的权重转移数量。最终的LLEP算法（算法4）将根据从LLA获得的路由计划执行分派-计算-合并操作。具体来说，对于每个设备，除了为原生专家计算GEMM外，LLEP还将为分配给该设备的外部专家计算GEMM。与其他方法不同，LLEP支持正确的梯度传播。在反向传播过程中，溢出的专家权重的梯度会返回到它们的原生设备，并与它们各自的原生梯度累加。

实现与优化。在实验中，我们使用标准的Torch NCCL来实现All-to-All和点对点（P2P）操作。LLA算法是用纯Python实现的。虽然我们简单的LLEP实现已经显示出显著的速度提升和内存节省，但仍有进一步优化和减少开销的机会。例如，通信操作可以写成底层的C++/Triton核函数，或者使用修改版的DeepEP【【7】Deepseek-v3 technical report，2024，arXiv】。这种融合操作也可以直接在未排序的张量Bp和Gp上执行All-to-All，从而避免了内存密集型的索引选择操作（算法4）。通信可以与计算重叠，或者隐藏在Grouped-GEMM操作之后。对于多节点设置，我们可以进一步修改LLEP，使其优先将工作溢出到节点内的设备，以限制更高的节点间通信开销。

算法4 LLEP dispatch-combine：每个设备p（从0开始索引）的操作，EP world size为P，每个设备的专家数量为M = N/P
输入：Bp ∈ R^(Bp×K×D)，Gp ∈ R^(Bp×K)，Ip ∈ Z^(Bp×K)，Wi ∈ R^(D×H) 对于 i = pM, pM + 1, ..., (p + 1)M - 1
l ← 所有GPU上全局专家的负载总和
if max(l)/mean(l) < λ then
    调用标准EP算法1
    输出：来自标准EP的H'p
end if
I¯p ← sort(flatten(Ip, dims=[Bp, K]))
B¯p ← index_select(flatten(Bp, dims=[Bp, K]), I¯p)
G¯p ← index_select(flatten(Gp, dims=[Bp, K]), I¯p)
// 构建路由计划和权重转移计划
A, W ← LLA(l, M)
{Bi|i ∈ [0, ..., P-1]} ← 从A构建B¯p的块
{Gi|i ∈ [0, ..., P-1]} ← 从A构建G¯p的块
// S是分配给此设备的外部专家的ID
{Bˆi|i ∈ [pM, (p + 1)M − 1] ∪ S} ← All-to-All({Bi})
{Gˆi|i ∈ [pM, (p + 1)M − 1] ∪ S} ← All-to-All({Gi})
{Wj|j ∈ S} ← P2P从其他GPU向此GPU转移权重
// 为本地和外部专家计算Grouped-GEMMs
{Hˆ i = Gˆi ⊙ BˆiWi|i ∈ [pM, (p + 1)M − 1] ∪ S}
// 执行合并：将输出路由回它们最初来源的地方
{Hi} ← All-to-All-reverse({Hˆi})
// 反转排序和重新索引
H¯p ← concat({Hi})
Hp ← reverse_sort(H¯p, I¯p)
Hp ← reshape(Hp, (Bp, K, H))
H'p ← sum(Hp, dim=K)
输出：H'p

A4 实验环境

• 数据集：

  • Megatron-Math【【2】Nemotron-math: Efficient long-context distillation of mathematical reasoning from multi-mode supervision，2025，arXiv】：用于端到端全模型速度测试，响应由gpt-oss-120b生成。
  • AIME'25：用于评估下游任务性能（准确率）。

• 模型架构：

  • MoE层：模拟了三种流行的MoE配置，分别来自gpt-oss-120b (N=128, D=H=2880, K=4), DeepSeek-V3 (N=256, D=7168, H=2048, K=8), 和 Kimi-K2 (N=384, D=7168, H=2048, K=8)。每个专家都是一个SwiGLU前馈模块。
  • 完整模型：gpt-oss-20b和gpt-oss-120b【【1】gpt-oss-120b & gpt-oss-20b model card，2025，arXiv】。

• 硬件配置：

  • GPU：P=8块 H200 GPU。

• 软件配置：

  • 实现：LLEP使用标准的Torch NCCL实现All-to-All和P2P通信。
  • 训练框架：在全参数微调实验中，使用了ZeRO-3和CPU offloading来处理梯度和优化器状态。

• 实验设置：

  • 批量大小：gpt-oss配置为每GPU 32K tokens，DeepSeek-V3和Kimi-K2配置为每GPU 16K tokens。
  • LLEP超参数：λ = 1.3, α = 1, m = 1024。
  • 不平衡场景模拟：在受控实验中，模拟了不同程度的路由不平衡，将30%到95%的token均匀地集中到1、4或16个专家上。

A4 实验结果

5.1 速度和内存性能分析

我们分析了不同规模和配置的流行MoE层的速度和内存性能，具体包括gpt-oss-120b【【1】gpt-oss-120b & gpt-oss-20b model card，2025，arXiv】、DeepSeek-V3【【7】Deepseek-v3 technical report，2024，arXiv】和Kimi-K2【【17】Kimi k2: Open agentic intelligence，2025，arXiv】中使用的MoE层。

• 实验内容：在8块H200 GPU上，对三种MoE架构的前向传播速度和每GPU峰值内存消耗进行基准测试，并模拟了从平衡到极端不平衡（95%的token集中于1个专家）的多种路由场景。
• 实验结果：如图4所示，LLEP在所有不平衡场景中均优于标准EP。
  • 加速比：不平衡程度越高，LLEP的加速效果越显著，在最极端的情况下（gpt-oss-120b，95%负载集中于1个专家）达到了6.11倍的加速。在负载均衡时，由于自适应比率λ的存在，LLEP的性能与EP相当。
  • 峰值内存：LLEP在所有不平衡场景下都保持了稳定且低水平的内存消耗。相比之下，标准EP的内存使用随不平衡度急剧增加，LLEP的内存节省高达5倍，这使得在不发生OOM的情况下可以提高吞吐量（批大小）。

5.2 真实场景下的端到端全模型速度分析

• 实验内容：我们使用真实的预训练模型gpt-oss-20b和gpt-oss-120b，在Megatron-Math数据集上进行了端到端的全模型前向传播吞吐量对比。此外，我们还进行了gpt-oss-20b的全参数微调实验，对比了达到相似下游性能（AIME'25准确率）所需的墙上时间。
• 实验结果：
  • 吞吐量（图1c）：由于预训练模型本身存在固有的路由不平衡，LLEP表现出显著优势。对于gpt-oss-20b和gpt-oss-120b，LLEP分别实现了高达2.2倍和1.88倍的吞吐量提升。并且随着GPU数量的增加，LLEP展现出更好的扩展效率和相对加速比。
  • 训练时间（图5）：在包含ZeRO-3和CPU offloading等额外开销的全参数训练中，LLEP依然实现了1.25倍的加速，即达到与EP相当的准确率所需的时间更短。

5.3 消融研究

• 实验内容：我们对LLEP的多个因素和超参数进行了消融研究，以探究其对性能的影响。
• 实验结论：
  • 批大小B（图6a）：LLEP的加速比随着批大小的增加而增大。大批次使得GPU容量饱和，计算负载均匀分布带来的收益超过了LLA算法和权重转移的开销。
  • 因子α（图6b）：较小的α值（即更严格的GPU容量阈值）能带来更高的加速比，表明在批次足够大时，LLEP倾向于实现更均衡的负载分配，即使通信成本可能更高。
  • 自适应比率λ（图7a）：在低不平衡度下，较高的λ值是有益的。这说明当路由分布足够均衡时，切换回标准EP可以避免LLEP因权重转移而引入的不必要开销。
  • 隐藏层大小D和H（图7b）：随着模型隐藏层维度的增加，LLEP的性能优势也随之扩大。这是因为更大的隐藏层维度提升了GEMM的计算效率，使得计算负载均衡的收益超过了数据和权重通信的成本。

A5 结论

我们提出了一种名为最低负载专家并行（LLEP）的新型专家并行算法，该算法能够动态地执行负载均衡，以解决MoE模型中普遍存在的不平衡路由现象，同时确保了MoE数学计算的精确性。实验结果表明，LLEP在MoE层上实现了高达5-6倍的加速和5倍的峰值内存消耗降低。对于端到端的完整模型，如gpt-oss-120b，它将吞吐量提高了多达90%。

A6 附录

A.1 分离式与融合式GROUPED-GEMM对比

cuBLAS与Triton Grouped-GEMM的性能比较。图8展示了使用cuBLAS实现的朴素for循环GEMM与一个在Triton中编写的、采用了张量内存加速器（TMA）的融合优化Grouped-GEMM核函数之间的计算时间对比。cuBLAS版本会启动N次GPU核函数，导致较高的开销，而Triton版本只启动一次。然而，如图所示，cuBLAS版本仍然胜过Triton版本，因为每个cuBLAS GEMM核函数都是针对特定硬件架构高度优化的，而Triton版本是一个通用实现。此外，尽管所有计算的总FLOPs相同，但随着专家数量的增加，计算时间显著增加。因此，计算少数几个大型GEMM比计算许多小型GEMM更优。专家并行和我们的方法（LLEP）都利用了这一原则，通过将专家权重分布到不同的EP rank上，使得每个rank只需计算少数几个专家。

A.2 专家数量

专家数量对性能的影响。与在批大小$B_i$和隐藏层大小H、D上观察到的趋势相似，图9显示，当MoE层中的专家数量（N）增加时，LLEP的效率更高，并表现出更大的加速比。