MoE大模型分布式训练：Switch Transformer与专家并行策略

点击 “AladdinEdu，同学们用得起的【H卡】算力平台”，注册即送-H卡级别算力，一站式沉浸式云原生集成开发环境，80G大显存多卡并行，按量弹性计费，教育用户更享超低价。

---

1 MoE与大模型分布式训练概述

混合专家模型（Mixture of Experts，MoE）作为大模型规模扩展的关键技术，正重塑着我们设计和训练超大参数规模人工智能模型的方式。传统的稠密Transformer模型在参数增加时面临计算成本平方级增长的瓶颈，而MoE通过稀疏激活的范式巧妙地解决了这一难题——模型总参数量可达万亿级别，但每个输入仅经过少数几个“专家”网络处理，使计算成本保持在与稠密模型相近的水平。这种稀疏性特性使得训练数百亿、甚至万亿参数级别的模型成为可能，而不需要相应倍数地增加计算资源。

然而，MoE模型的分布式训练面临着严峻挑战。随着模型规模从数百专家扩展到数千专家，专家并行（Expert Parallelism）策略的效率成为影响整体训练性能的关键因素。在万卡集群上，专家间的通信开销、负载不均衡问题以及跨节点的网络延迟相互交织，形成了一个复杂的优化难题。例如，Switch Transformer C-2048模型拥有1.6万亿参数和2048个专家，训练这样的模型需要高效的专家并行策略和精细的通信优化。

本文将深入剖析MoE大模型分布式训练的核心技术，重点聚焦Switch Transformer与专家并行策略，解析专家选择算法与负载均衡机制的内在原理，并探讨GShard与ST-MoE等先进框架在万卡集群环境下的通信优化实践。通过对这些前沿技术的系统分析，我们旨在为大规模MoE模型的高效训练提供深入见解和实践指南。

2 Switch Transformer架构解析

Switch Transformer作为MoE架构的里程碑式进展，其核心突破在于将条件计算（Conditional Computation）理念与Transformer架构深度融合。与传统稠密模型不同，Switch Transformer仅在每个输入样本上激活部分网络参数——具体来说是稀疏激活的前馈网络（FFN）层，这些层被重构为一组独立的“专家”。每个输入token通过路由机制被导向一个专家（Switch Transformer采用Top-1路由），而非所有专家，从而实现了模型容量与计算效率的优雅平衡。

2.1 稀疏激活与动态专家路由

Switch Transformer的稀疏性体现在其专家路由机制上。对于每个输入token，路由器计算其与所有专家的匹配分数，然后选择分数最高的专家进行处理。该过程可以形式化为：

$${\text{Expert}}_i=\arg \underset{i}{\max }(W_{router}\cdot x+b_{router})$$

其中$x$是输入token的表示，$W_{router}$和$b_{router}$是路由器的参数。与先前工作中使用Top-K路由（K>1）不同，Switch Transformer简化了路由算法，每次只选择一个专家，显著降低了计算和通信开销。

路由机制的核心是一个负载均衡损失函数，它鼓励专家间负载均衡，同时保持路由决策的合理性。这个损失函数由两部分组成：辅助损失（auxiliary loss）和专家均衡损失（expert balancing loss）。辅助损失确保模型训练稳定性，而专家均衡损失则防止少数专家垄断大多数token，造成负载不均。

2.2 专家容量与负载均衡

专家容量（expert capacity）是Switch Transformer中一个关键的超参数，它定义了每个专家单次处理的最大token数量。设置适当的专家容量对训练效率至关重要——容量过小会导致token被丢弃，影响模型质量；容量过大则会造成计算资源和显存的浪费。

表：Switch Transformer C-2048模型关键架构参数

参数	数值	说明	与T5-XXL对比
d_model	2080	隐藏层维度	高18%
num_heads	30	注意力头数	高25%
d_ff	6144	前馈网络维度	相当
num_layers	15	总层数	低25%
num_experts	2048	专家数量	新增维度
expert_capacity	64	专家容量	控制负载均衡
router_z_loss_coef	0.001	Z损失系数	稳定专家选择

在实际实现中，专家容量通过限制每个专家处理的token数量来实现负载均衡。例如，当某个专家被分配的token数量超过其容量时，超出的token会被丢弃或重新路由到其他专家。为了避免这种情况，Switch Transformer引入了专家容量因子，通过略微增加专家容量为路由提供一定的缓冲空间。

2.3 Switch Transformer的工程实现优化

在工程实现层面，Switch Transformer面临着一个重要挑战：如何处理被路由器丢弃的token？早期实现中存在一个逻辑缺陷——对于被丢弃的token，模型会保留其原始隐藏状态，并仍然应用路由器计算出的概率分数进行缩放。这种做法导致了多方面问题：行为不一致性、权重分配不合理、专家行为混淆等。

正确的实现方式是：初始化一个全零的张量来收集专家结果，只更新被分配到专家的token对应的位置，保持被丢弃token的结果为零。这种修改确保了模型行为与理论设计一致，专家专业化程度更高，训练过程更加稳定。虽然在实际应用中丢弃的token通常很少，这个优化对模型的正确性和稳定性有着重要意义。

3 GShard与ST-MoE的专家并行策略

随着MoE模型规模的不断扩大，单一的数据并行或模型并行已无法满足训练需求，专家并行（Expert Parallelism）作为一种专门针对MoE架构的分布式训练策略应运而生。专家并行的核心思想是将不同的专家分布到不同的计算设备上，每个设备仅负责一部分专家的计算，通过All-to-All通信操作实现token到对应专家的路由。

3.1 GShard的自动分片与路由优化

GShard是首个将MoE与Transformer架构深度结合的分布式训练框架，由Google在2020年提出。它的核心贡献在于通过条件计算和自动分片技术，实现了在保持模型性能的同时显著降低训练成本，参数规模可达6000亿。

GShard的技术方案包含几个关键创新：

1. 稀疏激活的MoE-Transformer架构：每个token仅经过两个专家，计算量从$O(N)$降为$O(2)$（N为专家总数）。它采用局部分组（Local Group Dispatching）策略，将batch内token分组并行处理，减少跨设备通信。

2. 基于SPMD的分布式训练框架：所有设备运行相同代码（单程序多数据模式），通过API标注张量分片策略，编译器自动生成高效并行代码。

3. 专家容量与负载均衡机制：GShard通过专家容量限制防止单个专家过载，结合设备级平衡损失均衡跨设备的计算负载。实验表明，这些优化可使训练速度提升30%，TPU利用率提高至90%以上。

表：GShard关键优化技术及其效果

技术	作用	效果
专家容量限制	防止单个专家过载	训练速度提升30%
设备级平衡损失	均衡跨设备的计算负载	TPU利用率提高至90%以上
随机路由	增加专家选择多样性	模型BLEU得分提升2.5

3.2 ST-MoE的分层路由与通信优化

ST-MoE（Switch Transformer MoE）在GShard基础上进一步优化了专家并行策略，特别是针对万卡集群的通信模式。ST-MoE引入了分层路由（Hierarchical Routing）机制，将路由决策过程分解为两个阶段：簇级别路由和专家级别路由。

分层路由的具体流程是：

• 第一阶段：簇级别路由。输入token首先与G个簇的“原型向量”计算相似度，通过Softmax选择最匹配的目标簇。这一步骤将路由的搜索空间从E个专家缩小至G个簇。
• 第二阶段：专家级别路由。在选定的簇内部，token再与该簇内的K个专家的路由权重计算相似度，通过另一次Softmax选择最终激活的Top-K个专家。

这种分层路由机制将路由计算复杂度从$O(E\cdot d)$降低到$O(G\cdot d+K\cdot d)$，在分布式环境中，数据仅需发送至托管目标簇的GPU子集，从而显著降低了All-to-All通信的数据交换量，直接缓解了系统的通信延迟瓶颈。

3.3 动态专家分组与结构化压缩

为了进一步优化专家并行效率，中科院自动化所研究团队提出了一种动态专家分组框架。他们发现，被语义相似的输入所激活的专家，其参数本身也存在着结构性冗余。基于这一发现，他们设计了在线双相似度聚类算法，周期性地将专家动态划分至若干专家簇。

该算法的聚类依据是一个融合相似度指标S，同时量化了专家的“结构相似性”与“功能相似性”：

• 结构相似性：通过计算两个专家权重矩阵向量化表示的余弦相似度，衡量它们在参数空间中的接近程度。
• 功能相似性：利用路由器的输出logit作为输入token的语义嵌入，通过每个专家的“激活质心”衡量其功能定位的相似度。

基于专家分组，研究团队进一步提出了共享基底与低秩残差压缩方法，将每个专家的权重矩阵分解为一个共享的公共部分和一个低秩的特有部分。在典型配置下（d=4096，K=8，r=16），专家簇内压缩比高达6.6倍，框架能在几乎不损失模型表达能力的前提下，大幅削减参数冗余。

4 万卡集群通信优化实战

在万卡集群环境下训练MoE模型，通信开销往往成为系统性能的主要瓶颈。特别是在专家并行模式下，All-to-All通信操作随着专家数量的增加而线性增长，其延迟可占单步训练时间的50%以上。因此，通信优化是实现MoE模型高效训练的关键。

4.1 All-to-All通信的瓶颈分析

在MoE模型的训练过程中，专家并行需要在多个设备间交换token和对应的专家输出，这一过程通过All-to-All通信实现。All-to-All属于高开销的集体通信操作，尤其在跨节点（跨服务器）场景下，其性能受限于节点间网络的带宽和延迟。

具体来说，MoE训练中的All-to-All通信包含两个阶段：

1. 前向传播中的token分发：将输入序列的token根据路由结果发送到对应的专家所在设备。
2. 反向传播中的梯度聚合：收集各专家计算的梯度，回传给相应设备以更新模型参数。

在万卡集群上，这种通信模式可能导致严重的性能问题。NVIDIA的研究显示，在DeepSeek R1模型的32 GPU训练中，某些层的MoE出现了严重的负载不均衡问题——从某一rank发送到另一rank的token数明显更多，形成了通信热点。这种不均衡在多次迭代中保持稳定，导致部分设备长期处于高负载状态，而其他设备则利用率不足。

4.2 通信与计算重叠技术

为了降低All-to-All通信的开销，业界提出了多种通信与计算重叠技术，即通过合理安排计算和通信任务，使它们尽可能并行执行，从而用计算时间掩盖通信时间。

4.2.1 基于1F1B的流水线重叠

NVIDIA和小红书Agi Infra团队合作提出了基于1F1B（One-Forward-One-Backward）流水线的MoE A2A通信计算重叠方案。该方案通过调整1F1B调度，在稳态阶段实现前向传播、反向传播与All-to-All通信的重叠。

具体实现包括两种变体：

• Conventional 1F1B：通过增加warmup step，创造前向传播与反向传播之间的通信重叠机会。但这种方法会导致warmup step过大，activation显存占用增加。
• Interleaved 1F1B：通过将稳态的第一个micro batch的fprop提前到warmup stage，实现在不增加峰值显存的前提下，实现不同micro batch前反向之间的EP A2A与计算的重叠。

进一步的优化是细粒度调度，通过拆分权重梯度计算（dw），将dispatch和combine操作的执行提前，从而更充分地重叠计算与通信。研究表明，这种拆分dw的方案比基线方法能更进一步减少整体执行时间。

4.2.2 DualPipe与DeepEP优化

DeepSeek V3论文提出了一种创新的流水线并行算法：DualPipe，专门针对MoE模型的通信优化。DualPipe实现了在两个micro batch forward和backward之间的EP A2A通信与计算重叠。

与1F1B方案相比，DualPipe具有不同的特性：

• PP bubble overhead：DualPipe的流水线泡沫为KaTeX parse error: Expected 'EOF', got '&' at position 14: (PP/2-1) × (F&̲B + B -3W)，而1F1B with A2A overlap为$(PP-1)\times (F+B)/vpp$。
• 内存占用：DualPipe的参数内存为2X，activation内存为$PP+1$；而1F1B方案分别为1X和$PP+(PP-1)/vpp$。

与此同时，DeepSeek开源的DeepEP通信库通过将通信拆分为机内通信和机间通信，并将二者进行overlap，实现了很高的通信效率。这一工作已集成到NVIDIA Megatron-LM中。

4.3 跨数据中心专家并行

随着MoE模型规模的持续扩大，单个数据中心的计算资源已无法满足需求，跨数据中心训练成为必然选择。然而，跨数据中心的专家并行面临着严峻的网络带宽限制。arXiv:2510.19470提出了一种名为HybridEP的建模指导框架，专门用于优化受限带宽下的专家并行。

HybridEP的核心思想是动态调整专家放置策略，通过混合专家/数据传输减少通信流量和频率。它包含两个关键技术：

1. 域基分区：在GPU级别构建混合模式与特定通信拓扑之间的映射关系。
2. 参数高效迁移：通过减少专家传输开销和扩大域大小来优化拓扑结构。

实验结果表明，HybridEP在受限带宽环境下优于现有最先进的MoE训练系统，最高可达5.6倍的性能提升。在大规模模拟中，HybridEP在1000个数据中心不同带宽条件下实现了最高1.45倍的加速。

表：MoE通信优化技术对比

优化技术	核心思想	适用场景	通信减少效果
分层路由	将路由分为簇级和专家级两步	大规模专家模型（专家数>100）	路由计算复杂度从O(E·d)降至O(G·d+K·d)
1F1B通信计算重叠	调整流水线调度，使计算掩盖通信	中等规模集群（单数据中心）	All-to-All通信时间可被大部分掩盖
DualPipe	专门设计的MoE流水线并行算法	超大规模模型训练	减少约50%的通信等待时间
HybridEP	动态专家放置，混合专家/数据传输	跨数据中心训练	在受限带宽下最高5.6倍性能提升

5 负载均衡机制的技术演进

负载均衡是MoE模型训练中的核心挑战，直接影响模型质量、训练效率和系统稳定性。不均衡的专家负载会导致多种问题：部分专家过拟合而其他专家欠拟合，计算资源利用率低，以及通信模式不稳定等。因此，负载均衡机制一直是MoE研究的重点。

5.1 负载不均的根源分析

要解决负载不均问题，首先需要理解其产生根源。根据NVIDIA对大规模EP推理的观察，负载不均主要源于以下几个方面：

1. 数据分布的特性：不同数据集会导致不同的专家激活模式。例如，在机器翻译数据集中，第36层的MoE出现了明显的热点专家；而在GSM8K数据集中，热点层变成了第57层。这表明专家负载失衡具有数据集敏感性。

2. 路由决策的集中性：由于Softmax函数的特性，路由器倾向于将高分赋予少数专家，导致“赢者通吃”现象。虽然Top-K路由（K>1）可缓解这一问题，但仍无法完全避免负载集中。

3. 专家初始化的随机性：专家网络的随机初始化导致其初始能力差异，在训练早期就可能形成强者恒强的马太效应。

4. 训练动态的变化：随着训练进行，数据分布和模型参数不断变化，导致专家激活模式也随之改变，静态的负载均衡策略难以适应这种动态变化。

值得注意的是，研究发现专家失衡分布在数十到数百次迭代中相对稳定，这为基于历史统计的负载均衡策略提供了理论基础。

5.2 静态负载均衡策略

早期MoE模型主要采用静态负载均衡策略，通过在损失函数中添加正则项来鼓励均衡的专家分配：

1. 辅助损失函数：Switch Transformer引入了路由器z损失（router z loss）和辅助损失系数（router aux loss coef），通过惩罚过大的路由器logits和均衡的专家选择分布来维持负载均衡。这些损失项在训练过程中与主损失函数一起优化，但不参与推理。

2. 专家容量限制：通过设置专家容量（expert capacity）硬性限制每个专家处理的token数量。当专家被分配的token超过其容量时，超出的token会被丢弃或重新路由。这一机制直接防止了单个专家过载，但可能导致token丢失和模型性能下降。

3. 随机路由：在训练初期引入一定随机性，避免早期固化。例如，GShard在训练前10%步骤中使用随机路由，随后逐渐过渡到Top-K路由。这种方法有助于专家网络的均衡初始化，但对训练后期的负载不均效果有限。

5.3 动态负载均衡技术

针对静态策略的不足，业界提出了多种动态负载均衡技术：

1. 在线专家聚类与重组：中科院自动化所的研究团队提出了在线双相似度聚类算法，周期性地根据专家的结构相似性和功能相似性对专家进行动态分组。这种方法不仅实现了粗粒度的负载均衡，还为后续的参数压缩奠定了基础。

2. 动态专家复制：中科大团队研发的基于专家选择分布预测的框架，通过复制高频专家、替换非重要专家的方式实现负载平衡。该框架依据token对专家选择偏好的预测结果，精准判断出专家动态访问的频率和重要性，进而调整专家分布。

3. 在线负载均衡器：NVIDIA在TensorRT-LLM中实现了在线专家负载均衡器，动态平衡专家负载分布并适应在线流量模式的变化。该均衡器包含统计组件、复制与放置计算组件以及更新权重与放置组件，能够实时监控专家负载并调整专家放置策略。

表：负载均衡技术对比

技术类型	代表方法	优点	缺点
静态策略	辅助损失函数、专家容量限制	实现简单，计算开销小	适应能力差，可能影响模型性能
动态重组	在线双相似度聚类、专家复制	适应数据分布变化，均衡效果好	实现复杂，需要额外计算
在线均衡	TensorRT-LLM负载均衡器	实时适应，适合在线推理	系统复杂度高，需要专门架构

5.4 负载均衡的未来趋势

从技术发展角度看，MoE负载均衡机制正呈现几个明显趋势：

1. 从被动补偿到主动预测：传统方法多在负载不均发生后进行补偿，而新技术如中科大团队的专家选择分布预测框架则尝试预测负载模式并提前调整。

2. 从全局均衡到局部优化：早期方法关注全局负载均衡，而越来越多的工作开始考虑数据局部性和网络拓扑，实施基于物理约束的局部优化。

3. 从单一目标到多目标权衡：负载均衡不再孤立考虑，而是与模型性能、通信开销、能源效率等多目标共同优化，如中科院自动化所提出的“优化三难困境”解决方案。

4. 从手工设计到自主学习：随着强化学习等技术的发展，部分研究开始探索自动学习负载均衡策略的可能性，减少对手工设计规则的依赖。

6 MoE技术的未来展望

随着大模型规模的持续扩展，MoE技术正处于快速演进的过程中。从当前的Switch Transformer、GShard到未来的发展方向，我们可以观察到几个明显趋势。

硬件与软件的协同优化将更加深入。专为MoE设计的新型硬件架构正在涌现，如NVIDIA的Grace Blackwell系统提供了高带宽的GPU间互联，为专家并行提供了更好的基础设斀。同时，软件栈也在不断优化，如TensorRT-LLM对大规模专家并行的支持，以及HybridEP针对跨数据中心场景的优化。这种硬件与软件的紧密协同将极大提升MoE模型的训练和推理效率。

自动化与自适应成为关键方向。当前MoE模型中的专家数量、容量因子等超参数多凭经验设置，缺乏理论指导。未来，我们预计会出现更多自动化MoE架构搜索技术，根据任务需求和硬件环境自动确定最优的专家配置。同样，在路由算法和负载均衡策略方面，自适应性将更强，能够根据数据分布动态调整专家选择和行为。

跨模态与跨域扩展潜力巨大。当前MoE研究主要集中在语言模型，但其稀疏激活的思想同样适用于视觉、多模态等其它领域。将MoE技术与不同模态的网络架构结合，处理图像、视频、音频等数据，是一个充满希望的方向。专家网络可能根据不同模态、不同概念或不同任务自动专业化，形成更加智能的条件计算模式。

理论突破与效率极限的探索仍需深入。虽然MoE在实践中已展现出巨大价值，但其理论基础仍相对薄弱。如何理论分析MoE的泛化能力、优化动态和缩放定律，是未来研究的重要课题。同时，在模型规模持续增长的同时，我们必须考虑能耗和环境影响，寻找规模与效率的最佳平衡点。

总之，MoE技术作为大模型规模扩展的关键路径，正在从多个维度持续进化。从架构创新到并行策略，从通信优化到负载均衡，每一次突破都为更大规模、更高效的AI模型开辟了道路。随着技术的不断成熟，MoE有望成为实现下一代超大规模人工智能的核心技术支撑。

---

点击 “AladdinEdu，同学们用得起的【H卡】算力平台”，注册即送-H卡级别算力，一站式沉浸式云原生集成开发环境，80G大显存多卡并行，按量弹性计费，教育用户更享超低价。