Cerebras 发布 MoE 模型调试实战教程：从零训练小型 MoE 并修复路由梯度断裂

本视频展示了训练一个小型 MoE 模型并调试路由器问题的完整分步讲解。如果你无法打开上方显示的视频，请使用此链接在 YouTube 上打开：https://youtu.be/phXUzFt7hrs?si=6FAHkA00Tvpjz4LR

我们期望在最后看到什么

所以，我敢打赌，当你听到 Mixture-of-Experts（MoE）时，你立刻会想“又是一件只有 Google 才负担得起训练的东西”，对吧？这正是我们今天想要破除的迷思。没错，那些著名的 MoE 模型确实巨大——我们说的是万亿参数规模（Kimi Team, 2025）。但这就像因为 GPT-4 存在就避开神经网络一样。你可以在不到 20 行代码里从零创建一个感知机网络。

不幸的是，一个常见的迷思是训练 MoE 模型对大多数人来说并不真正可行。事实上，在我们撰写这篇博客时，这就是 Aman 提出的第一个问题。我们需要一个巨大的计算集群来做这件事吗？但这里有一个反直觉的事实——MoE 不仅仅关乎规模，它关乎一种你可以缩小规模来观察同样收益的技术。

在这篇博客的最后，你将学会如何训练一个小型 MoE 模型，并在 GPT-2 规模上击败一个稠密模型。至于算力——你只需要你的笔记本电脑（好吧，也许需要一个小型加速节点，但绝对不需要数据中心！）。让我们开始吧。

前置知识

如果你不熟悉稠密 GPT-2 风格的模型，我们强烈建议先观看 Andrej Karpathy 的 Let’s reproduce GPT-2 (124M)，然后再回到这篇博客，我们会在其基础上构建，并特别聚焦于 MoE。

环境搭建

让我们为大型运行做好准备！我们基于 modded-nanogpt 代码库（Jordan et al., 2024a）构建，它为我们提供了一个在 FineWeb 数据集（Penedo et al., 2024）上训练的、扎实的 124M GPT-2 风格基线。我们的模型是一个仅解码器 transformer，使用了 （So et al., 2022）、RoPE（Su et al., 2023）、RMSNorm（Zhang and Sennrich, 2019）和 QK-Norm（Henry et al., 2020），并使用 Muon 优化器（Jordan et al., 2024b）训练。它基本上是一个现代化的 GPT-2，可在 11 分钟内完成训练。如果你能使用 Cerebras 硬件，可以在那里跟着做，否则在 GPU 上也能正常运行。要开始，请下载 NanoGPT 的 MoE 分支：

在上述步骤下载 FineWeb 训练文件后，在你最喜欢的编辑器或 IDE 中打开 train_gpt_moe.py（这个文件包含我们完整的 MoE 实现，适合想立即运行实验的人）。或者，如果你更愿意跟着一步步编写代码并实现 MoE 组件，请改为打开 train_gpt.py。此外，一切都可以缩小规模（包括模型大小！），我们鼓励你动手尝试——把这篇博客当作你的起点。打破东西，尝试奇怪的设置，看看会发生什么。

专家网络住在哪里

好了，你已经拉取了 modded-nanogpt 代码库，正盯着代码看。激动人心的部分来了——让我们把在（Soboleva, 2025a）中介绍的 MoE 概念变成你真正可以训练的东西！第一步相当直接。我们需要把单个前馈网络（FFN）替换为多个专家网络。

图 1：每个专家只是一个 GPT-2 的 FFN。你可以在 train_gpt_moe.py#L184-L196 中找到专家类定义，在 train_gpt_moe.py#L263 中找到实例化逻辑。

如图 1 所示，每个专家实际上就是 GPT-2 的 FFN 直接复制粘贴（两个线性层之间夹着非线性激活）。那么，我们如何使用这些专家网络呢？图 2 展示了我们在处理批量 token 时如何组合专家的输出，这是对我们在（Soboleva, 2025a）中介绍的单 token 方法的扩展。请注意，这个实现并不快，而且会浪费计算资源，但很适合用于教学。

图 2：批处理中的专家混合（参见训练代码 train_gpt_moe.py#L208-L225）。

正如本系列第 2 部分所指出的，MoE 路由很棘手，可能会完全抵消这些模型在理论上带来的好处。在接下来的章节中，我们将保持专家架构和混合逻辑不变，只调整路由，并观察它如何改变我们模型的质量。事先提醒：我们会先把事情搞砸、犯一些错误，就像现实生活中实际发生的那样。

学习式路由

学习式路由是行业标准，它支撑着 Google、OpenAI、Anthropic 等发布的那些出色的 MoE 模型。现在你将亲手构建它！我们已经在第 2 部分介绍了理论，这里我们将聚焦于实际实现。图 3 展示了学习式路由有两个目标：第一，通过可学习权重为每个 token 找到最佳专家；第二，通过负载均衡保持所有专家都被利用。这个双重目标通常是我们容易搬起石头砸自己脚的地方，所以让我们看看启动这次训练运行后会发生什么。将辅助损失贯穿模型，对所有层求和，并以 0.01 的权重加到交叉熵损失上。

图 3：学习式路由实现：（上）用于 token-专家分配的路由器逻辑，（下）用于负载均衡的辅助损失。路由器定义见 train_gpt_moe.py#L228-L240，负载均衡逻辑见 train_gpt_moe.py#L242-L250。

是时候检查损失曲线了。图 4 显示我们的 MoE 已经击败了稠密 GPT-2 基线。恭喜！你实现了第一个能正常工作的 MoE 模型！但等等——我们提到过会搬起石头砸自己的脚，可一切却都正常工作了。那为什么还要写这份指南呢？你可能已经注意到，我们在 MoE 101 系列第 2 部分中承诺过，在 4 个专家的情况下，使用学习式路由相比稠密 GPT-2 基线大约能提升 2%。而我们现在只得到了其中一半的收益。另外 1% 去哪了？这正是大多数研究者在复现别人的论文却发现效果不如承诺时所经历的时刻。在本指南中，我们将解决这个问题。让我们一起正确地调试它。

图 4：使用 4 个专家和学习式路由的 MoE 优于稠密 GPT-2 模型。损失改善在整个训练过程中不断扩大，在我们目标的 24 亿 token 处达到 1%。要复现此图，请运行 plot.ipynb。

记住，我们的路由器有两个目标：为每个 token 找到最佳专家，并且保持所有专家都被利用。第二个目标很容易检查，我们只需查看专家利用率统计就能知道做得如何。图 5 显示每一层几乎完美均衡。双重目标中的负载均衡部分似乎工作得很好——问题不在那里。

图 5：专家利用率（跨层平均）迅速收敛到接近 25%，在总共 4 个专家的情况下，这是一个完全均衡的情形。要复现此图，请运行 plot.ipynb。

回到“为每个 token 找到最佳专家”这个可疑点。为了调试 router 在学习什么，我们可以查看梯度范数。看看图 6 中的问题。我们可以看到，router 从损失函数的交叉熵部分获得的梯度实际上为零！这意味着你的 router 只是在直接学习如何做负载均衡。而它本应学习哪个专家对每个 token 实际上最佳的那部分呢？那个信号根本没有到达 router。在你向下滚动查看解决方案之前，先看看图 3 中的代码。你能发现阻止我们的 router 从损失函数的交叉熵部分获得梯度的 bug 吗？小提示，特别考虑 top_k = 1 的情况。

图 6：top_k=1 的学习型 router 从交叉熵损失中获得的梯度为零。梯度仅来自目标函数的负载均衡部分。要复现此图，请运行 plot.ipynb.

修复学习型路由

我们刚刚了解到的事情很重要。尽管负载均衡完美，学习型路由也可能完全无用。虽然 router 在学习一些东西，但它实际上只学会了负载均衡，而不是如何路由以创建专门的专家（改善交叉熵）。不幸的是，许多 MoE 论文从未提及这个问题。它们大多关注负载均衡问题，从而在性能上留下了很大的提升空间……

好了，让我们诊断这个 bug，你可以检查自己是否猜对了。当 top_k = 1 时，我们有 gate.sum(dim=-1, keepdim=True) = gate，所以我们的归一化 gate = gate / gate.sum(...) 变成了 gate = gate / gate = 1。这意味着交叉熵梯度永远无法到达 router 权重。学习信号被阻断了！注意，这特别是 top_k = 1 的问题。然而，当 top_k = 2 时，我们可以在图 7 中看到梯度流看起来正常。为什么？因为有两个专家激活时，gate 在归一化过程中保持了相对差异，router 仍在学习哪个专家应该获得更多权重。这就像如果你问“你有兴趣申请数学学位吗？”而没有其他选项，就很难做出自信的选择。但如果你说：“你对数学学位更感兴趣，还是对物理更感兴趣？”突然之间，你对数学的 0.9 信心和对物理的 0.1 信心就大大简化了决策。

图 7：与获得零梯度的 top_k=1 情况不同，top_k=2 的学习型 router 从交叉熵损失中获得了大量梯度。要复现此图，请运行 plot.ipynb。

这暗示了我们的修复方案。对于 top_k = 1，我们将引入一个“空专家”——一个永远不会被激活的幻影专家，但它为你的路由器提供了一个与之对比预测的选项。与增加参数数量和计算量的共享专家方法不同，我们的幻影专家增加了零计算开销和零额外参数！类似于 Bachlechner 等人 2020 年解决神经网络深层梯度消失问题的方式，我们将增强进入路由器的梯度流。图 8 展示了我们如何实现它。

图 8：添加一个“空专家”为路由器提供对比对象，从而实现更好的梯度流。

null_expert_bias 是一个可调参数，经过最少调优后，我们在实验中将其设为 10。现在路由器必须证明为什么选择真实专家比幻影选项更好。这正是我们想要的对比。让我们看看这是否有助于梯度流。图 9 显示，我们终于从对偶目标的交叉熵部分接收到了非平凡的梯度量。并且它与图 7 中 top_k = 2 观察到的行为一致。“学习”已经开始了！结果，图 10 显示我们的 MoE 模型现在比密集基线好了 2%！一行修复带来了巨大收益（根据我们的经验，在更大规模下效果会更加明显）。

图 9：在 top_k=1 情况下使用“空专家”提升了从对偶目标交叉熵部分进入路由器的梯度流。要复现此图，请运行 plot.ipynb。

花点时间想想刚刚发生了什么。你拿了一个有隐蔽缺陷的 MoE 模型。学习信号完全停止了。你找到了根本原因，应用了针对性修复，然后——信号又流动起来了！我们现在几乎将相对于密集 GPT-2 基线获得的收益翻倍。注意，你不需要增加参数数量或计算量，甚至不需要超参数扫描。你只是修复了 MoE 网络的一个微小组件——路由器。它就极大地改变了模型的质量。

图 10：具有 4 个专家和带“空专家”的学习路由的 MoE 在目标 2.4B token 上比 GPT-2 模型高出 2%。“空专家”的存在几乎使相对于密集基线获得的收益翻倍。要复现此图，请运行 plot.ipynb。

我们取得的成果总结

还记得我们是怎么开始的吗？你当时可能认为 MoE 只是另一种“仅限谷歌规模”的技术。我们着手证明这是错的。我们拿了一个仅 124M 参数的微型 GPT-2 模型——那种你可以用最低计算预算训练的模型——并让它工作起来。学习路由起初并没有“学习”，然后通过一个小修复恢复了生机。我们观察到，仅仅调整路由器带来的收益就是巨大的。我们没有调优超参数，也没有在批量大小扫描上花费计算。我们只是诊断了实际损坏的地方，应用了针对性修复，并解决了它。

MoE 很棘手，因为它们可能以隐蔽的方式失败。这是一个不受欢迎的观点，但完美的负载均衡并不意味着路由器在“学习”。现在你也知道了。但 MoE 是真实的，因为当你修复实际损坏的地方时，收益是立竿见影且显著的。

我们的调试方法论（假设驱动的调查、衡量正确的指标并理解失效模式）适用范围远超 MoE。当你的模型崩溃时（它们一定会崩溃），你将知道如何让它们起死回生，而且你并不真的需要为此准备一个数据中心。

好吧，你能够训练一个 4 专家的 MoE 模型，但试着把它扩展到 64 个专家，你就会遇到 OOM 错误。在 MoE at Scale: Making Sparse Models Fast on Real Hardware 中，我们将调试为什么会发生这种情况并修复它。

引用

有问题？你可以在这里找到 Daria：https://soboleva-daria.github.io/

参考文献

Bachlechner, T., Majumder, B. P., Mao, H. H., et al. (2020). ReZero is all you need: Fast convergence at large depth. arXiv preprint arXiv:2003.04887. https://doi.org/10.48550/arXiv.2003.04887

Henry, A., Dachapally, P. R., Pawar, S., et al. (2020). Query-key normalization for transformers. arXiv preprint arXiv:2010.04245. https://doi.org/10.48550/arXiv.2010.04245

Jordan, K., Bernstein, J., Rappazzo, B., et al. (2024). modded-nanogpt: Speedrunning the NanoGPT baseline. GitHub Repository. https://github.com/KellerJordan/modded-nanogpt

Jordan, K., Jin, Y., Boza, V., et al. (2024). Muon: An optimizer for hidden layers in neural networks. Blog Post. https://kellerjordan.github.io/posts/muon/

Kimi Team (2025). Kimi K2: Open agentic intelligence. arXiv preprint arXiv:2507.20534. https://doi.org/10.48550/arXiv.2507.20534

Penedo, G., Kydličák, H., Ben Allal, L., et al. (2024). The FineWeb datasets: Decanting the web for the finest text data at scale. In Proceedings of the Thirty-Eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track. https://openreview.net/forum?id=n6SCkn2QaG

Radford, A., Wu, J., Child, R., et al. (2019). Language models are unsupervised multitask learners. OpenAI Technical Report. https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf

So, D. R., Mańke, W., Liu, H., et al. (2021). Primer: Searching for efficient transformers for language modeling. arXiv preprint arXiv:2109.08668. https://doi.org/10.48550/arXiv.2109.08668

Soboleva, D. (2025). MoE fundamentals: Sparse models are the future. Cerebras Blog. https://www.cerebras.ai/blog/moe-guide-why-moe

Soboleva, D. (2025). Router wars: Which MoE routing strategy actually works. Cerebras Blog. https://www.cerebras.ai/blog/moe-guide-router

Su, J., Lu, Y., Pan, S., et al. (2021). RoFormer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.09864. https://doi.org/10.48550/arXiv.2104.09864

Zhang, B., & Sennrich, R. (2019). Root mean square layer normalization. arXiv preprint arXiv:1910.07467. https://doi.org/10.48550/arXiv.1910.07467

讨论

我们在制作这篇博文过程中发现的其他一些零散的事情：

• 仅仅移除归一化 (gate = gate / gate.sum(...)) 并没有帮助（至少在 top_k = 1 的情况下）。尽管有来自交叉熵损失的梯度流动，但“学习”信号仍然微乎其微，导致性能比稠密 GPT-2 基线更差。
• 由于我们是在 modded-nanogpt 的基础上构建的，我们使用了 Muon 优化器。我们还基于 EleutherAI/nanoGPT-mup 做了一些运行，它使用 Adam 和一个稍旧一些的 transformer 架构。我们发现，在使用 Adam 时，仅从负载均衡损失获得梯度信号的学习路由（即存在归一化 bug 的情况）无法超越稠密 GPT-2 基线，而在使用 Muon 时它确实可以。也许 Muon 能帮助专家网络在更少的路由器监督下学习，不过我们尚未深入研究这一现象。
• 我们尝试让哈希路由正常工作——即我们在（Soboleva, 2025b）中引入的版本。令人惊讶的是，无论是使用 Adam 还是 Muon 优化器，我们都无法让哈希路由超越稠密 GPT-2 基线（除非我们只使用 4 层深的 GPT-2 风格网络）。我们没有花时间去弄清楚哈希实现哪里出了问题，所以如果你确实尝试了并找到了 bug——请告诉我们。