MoE LLMs中负载均衡策略演进的回顾：陷阱与经验

大家好，欢迎来到我对专家混合模型（MoE）多年来演进的充满乐趣又略显书呆子气的探索——我将特别关注研究人员在负载均衡方面所采用的巧妙、有时混乱但始终有趣的方法。这篇文章更像是我持续探索的“实验室笔记”：学术分析与个人反思的混合体。

故事从GShard开始——那时人们意识到，拥有数十亿（或数万亿）参数的模型可以通过巧妙的“稀疏化”来更快地训练，同时保持高精度。从那时起，我们见证了一系列创新。在这里，我想将我们如何从GShard发展到最新的创新，例如DeepSeek-V3——每一个的贡献、遇到的陷阱以及仍未解决的重大问题——串联起来。

1. 引言

1.1 为什么选择稀疏专家混合模型（MoE）？

我们先来介绍一下背景。当人们意识到可以大幅增加模型容量（参数）而无需线性增加计算量（FLOPs）时，MoE架构风靡全球。其核心思想是，对于每个token，我们只“激活”一小部分总参数（即几个专家），而不是强制每个参数都参与其中。

然而，这种方法的阴暗面很快显现出来：如果你只将token发送给一小部分专家，如何保持负载“均衡”，以避免单个专家被token压垮而其他专家闲置？这就是负载均衡的核心问题，在大规模部署时很难解决。

1.2 本文内容

我将回顾一些里程碑式的MoE系统，从GShard（最早大规模主流的MoE系统）开始，一直到全新的DeepSeek-V3。虽然我们会涵盖常见的模型（如Switch Transformer和GLaM），但我想强调每个系统遇到的陷阱以及新架构如何克服这些陷阱。

如果你阅读本文是为了获取实际经验，那太好了：我将尽量保持足够的学术严谨性，使其对高级从业者和研究人员有所帮助。但希望它足够轻松愉快，让你读几段后不会睡着——毕竟这是一个博客，而不是期末考试！

1.3 我注意到的关键主题

1. 路由方法：top-2门控、单专家门控、top-K门控、关联感知门控……是的，我们喜欢门控术语！
2. 辅助损失：有助于推动专家的均衡使用，但如果干预过多也会损害性能。
3. 容量约束：“容量因子”是“每个专家在丢弃多余token之前可以处理多少token”的雅称。
4. 实现细节：从“随机调度”到分层全对全。HPC（高性能计算）视角非常相关。
5. 可扩展性：在某些情况下，我们谈论的是成千上万的专家，因此分布式计算开销不容忽视。

2. 历史演进：从GShard到Switch

2.1 GShard⁰：先驱

GShard（由Google推出）被广泛认为是首批大规模、超稀疏的MoE框架之一。它通过展示如何仔细分片层并在专家之间平衡token，从而训练约6000亿参数的模型，改变了人们的看法。

GShard的门控方法通常为每个token选择前2个专家。我们用以下公式表示：

$$\text{GATE}(x)=\text{Top2}(W_{gate}\cdot x)\backslash text\{GATE\}(x)=\backslash text\{Top2\}(W\_\{gate\}\backslash cdot\; x)$$

其中$xx$是token嵌入，$W_{gate}W\_\{gate\}$是路由器的权重矩阵。只有前2个专家会被激活。然而，为了避免每个专家过载，我们需要引入：

1. 专家容量，$C\approx \frac{2N}{E}C\; \backslash approx\; \backslash frac\{2N\}\{E\}$，如果存在$NN$个token和$EE$个专家。如果一个专家超过容量限制而过载，一些token会被丢弃（或溢出到下一层）。
2. 辅助负载均衡损失，通常形式为

$${\mathcal{L}}_{\text{aux}}\sum _{e=1}^E{f}_e{P}_e\backslash mathcal\{L\}\_\{\backslash text\{aux\}\}\; \backslash \; \backslash sum\_\{e=1\}^E\; f\_e\; P\_e$$

其中$f_{e}f\_e$是分配给专家$ee$的token比例，$P_{e}P\_e$是专家$ee$的平均门控概率。这个损失促使系统更均匀地分配token给各个专家。

3. 局部组，这样并非每个token都与全局中的每个其他token竞争。

陷阱：你猜对了——丢弃token并不是什么光彩的事。如果token超过容量，它们可能会得到不完整的处理。此外，top-2门控和随机调度的开销在大规模部署时可能会很重。而且，过度依赖辅助损失有时会强制“虚假”的token分布，从而损害专业化学习。但即便如此，GShard仍然证明了MoE是可行的，并且值得付出努力。容量约束的概念非常准确，我们在几乎所有后续的MoE方法中仍然可以看到它。

2.2 Switch Transformer¹：“少即是多”的时代

Switch Transformer本质上说：“嘿，我们只将每个token路由到一个专家。”这使得门控更简单（选择门控logits最高的专家），并大大减少了计算开销。门控函数如下：

$$g_i(x)=\text{softmax}(W_{\text{router}}\cdot x{)}_{i}g\_i(x)\; =\; \backslash text\{softmax\}(W\_\{\backslash text\{router\}\}\; \backslash cdot\; x)\_i$$

我们选择

$$\text{expert index}(x)={\text{argmax}}_i{g}_i(x)\mathrm{.}\backslash text\{expert\; index\}\; (x)=\{\backslash text\{argmax\}\}\_i\; g\_i(x).$$

Switch Transformer的主要创新是其单专家路由，因为激活的专家越少，通常代码越简单，训练速度也越快。为了更好地平衡负载，他们保留了类似于GShard方法的辅助负载均衡损失。他们还定义了一个容量因子，让专家处理比简单比例更多的token。例如，

$$C=\text{CF}\times \frac{\text{tokens per batch}}{\text{number of experts}}C\; =\; \backslash text\{CF\}\; \backslash times\; \backslash frac\{\backslash text\{tokens\; per\; batch\}\}\{\backslash text\{number\; of\; experts\}\}$$

Switch Transformer的收益与权衡相当明显：速度更快，因为每个token只执行一次前馈路径，但可能会面临更大的token溢出风险（你只有一个专家来处理它们！）。一些token会被“丢弃”或强制通过残差通路。

陷阱与经验教训：单专家路由在概念上更简单，通常也更快。但如果容量因子（CF）设置不正确，可能会导致过多的token被丢弃，或者过多的token被分配给一个专家。Switch Transformer基本上阐明了如何通过精心选择的超参数调优来发挥奇效。Switch简化了MoE门控——表明即使是top-1路由也可能实现扩展。这激发了后续关于“哪个K最好？”和“如何最好地处理溢出？”的研究。

3. 改进与变体：GLaM、DeepSpeed-MoE、ST-MoE、Mixtral

3.1 GLaM²：重新审视Top-2，注重效率

GLaM（通用语言模型）重新引入了top-2门控，但对能效提出了新的看法——据报道，它仅使用GPT-3约1/3的训练能量，却实现了更好的零样本性能。他们使用了以下公式：

$$y=\sum _{i=1}^2{g}_i\cdot E_i(x),y\; =\; \backslash sum\_\{i=1\}^2\; g\_i\; \backslash cdot\; E\_i(x),$$

其中$g_{i}g\_i$是门控权重，$E_i(x)E\_i(x)$是选定的两个专家。类似地，GLaM引入了一个精心调整的辅助损失，以鼓励token在专家之间均匀分布。此辅助损失通过优化专家利用率来惩罚不平衡路由：

$${\mathcal{L}}_{\text{aux}}=\alpha \cdot \sum _{i=1}^E{f}_i\cdot p_i,\backslash mathcal\{L\}\_\{\backslash text\{aux\}\}=\backslash alpha\; \backslash cdot\; \backslash sum\_\{i=1\}^E\; f\_i\; \backslash cdot\; p\_i,$$

其中$f_{i}f\_i$是路由到专家$ii$的token比例，$p_{i}p\_i$是专家$ii$的平均门控概率，$\alpha \backslash alpha$是权重因子。为了防止专家过载，GLaM还引入了容量约束，其中每个专家的最大token容量定义为：

$$C=\frac{\text{tokens per batch}}{\text{number of experts}}\cdot \text{capacity factor}\mathrm{.}C\; =\; \backslash frac\{\backslash text\{tokens\; per\; batch\}\}\{\backslash text\{number\; of\; experts\}\}\; \backslash cdot\; \backslash text\{capacity\; factor\}.$$

超过此容量的token将被丢弃，并通过残差连接传递到下一层。通常使用$1.251.25$的容量因子来平衡token溢出和计算效率。

陷阱与经验教训：GLaM强调了当只激活模型参数的一小部分时，能耗节省可以有多大。（他们将其与GPT-3进行比较，并表示：“看，我们只用了它一小部分能量。你们都应该注意了！”）尽管GLaM发现确实可以超越密集计算的成本，但必须注意专家使用中潜在的不平衡——特别是在实际文本分布中。模型精心调整的门控和容量约束有助于防止专家过载。

3.2 DeepSpeed-MoE³：聚焦推理

微软的DeepSpeed-MoE是负载均衡如何成熟以应对训练期间的token分布挑战和推理期间的有效专家利用的典型例子。DeepSpeed-MoE吸取了早期MoE系统的教训，引入了几项创新来解决token负载不平衡问题。

核心思想。DeepSpeed-MoE的核心在于通过灵活的多专家和多数据并行设计来扩展MoE框架，以优化负载均衡，特别是关注token级别的专家间分布。目标很明确：确保没有专家过载，同时保持训练高效和可扩展到分布式GPU。

DeepSpeed-MoE沿袭Switch Transformer，采用top-1门控机制。这简化了路由，并相对于top-2或top-k门控降低了计算开销。为了防止token不平衡，DeepSpeed-MoE添加了一个辅助负载均衡损失。该损失促使token在专家之间更均匀地分布：

$${\mathcal{L}}_{aux}=\alpha \sum _{i=1}^E\mathrm{\mid }{f}_i-\frac{1}{E}\mathrm{\mid },\backslash mathcal\{L\}\_\{aux\}\; =\; \backslash alpha\; \backslash sum\_\{i=1\}^E\; |f\_i\; -\; \backslash frac\{1\}\{E\}|,$$

其中${𝑓}_i𝑓\_i$是路由到专家$ii$的token比例，$EE$是专家总数，$\alpha \backslash alpha$是可调权重。此项阻止token过度集中于少数专家。DeepSpeed-MoE还采用动态token再分配策略：在训练期间，DeepSpeed-MoE动态重新分配token，以防止任何单个专家成为瓶颈。超出专家容量的token会被重新路由到其他不那么繁忙的专家，而不是被丢弃或传递到残差通路。为了进一步减轻token分布不均的影响，DeepSpeed-MoE引入了Residual-MoE架构。在此架构中，密集MLP的输出与选定专家的输出结合，将专家输出视为“残差修正”：

$$y=\text{MLP}(x)+g\cdot E(x),y=\backslash text\{MLP\}(x)\; +\; g\; \backslash cdot\; E(x),$$

其中$gg$是门控得分，$E(x)E(x)$是专家输出。这确保即使未充分利用的专家也能对模型的总体输出做出有意义的贡献。

跨GPU负载均衡。 DeepSpeed-MoE利用了更深层受益于更多专家的观察，在靠后的层中使用更多专家。虽然这确保了高效的参数使用和改进的模型质量，但这可能导致层之间专家数量的差异。在这种情况下，均匀的并行度是低效的，因为：

• 将并行度设置为最小专家数量会导致批量大小减少，并增加处理较大层的GPU的内存需求。
• 将并行度设置为最大专家数量会导致负载不平衡，其中一些GPU处理的专家多于其他GPU。

DeepSpeed-MoE系统通过动态调整跨层的并行度并优化工作负载分配来解决此问题。对于给定的模型，系统允许模型的不同部分使用不同程度的专家并行和数据并行。例如，具有32个专家的层可能使用32路专家并行和4路数据并行，而具有128个专家的层可能使用128路专家并行和1路数据并行。这确保了每个GPU每层只处理一个专家，无论该层中的专家总数是多少。通过将专家并行与每层中的专家数量对齐，系统避免了某些GPU处理的专家多于其他GPU的情况。这避免了瓶颈并确保了资源的最大利用率。

陷阱与经验教训。虽然DeepSpeed-MoE在平衡token负载方面取得了令人印象深刻的成果，但仍存在一些权衡：

• 配置复杂性：平衡容量因子、辅助损失权重和专家并行设置需要仔细调整。
• 真实世界数据的边缘情况：NLP任务中的文本分布可能高度偏斜，如果未仔细调整，这仍然会给门控机制带来压力。

尽管如此，DeepSpeed-MoE证明了token负载均衡不仅仅是理论上的优化——它是训练大规模MoE系统的实际必要条件。通过将路由创新与系统级优化相结合，它为MoE训练的效率和可扩展性设定了新标准。即使你拥有出色的训练流水线，你仍然需要很好地处理推理——特别是如果你想要实时或交互式应用程序。

3.3 ST-MoE⁴：容量因子调整与路由器Z损失

ST-MoE（稳定可迁移的专家混合模型）标志着稀疏专家模型向前迈出了重要一步，为训练稳定性和可迁移性方面的一些长期挑战提供了解决方案。虽然像Switch Transformer和GLaM这样的早期模型奠定了基础，但ST-MoE通过架构创新和超参数优化相结合的方式，改进了这些思想，解决了陷阱。

ST-MoE的一项突出贡献是路由器Z损失，旨在稳定训练而不降低质量。稀疏模型通常因路由中的指数函数而导致不稳定性，这些函数会放大微小的数值误差。路由器Z损失通过对路由网络中的大logits施加惩罚来缓解这种情况，从而有效地控制其幅度：

$${\mathcal{L}}_z=\frac{1}{B}\sum _{i=1}^B(\text{log}\sum _{j=1}^N\text{exp}(x_{ij}){)}^2\backslash mathcal\{L\}\_z\; =\; \backslash frac\{1\}\{B\}\; \backslash sum\_\{i=1\}^B(\backslash text\{log\}\backslash sum\_\{j=1\}^N\; \backslash text\{exp\}(x\_\{ij\}))^2$$

其中，$BB$是批次大小，$NN$是专家数量，$x_{ij}x\_\{ij\}$是用于路由的 logits。这种损失不仅减少了不稳定性，而且略微提高了模型质量——对于稀疏模型训练来说是双赢。

调整容量因子。ST-MoE 还强调了容量因子（CF）在平衡效率和性能方面的关键作用。为了进一步改进负载均衡，ST-MoE 引入了类似于 DeepSpeed-MoE 的辅助损失，确保 tokens 在专家之间均匀分布。

陷阱与经验教训：ST-MoE 实现了改进的稳定性与质量权衡：GLaM 和 DeepSpeed-MoE 等早期方法在负载均衡方面取得了进展，但通常需要在模型质量或可扩展性方面做出妥协。ST-MoE 的路由器 z-loss 表明，无需此类权衡即可实现稳定性。然而，ST-MoE 并非没有局限性。调整 CF 和 z-loss 权重等超参数的复杂性需要仔细的实验。总而言之，ST-MoE 代表了 MoE 架构演变的新篇章，将稳健的设计原则与解决长期挑战的创新解决方案相结合。

3.4 Mixtral 8x7B⁵：时间局部性与专用稀疏核

Mixtral 8x7B 是一款创新的稀疏专家混合（SMoE）语言模型，旨在解决 MoE 架构中长期存在的负载均衡挑战。让我们深入探讨其在每个专家 token 负载均衡问题上的独特方法，并揭示它提供的经验教训。

Mixtral 的核心是采用 Top-2 门控机制来路由 tokens：每层包含 8 个专家，在给定时间每个 token 仅激活 2 个专家。这种方法通过将每个 token 限制为仅两个专家，有效地将每个 token 的活跃参数数量限制在 13B，与 Llama 2 70B 等密集模型相比，显著减少了参数。同时，所选专家可以因 token 和层而异，增强了模型对不同输入模式的适应性。

专家分配中的时间局部性。路由分析最显著的发现之一是观察到的专家分配中的时间局部性。Tokens 在连续位置上通常保持相同的专家分配，特别是在更深的层中：在第 15 层和第 31 层中，连续 tokens 分配给相同专家的频率远高于随机分布的预测。这种现象被称为高重复率，表明了结构化行为，可能与输入的句法或位置特征有关。时间局部性既提供了机会也带来了挑战：它确保了 token 分配更平滑的转换，最大限度地减少了特定专家的突然工作负载峰值，但它也可能导致 tokens 在一部分专家上过度集中，特别是在具有句法或位置规律性的数据集中。与 DeepSpeed-MoE 类似，Mixtral 也采用了动态 token 再分配策略：当专家超出其 token 容量时，多余的 tokens 通过重新分配给其他负载较轻的专家来有效处理。

使用稀疏核减轻 GPU 过载。Mixtral 采用专门的稀疏核（例如 Megablocks）来缓解 token 过载。Megablocks 有效地处理可变 token 分配，利用高算术强度来加速计算。用于特定专家的 tokens 会在 GPU 之间动态路由。这种分区策略虽然有效，但需要仔细的负载均衡以避免 GPU 过载。

陷阱与经验教训。Mixtral 对不同数据集专家使用情况的分析强调了理解特定领域 token 分布的重要性。如果数据集分布发生变化（例如，从新闻文章变为代码），“局部性”可能会消失。因此，每种方法都对您的数据有假设。

4. 下一代方法：OpenMoE、DeepSeekMoE、JetMoE 等

4.1 OpenMoE⁶：上下文无关专业化与末端丢弃

OpenMoE 是标准 top-k 门控公式的另一个有趣的变体，它具有容量限制和辅助平衡损失。但它以识别 MoE 系统在大量训练运行中出现的某些古怪行为而闻名，即上下文无关专业化和末端丢弃。

• 上下文无关专业化：Tokens 的路由更多地可能取决于 token ID 或表面模式，而不是更深层次的语义属性，尤其是在预训练早期。
• 末端丢弃：在长序列中，容量约束通常在序列后期触发，因此那些后期的 tokens 更容易被丢弃。这显然会损害依赖序列末尾上下文的任务的性能。

与许多其他 MoE 一样，OpenMoE 采用 top-k 选择，其中 $k=2k=2$。与 GShard 和 Switch 类似，采用了以下形式的负载均衡损失：

$${\mathcal{L}}_b=E\cdot \sum _{i=1}^E{m}_i\cdot P_i,\backslash mathcal\{L\}\_b\; =\; E\; \backslash cdot\; \backslash sum\_\{i=1\}^E\; m\_i\; \backslash cdot\; P\_i,$$

其中 $m_{i}m\_i$ 是路由到专家 $ii$ 的 token 比例，$P_i$ 是专家 $ii$ 的平均门控概率。为了稳定训练，他们还通过惩罚大的 logits 引入了路由器损失：

$${\mathcal{L}}_z=\frac{1}{B}\sum _{j=1}^B(\text{log}\sum _{i=1}^E(f(x_j{)}_i))\mathrm{.}\backslash mathcal\{L\}\_z\; =\; \backslash frac\{1\}\{B\}\; \backslash sum\_\{j=1\}^B\; (\backslash text\{log\}\; \backslash sum\_\{i=1\}^E\; (f(x\_j)\_i)).$$

为了保持工作负载平衡，OpenMoE 对每个专家都施加了容量限制。这可以确保在通过专家并行（即，将不同的专家分配到不同的 GPU）训练和部署 MoE 模型时的吞吐量。然而，OpenMoE 首次发现了末端丢弃问题，即如果先前的 tokens 已经填满了专家，则序列后期的 tokens 将被丢弃。在仅解码器 MoE 架构中，由于自回归特性，序列中后期的 tokens 更容易被丢弃。这对于像指令遵循这样的序列任务来说尤其成问题，因为后期 tokens 可能携带关键信息。

陷阱与经验教训：OpenMoE 告诉我们要警惕分布上的怪癖，特别是如果你关注依赖于完整序列覆盖的任务或希望实现强大的领域适应。如果门控功能捕捉到肤浅的模式（如 token ID），它可能无法很好地适应新领域。由于容量限制是批处理机制，批处理末尾的 tokens 可能会饿死。

4.2 DeepSeekMoE⁷：细粒度专家与共享专家

在我们讨论最新版本（DeepSeek-V3）之前，让我们先讨论 DeepSeekMoE。它因将每个专家分成更细的子专家并隔离一些“共享专家”（这些专家始终激活，即绕过门控）而闻名。这种方法旨在减少参数冗余，同时仍为专用子专家提供足够的 다양성。

细粒度专家分割。DeepSeekMoE 引入了细粒度专家分割的概念，以增强专家专业化。这是通过将每个专家分割成更小的单元来实现的，同时保持总参数数量和计算成本不变

$$h_t^l=\sum _{i=1}^{mN}{g}_{i,t}\cdot {\text{FFN}}_i(u_t^l)+u_t^l,h\_t^l\; =\; \backslash sum\_\{i=1\}^\{mN\}\; g\_\{i,t\}\; \backslash cdot\; \backslash text\{FFN\}\_i\; (u\_t^l)\; +\; u\_t^l,$$

其中 $mNmN$表示细粒度专家的总数，$g_{i,t}g\_\{i,t\}$是专家 $ii$ 的门控值。路由机制为每个 token 选择 top-mK 个专家。

假设您有 $mNmN$ 个总子专家，其中包含 $N_r=mNN\_r=mN$ 个“路由”专家和 $N_{s}N\_s$ 个“共享”专家。对于第 $tt$ 个 token $u_t^{l}u\_t^l$ 在第 $ll$ 层：

$$h_t^l=u_t^l+\sum _{i=1}^{N_s}{\text{FFN}}_i^{(s)}(u_t^l)+\sum _{j=1}^{N_r}{g}_{j,t}\cdot {\text{FFN}}_j^{(r)}(u_t^l),h\_t^l\; =\; u\_t^l\; +\; \backslash sum\_\{i=1\}^\{N\_s\}\; \backslash text\{FFN\}\_i^\{(s)\}\; (u\_t^l)\; +\; \backslash sum\_\{j=1\}^\{N\_r\}\; g\_\{j,t\}\; \backslash cdot\; \backslash text\{FFN\}\_j^\{(r)\}\; (u\_t^l),$$

其中 $g_{j,t}g\_\{j,t\}$ 是子专家 j 的门控值，通常在 top-$K_r\) 中选择。

DeepSeekMoE 采用了两级负载均衡损失，以解决潜在的路由崩溃和计算瓶颈问题

• 专家级平衡损失：此损失鼓励 token 在专家之间均匀分布

$${\mathcal{L}}_{\text{ExpBal}}={\alpha }_1\sum _{i=1}^{mN-K_s}{f}_i\cdot P_i,\backslash mathcal\{L\}\_\{\backslash text\{ExpBal\}\}\; =\; \backslash alpha\_1\; \backslash sum\_\{i=1\}^\{mN-K\_s\}\; f\_i\; \backslash cdot\; P\_i,$$

其中 $f_{i}f\_i$ 是路由到专家 $ii$ 的 token 比例，$P_i$ 是专家 $ii$ 的平均路由概率。

• 设备级平衡损失：它确保设备之间的计算负载平衡

$${\mathcal{L}}_{\text{DevBal}}={\alpha }_2\sum _{i=1}^D{f}_i^{\mathrm{\prime }}\cdot P_i^{\mathrm{\prime }},\backslash mathcal\{L\}\_\{\backslash text\{DevBal\}\}\; =\; \backslash alpha\_2\; \backslash sum\_\{i=1\}^\{D\}\; f\text{'}\_i\; \backslash cdot\; P\text{'}\_i,$$

其中 $DD$ 是设备数量，$f_i^{\mathrm{\prime }}f\text{'}\_i$ 和 $P_i^{\mathrm{\prime }}P\text{'}\_i$ 分别表示设备 $ii$ 的平均 token 比例和概率。

4.3 JetMoE⁸：无丢弃 MoE 与流水线并行

大多数 MoE 方法在超出容量时考虑丢弃 tokens，而 JetMoE 则尝试一种“无丢弃”方法。其设计确保永远不会完全丢弃任何 tokens

1. 无丢弃 MoE：门控机制经过精心管理，以不超过每个专家的最大容量。
2. 流水线并行：JetMoE 不会将专家分散到多个设备上，而是将一层的所有专家保留在同一设备上，为不同层形成流水线。

JetMoE 采用 top-2 路由机制，并具有先前定义的所有负载均衡功能，例如基于频率的辅助负载均衡损失和路由器 z-Loss。与以前使用固定容量因子方法不同，JetMoE 继承自 MegaBlocks，它用块稀疏矩阵操作取代了传统的 token 丢弃方法。MegaBlocks¹¹ 实现了自定义块稀疏 GPU 内核，以有效地处理 MoE 计算的动态和负载不平衡性质。通过根据专家分配动态构建块稀疏矩阵拓扑，该框架确保所有 tokens 都得到处理而不会被丢弃，这与使用固定容量因子的传统方法不同。

陷阱与经验教训：实现无丢弃可能会变得复杂。您可能会看到开销或次优门控。如果您做得好，您将拥有一致的 token 覆盖。这对于丢弃 tokens 是灾难性任务（如问答或代码生成）很有吸引力。但您必须实时处理容量受限门控的复杂性。

4.4 Skywork-MoE⁹：门控 Logit 归一化与自适应辅助

Skywork-MoE 是一个高性能专家混合（MoE）模型，拥有 1460 亿参数和 16 个专家。该模型利用 Skywork-13B（一个密集语言模型）的架构，使用其预训练的密集检查点进行初始化。Skywork-MoE 结合了门控 Logit 归一化和自适应辅助损失系数等先进技术，以改善专家多样化和层特定负载均衡。它引入了两个巧妙的理念来解决专家不平衡问题。

1. 门控 Logit 归一化：他们在 softmax 之前对门控 logits 进行标准化，控制“锐度”。
2. 自适应辅助损失系数：如果某一层丢弃了过多的 tokens，则会自动增加平衡惩罚。

MoE 层将 transformer 中的标准 FFNs 替换为多个专家，为每个输入 token 选择性地激活 top-k 个最相关的专家。

用于负载均衡的辅助损失。为了防止路由崩溃（少数专家占据主导地位），Skywork-MoE 采用辅助损失：

$${\mathcal{L}}_{\text{aux}}=\sum _{j=1}^n(\frac{1}{n}-\frac{1}{T}\sum _{i=1}^T{g}_{ij}{)}^2,\backslash mathcal\{L\}\_\{\backslash text\{aux\}\}=\backslash sum\_\{j=1\}^n\; (\backslash frac\{1\}\{n\}\; -\; \backslash frac\{1\}\{T\}\backslash sum\_\{i=1\}^T\; g\_\{ij\})^2,$$

其中 $nn$ 是专家数量，$TT$ 是 token 批处理大小，$g_{ij}istheprobabilityoftoken(ig\_\{ij\}\; is\; the\; probability\; of\; token\; \backslash \backslash (i$ 路由到专家 $jj$ 的概率。辅助损失确保 token 在专家之间均匀分布。

门控 Logit 归一化。为了提高专家辨别能力，Skywork-MoE 引入了门控 Logit 归一化：

$$z=Wx+b,z^{\mathrm{\prime }}=\lambda \frac{z-\mu }{\sigma },z=Wx+b,\; z\text{'}=\backslash lambda\; \backslash frac\{z-\backslash mu\}\{\backslash sigma\},$$

$$g=\text{softmax}(z^{\mathrm{\prime }}),g\; =\; \backslash text\{softmax\}(z\text{'}),$$

其中 $\mu \backslash mu$ 和 $\sigma \backslash sigma$ 是 $zz$ 的均值和标准差；$\lambda \backslash lambda$ 是控制输出分布“锐度”的缩放因子。这种归一化增强了门控机制区分专家的能力，减少了门控输出的熵。

自适应辅助损失系数。Skywork-MoE 采用动态方法来调整每个 MoE 层 $ll$ 的辅助损失系数 ${\alpha }^{(l)}\backslash alpha^\{(l)\}$。

$${\alpha }_{i+1}^{(l)}=\beta {\alpha }_i^{(l)}+(1-\beta )\xi d_i^{(l)},\backslash alpha\_\{i+1\}^\{(l)\}\; =\; \backslash beta\; \backslash alpha\_i^\{(l)\}\; +\; (1-\; \backslash beta)\backslash xi\; d\_i^\{(l)\},$$

其中$d_i^{(l)}d\_i^\{(l)\}$是迭代$ii$时第$ll$层的标记丢弃率，$\xi \backslash xi$是敏感度参数，$\beta \backslash beta$是平滑因子。这种适应确保了负载均衡的分布，而不会过度正则化已经均衡的层。

陷阱与教训。一方面，调整${\alpha }^{(l)}\backslash alpha^\{(l)\}$是有帮助的——有些层可能已经平衡，而其他层需要更强劲的推动来分配标记。因此，一刀切的辅助损失可能不是最优的。同时，门控逻辑归一化中的超参数调整可能很棘手。如果$\lambda \backslash lambda$设置得过高，门控概率可能会变得过于“尖锐”，迫使标记进入极端分布。过低则专家可能不够专业。

4.5 DeepSeek-V3¹⁰：基于偏置的无辅助损失策略

最后，DeepSeek-V3是最新版本，它被认为是尖端技术，因为它试图消除大量的辅助损失，并用更直接的、基于偏置的均衡方法取而代之。如果您想了解高级负载均衡，DeepSeek-V3 是一个很好的例子。

模型架构。 DeepSeek-V3 采用 DeepSeekMoE 架构用于前馈网络 (FFN)。与 GShard 等传统 MoE 架构相比，DeepSeekMoE 引入了更细粒度的专家，并将一些专家隔离为共享专家。第 $tt$ 个标记的 FFN 输出，表示为 $h{\text{′}}_{t}h\prime \_t$，计算如下：

$$h_t^l=u_t^l+\sum _{i=1}^{N_s}{\text{FFN}}_i^{(s)}(u_t^l)+\sum _{j=1}^{N_r}{g}_{j,t}\cdot {\text{FFN}}_j^{(r)}(u_t^l),h\_t^l\; =\; u\_t^l\; +\; \backslash sum\_\{i=1\}^\{N\_s\}\; \backslash text\{FFN\}\_i^\{(s)\}\; (u\_t^l)\; +\; \backslash sum\_\{j=1\}^\{N\_r\}\; g\_\{j,t\}\; \backslash cdot\; \backslash text\{FFN\}\_j^\{(r)\}\; (u\_t^l),$$

其中

$$g_{i,t}=\frac{g_{i,t}^{\mathrm{\prime }}}{\sum _{j=1}^{N_r}{g}_{j,t}^{\mathrm{\prime }}},g\_\{i,t\}\; =\; \backslash frac\{g\text{'}\_\{i,t\}\}\{\backslash sum\_\{j=1\}^\{N\_r\}\; g\text{'}\_\{j,t\}\},$$

$$g_{i,t}^{\mathrm{\prime }}=s_{i,t},\text{if}{s}_{i,t}\in \text{TopK}(\{s_{j,t}\mathrm{\mid }1\le j\le N_r\},K_r),\text{else}0g\text{'}\_\{i,t\}=s\_\{i,t\},\; \backslash ,\; \backslash text\{if\}\; \backslash ,\; s\_\{i,t\}\backslash in\; \backslash text\{TopK\}(\backslash \{s\_\{j,t\}|1\; \backslash leq\; j\; \backslash leq\; N\_r\backslash \},\; K\_r),\; \backslash ,\; \backslash text\{else\}\; \backslash ,\; 0$$

$$s_{i\mathrm{.}t}=\sigma (u_t^{\mathrm{\top }}{e}_i)\mathrm{.}s\_\{i.t\}=\backslash sigma(u\_t^\backslash top\; e\_i).$$

这里，与 DeepSeek-MoE 类似，$N_{s}N\_s$和$N_{r}N\_r$是共享专家和路由专家的数量；$K_{r}K\_r$是激活路由专家的数量；$g_{i,t}g\_\{i,t\}$是门控值；$s_{i,t}s\_\{i,t\}$表示标记与专家的亲和力；$e_{i}e\_i$是第$ii$个路由专家的质心向量；$\sigma \backslash sigma$是激活函数。

无辅助损失的负载均衡策略。传统 MoE 模型常因专家负载不均衡而导致路由崩溃，降低计算效率。传统解决方案通过辅助损失来促进均衡，但若过度强调则可能损害性能。为解决此问题，DeepSeek-V3 引入了一种无辅助损失策略，为每个专家添加偏置项$b_{i}b\_i$，以调整亲和力分数：

$$g_{i,t}^{\mathrm{\prime }}=s_{i,t},\text{if}{s}_{i,t}+b_i\in \text{TopK}(\{s_{j,t}+b_i\mathrm{\mid }1\le j\le N_r\},K_r),\text{else}0.g\text{'}\_\{i,t\}=s\_\{i,t\},\; \backslash ,\; \backslash text\{if\}\; \backslash ,\; s\_\{i,t\}\; +\; b\_i\backslash in\; \backslash text\{TopK\}(\backslash \{s\_\{j,t\}\; +\; b\_i|1\; \backslash leq\; j\; \backslash leq\; N\_r\backslash \},\; K\_r),\; \backslash ,\; \backslash text\{else\}\; \backslash ,\; 0.$$

偏置项$b_{i}b\_i$在训练期间动态更新

$$b_i\leftarrow b_i-\gamma ,\text{if expert }i\text{ is overloaded},b\_i\; \backslash leftarrow\; b\_i\; -\; \backslash gamma,\; \backslash quad\; \backslash text\{if\; expert\; \}\; i\; \backslash text\{\; is\; overloaded\},$$

$$b_i\leftarrow b_i+\gamma ,\text{if expert }i\text{ is underloaded},b\_i\; \backslash leftarrow\; b\_i\; +\; \backslash gamma,\; \backslash quad\; \backslash text\{if\; expert\; \}\; i\; \backslash text\{\; is\; underloaded\},$$

其中$\gamma \backslash gamma$是偏置更新速度。该策略确保了在整个训练过程中专家负载的均衡，而不会因辅助损失而导致性能下降。

互补的序列级辅助损失。为了防止单个序列内的极端不平衡，还采用了序列级平衡损失。

$${\mathcal{L}}_{\text{Bal}}=\alpha \sum _{i=1}^{N_r}{f}_i{P}_i,\backslash mathcal\{L\}\_\{\backslash text\{Bal\}\}\; =\; \backslash alpha\; \backslash sum\_\{i=1\}^\{N\_r\}\; f\_i\; P\_i,$$

其中

$$f_i=\frac{N_r}{K_{r}T}\sum _{t=1}^T\mathbb{I}(s_{i,t}\in \text{TopK}(\{s_{j,t}\mathrm{\mid }i\le j\le N_r\}),K_r),f\_i\; =\; \backslash frac\{N\_r\}\{K\_r\; T\}\backslash sum\_\{t=1\}^T\; \backslash mathbb\{I\}(s\_\{i,t\}\; \backslash in\; \backslash text\{TopK\}(\backslash \{s\_\{j,t\}\; |\; i\; \backslash leq\; j\; \backslash leq\; N\_r\backslash \}),\; K\_r),$$

$$s_{i,t}^{\mathrm{\prime }}=\frac{s_{i,t}}{\sum _{j=1}^{N_r}{s}_{j,t}},P_i=\frac{1}{T}\sum _{t=1}^T{s}_{i,t}^{\mathrm{\prime }}\mathrm{.}s\text{'}\_\{i,t\}=\backslash frac\{s\_\{i,t\}\}\{\backslash sum\_\{j=1\}^\{N\_r\}\; s\_\{j,t\}\},\; P\_i\; =\; \backslash frac\{1\}\{T\}\backslash sum\_\{t=1\}^T\; s\text{'}\_\{i,t\}.$$

这里，$\alpha \backslash alpha$是一个小值超参数，$\mathbb{I}\backslash mathbb\{I\}$是指示函数，$TT$表示序列长度。

动态路由和节点限制策略。DeepSeek-V3 还采用了节点限制路由机制，以减少训练期间的通信成本。每个标记最多发送到 $MM$ 个节点，这由分布在每个节点上的专家的最高 $K_r\mathrm{/}MK\_r/M$ 亲和力得分确定。这种方法保持了几乎完全的计算-通信重叠，同时提高了可伸缩性。

陷阱与教训。如果$\gamma \backslash gamma$（偏置更新速度）过大，门控可能会不稳定。如果过小，可能无法快速适应标记分布的变化。然而，这种方法可以在最小化对主要训练目标的干扰下保持负载均衡。这可以说比强硬的全局辅助项更清晰。DeepSeek-V3 示例了 MoE 的新浪潮——从大量辅助正则化转向更微妙、动态和局部修正的平衡方式。

5. 新兴趋势与观察

从 GShard 到 DeepSeek-V3 的发展历程中，一些总体趋势变得清晰：

• 门控变得更巧妙
  我们从简单的 Top-2 门控（GShard）开始，发展到单专家门控（Switch），然后探索了基于相关性、基于偏置以及更复杂的路由方式。研究人员不断寻求复杂性和效率之间的最佳平衡点。

• 重新思考辅助损失
  早期，GShard 和 Switch 等方法严重依赖辅助损失来防止专家过载。最近，一些方法（如 DeepSeek-V3）正在最小化或取消它，转而采用更直接、动态的解决方案来管理平衡。

• 容量限制和丢弃
  在 JetMoE 等“无丢弃”方法和严重依赖容量因子（Switch、GLaM）的设计之间存在一个光谱。任何一个极端都不是一刀切的解决方案；每个数据集或用例都可能以不同的方式倾斜平衡。

• 训练与推理
  训练阶段的负载均衡并不总是能解决推理阶段的瓶颈。DeepSpeed-MoE 等系统强调专门的策略（标记分组、动态节点并行）以避免推理成为一场噩梦。

• 多维并行
  流水线并行、张量并行、专家并行：HPC 现在是 MoE 的常态。我们看到了更多灵活的组合这些并行性的方式，根据层进行调整，以榨取每一丝性能。

5.1 主要 MoE 方法快速比较表

| 方法 | 路由 | 容量因子 | 核心思想 | 陷阱 | |:-------------: |------------------------------------------- |--------------------------------------------------------------- |------------------------------------------------------------------------------------------ |----------------------------------------------------------------------------------------------------- |

| GShard⁰ | Top-2 门控，局部组 | 引入容量约束概念以减少溢出 | 早期大规模 MoE (~600B 参数)，Top-2 门控，随机分发 | 过度依赖辅助损失，标记丢弃可能降低性能 |

| Switch¹ | Top-1 门控 (argmax) | 是（关键超参数） | 单专家路由，代码更简单，开销低于 Top-2 门控 | Top-1 溢出风险更大，需要仔细调整容量因子 | | GLaM² | Top-2 门控，注重能效 | 是（容量因子通常为 1.25） | 强调降低能耗（~GPT-3 训练成本的 1/3），强大的零样本性能 | 实际文本分布可能存在潜在不平衡 | | DeepSpeed-MoE³ | Top-1 门控，动态重路由 | 是（动态重新分配而非丢弃） | 多专家和多数据并行，针对训练和推理进行 HPC 优化 | 配置复杂，文本倾斜如果未仔细调整仍可能破坏负载均衡 | | ST-MoE⁴ | 带有路由器 z-损失的 Top-1 门控 | 是 | 通过 z-损失解决训练不稳定性，改进容量因子调整 | 超参数调整复杂；如果 z-损失过高或过低，可能导致不稳定或正则化不足 | | Mixtral⁵ | Top-2 门控 | 是（动态重新分配） | 观察到专家使用中的时间局部性，专用稀疏核 (Megablocks) | 如果数据分布倾斜，可能过度集中于某些专家 | | JetMoE⁸ | Top-2 门控 | 灵活的“无丢弃”方法 | 无丢弃流水线并行，无标记丢弃，块稀疏核优化 | 实现复杂性，块稀疏矩阵操作的开销 | | DeepSeek-V3¹⁰ | 带有基于偏置均衡的 Top-K 门控 | 最小化或消除大型辅助损失 | 细粒度专家 + 共享专家，节点限制路由，动态门控偏置更新 | 偏置更新速度调整可能棘手；超参数选择不当可能导致“门控抖动”风险 |

6. 陷阱与经验教训

MoE 中的负载均衡是一把双刃剑——走得太远会阻碍模型的主要目标；走得太轻，一半专家可能会闲置。以下是主要的陷阱和我们可以学到的经验教训：

• 路由崩溃与过度专业化
  如果少数专家吸收了大部分标记，你就浪费了参数。良好的门控加上温和的均衡损失（或偏置校正）可以避免崩溃。

• 容量因子调整
  设置得太高，丢弃量最小但浪费计算资源。设置得太低，会大量丢弃标记。容量因子 (CF) 的调整是一门艺术——尤其是在处理大型或倾斜数据集时。

• 过度依赖辅助损失
  强烈的平衡损失可能会掩盖语言建模的目标。平衡至关重要，但过度使用可能会阻碍专业化学习。

• 推理时瓶颈
  为训练进行的平衡并不自动转化为平衡的推理。如果在推理时某些专家负载过重，那会大大影响延迟。分层路由和动态标记分组（如 DeepSpeed-MoE）等策略会有帮助。

• 领域适应挑战
  预训练后的门控（Gating）通常会锁定某些模式。如果领域发生变化（例如，从新闻领域转向代码领域），除非你仔细重新训练或调整，否则该门控逻辑可能无法很好地适应。

7. 结论

从 GShard 到 DeepSeek-V3 的历程表明，MoE 中的负载均衡已经从一个次要问题发展成为一个核心难题。GShard 推广了 top-2 门控方法和容量限制；Switch Transformer 通过 top-1 简化了路由；GLaM 专注于能源效率；DeepSpeed-MoE 展示了训练和推理的鲁棒均衡；ST-MoE 引入了 z-loss 以提高稳定性；Mixtral 利用了时间局部性；等等——最终形成了更动态、基于偏差或基于相关性的方法，例如 DeepSeek-V3。

主要启示：完美的负载均衡是一个不断变化的目标。用力过猛，会损害模型性能；置之不理，你的超级巨型模型最终会让一半的专家闲置。我们可能会看到与 HPC 策略的进一步整合、更具适应性的门控机制，以及针对棘手的推理瓶颈的新解决方案。

随着该领域的不断发展，我们可能会看到与 HPC 技术的更多协同作用、更具适应性的门控网络，以及处理推理时间约束的新方法。对于 MoE 研究人员来说，这是一个激动人心的时代——我无疑期待着下一波突破。

感谢您的阅读，如果您有任何想法、问题或改进意见，请随时与我联系。期待下一次 MoE 冒险——愉快的门控！

参考

[0] Lepikhin, et al. **GShard: 使用条件计算和自动分片扩展巨型模型。** *arXiv 预印本 arXiv:2006.16668*, 2020。[1] Fedus, et al. **Switch Transformers: 使用简单高效的稀疏性扩展到万亿参数模型。** *arXiv 预印本 arXiv:2101.03961*, 2021。[2] Du, et al. **GLaM: 使用专家混合模型高效扩展语言模型。** *arXiv 预印本 arXiv:2112.06905*, 2021。[3] Rajbhandari, et al. **DeepSpeed-MoE: 推进专家混合模型推理和训练，助力下一代人工智能规模。** *arXiv 预印本 arXiv:2201.05596*, 2022。[4] Zoph, et al. **ST-MOE: 设计稳定可迁移的稀疏专家模型。** *arXiv 预印本 arXiv:2202.08906*, 2022。[5] Jiang, et al. **Mixtral of Experts。** *arXiv 预印本 arXiv:2401.04088*, 2024。[6] Xue, et al. **OpenMoE: 开放专家混合语言模型的早期尝试。** *arXiv 预印本 arXiv:2402.01739*, 2024。[7] Dai, et al. **DeepSeekMoE: 迈向专家混合语言模型中的终极专家专业化。** *arXiv 预印本 arXiv:2401.06066*, 2024。[8] Shen, et al. **JetMoE: 以 10 万美元达到 Llama2 性能。** *arXiv 预印本 arXiv:2404.0741*, 2024。[9] Wei, et al. **Skywork-MoE: 深入探讨专家混合语言模型的训练技术。** *arXiv 预印本 arXiv:2406.06563*, 2024。[10] DeepSeek-AI. **DeepSeek-V3 技术报告。**[11] Gale, et al. **MegaBlocks: 使用专家混合模型进行高效稀疏训练。** *arXiv 预印本 arXiv:2211.15841*, 2021。