容量因子及其对性能的影响

专家混合模型中的门控网络动态地将每个令牌路由到选定的专家。这种动态的、依赖数据的路由带来了工程上的挑战：虽然路由灵活，但GPU计算在操作具有固定、可预测形状的张量时效率最高。如果一个专家突然被分配了1,000个令牌而另一个只被分配了10个，我们不能简单地在单个批次内动态分配不同数量的内存和计算。

为了解决此问题，MoE的实现会为每个专家预分配一个固定大小的缓冲区，或称“容量”。每个专家都被配置为在每个批次中处理特定最大数量的令牌。容量因子是控制此缓冲区大小的超参数。它直接决定了计算浪费和信息损失之间的权衡。

每个专家的容量是相对于令牌的完全均匀分布计算的。如果你的批次中有 $N$ 个令牌， $E$ 个专家，一个完全平衡的系统会将 $N/E$ 个令牌发送给每个专家。容量因子 $C$ 是应用于这个理想平均值的乘数。

专家的令牌容量公式为：

\text{专家容量} = \text{int}\left( \frac{\text{令牌数量}}{\text{专家数量}} \times C \right)

容量因子 $C=1.0$ 意味着每个专家缓冲区可以恰好容纳平均数量的令牌。而 $C=1.25$ 的因子提供了25%的缓冲区，允许每个专家处理比批次平均多达25%的令牌。

容量因子的选择会造成直接且不可避免的权衡。因为路由是动态的，在给定批次中，一些专家不可避免地会比其他专家更受欢迎。这会导致两种对立的结果。

低容量因子 (例如 $C=1.0$ )： 如果一个专家的缓冲区太小，当它接收到的令牌超出其处理能力时会溢出。缓冲区满后到达的令牌被视为“丢弃”的。这些丢弃的令牌不会被专家处理。相反，它们直接通过MoE层的残差连接。这节省了计算资源，但以模型性能为代价，因为模型失去了对这些令牌应用专业知识的机会。
高容量因子 (例如 $C=2.0$ )： 大型缓冲区显著降低了丢弃令牌的风险。然而，它强制系统为很少使用的容量分配内存和计算资源。如果一个专家被分配了200个令牌的容量但只接收了50个，那么150个未使用的槽位的计算和内存就被浪费了。这种填充增加了训练成本并降低了执行速度。

下图展示了此过程。输入令牌被路由到三个专家，每个都有固定容量。专家2溢出，导致一个令牌被丢弃。专家3未充分利用，导致容量浪费（填充）。

固定容量下令牌路由的图示。专家2接收三个令牌但只能处理两个，导致一个令牌被丢弃。专家3只接收一个令牌，留下一个槽位作为计算浪费。

容量因子、丢弃令牌和计算浪费之间的关系在绘图时很清晰。随着容量因子的增加，丢弃令牌的百分比急剧下降，但浪费的计算量线性增加。

容量因子直接控制了模型质量（更少的丢弃令牌）与计算效率（更少的浪费）之间的平衡。

容量因子没有一个“最佳”值；它是一个重要的超参数，取决于你的具体模型、数据和训练目标。然而，这里有一些常见做法：

从常见基线开始： 大多数研究和生产系统使用介于1.25和2.0之间的容量因子。对于许多应用， $C=1.25$ 是一个合理的起始点。
监控丢弃的令牌： 在训练期间，记录每个批次中丢弃令牌的百分比。高百分比（例如，持续高于1%）表明你的模型学习受到影响。这表明应增加容量因子。
分析计算浪费： 如果你的丢弃令牌率几乎为零但训练速度慢，你的容量因子可能不必要地高。分析工具可以帮助你衡量填充量，这代表着浪费的计算。降低容量因子可以提高吞吐量。

容量因子并非独立运行。它的效果与第1章中讨论的辅助负载均衡损失紧密关联。设计良好的负载均衡损失鼓励门控网络更均匀地在专家之间分配令牌。

最终，管理容量因子是训练MoE模型的一项重要技能。它需要仔细调整和监控，以找到在最大化模型质量的同时尊重计算预算的正确平衡。

这部分内容有帮助吗？

参考文献

Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc Le, Geoffrey Hinton, Jeff Dean, 2017 International Conference on Learning Representations (ICLR) DOI: 10.48550/arXiv.1701.06538 - 这篇基础论文介绍了稀疏门控混合专家（MoE）层，详细阐述了专家容量的概念以及用于确保均匀令牌分布的辅助负载平衡损失。
GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding, Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, Zhifeng Chen, 2021 International Conference on Learning Representations (ICLR), Vol. 139 (OpenReview) DOI: 10.5555/3524938.3525287 - 描述了谷歌大规模实现 MoE 的方法，提供了在实际场景中管理专家容量、高效处理令牌路由和防止令牌丢弃的实用细节。
Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, William Fedus, Barret Zoph, Noam Shazeer, 2021 International Conference on Learning Representations (ICLR) DOI: 10.48550/arXiv.2101.03961 - 介绍了 Switch Transformer 架构，讨论了专家容量的实际影响以及高效路由和负载平衡对于将稀疏 MoE 模型扩展到超大规模的重要性。