专家容量与规模考量

选择合适的专家数量、它们各自的尺寸以及分配给每个专家的处理能力，是设计高效专家混合（MoE）模型的根本。这些选择直接影响模型的参数数量、计算需求（浮点运算次数）、分布式环境中的通信开销，并最终影响其学习专用功能的能力。恰当设置这些超参数，需要平衡模型性能、训练稳定性与计算效率等多方面因素。

理解专家容量

在MoE模型设计中，一个主要理念是专家容量。由于计算通常是批处理和并行化的，特别是在GPU等加速器上，我们需要为每个专家在给定计算批次（例如，特定设备组中的每次前向传播）内可以处理的令牌数量定义一个固定缓冲区大小。这个缓冲区大小就是专家的容量，表示为$C$。

容量通常通过**容量系数（CF）**来确定，这是一个通常大于1的超参数。此计算确保容量略大于令牌在专家之间完美分配时所需的容量。对于正在处理的$T$个令牌组以及该组可用的$N$个专家，每个专家的容量$C$通常计算为：

$$C = \text{取整}\left( \text{CF} \times \frac{T}{N} \right)$$

如果路由机制在处理组内分配给某个特定专家的令牌超过$C$个，则超出容量的令牌被视为“丢弃”。丢弃的令牌通常会跳过该层的专家计算，这意味着它们的表示会直接通过残差连接传递。这会导致信息损失。

设置容量涉及直接的权衡：

• 低容量（低CF，例如CF接近1.0）：
  • 优点： 减少计算成本和内存占用，因为专家计算张量需要更少的填充。
  • 缺点： 增加丢弃令牌的可能性，特别是当路由不平衡时。这会妨碍学习并降低模型质量。需要良好的负载平衡机制（在第三章中讨论）才能运行良好。
• 高容量（高CF，例如CF >= 2.0）：
  • 优点： 最大限度地减少丢弃令牌，提供更稳定的训练信号，并可能带来更高的模型质量。对路由不平衡更具稳定性。
  • 缺点： 大幅增加计算成本和内存使用。如果容量远高于实际路由到大多数专家的令牌数量，计算资源会浪费在填充令牌上。在极端情况下，计算量接近密集层。

最佳CF值通常通过经验确定，旨在平衡可接受的丢弃令牌百分比与计算预算。1.25或1.5等值是常见的起始点。

增加容量系数可以减少丢弃的令牌，特别是在负载平衡不理想的情况下。然而，较高的CF会由于填充而增加计算开销。

确定专家数量（N）

专家总数$N$是一个主要的架构选择。增加$N$可带来以下益处：

• 更多参数： MoE层能高效地扩展模型规模。增加专家会使得总参数量大幅增加，但如果$k$（每个令牌路由到的专家数量）保持较小，每个令牌的计算成本（浮点运算次数）只会适度增加。
• 更精细的专业化： 更多专家可能使得模型能够学习针对不同类型输入或上下文的更独特、专业的功能。
• 硬件映射： 在使用专家并行（第四章）的分布式环境中，$N$通常根据可用处理设备的数量（例如GPU）选择，每个设备分配一个或多个专家。大型模型中，$N$的常见值包括8、16、64、128甚至更多。

然而，增加$N$也带来以下挑战：

• 通信开销： 专家并行需要全局通信（All-to-All，第四章）以在设备之间混洗令牌。更多专家通常意味着更多的通信开销，这可能成为一个瓶颈。
• 负载平衡难度： 确保更多专家都能得到良好使用，对路由器和辅助损失函数（第三章）来说可能更具挑战性。
• 收益递减： 在某个时点，如果数据不能自然地分解成那么多独特的专业，或者路由器难以良好分配令牌，增加更多专家可能不会带来明显的模型质量提升。

单个专家规模

与专家数量互补的是它们各自的尺寸，通常指的是专家网络中的隐藏维度（例如，标准Transformer前馈网络中的中间维度）。

• 专家尺寸与密集层等效物： MoE层替换单个密集FFN块。如果密集FFN具有中间维度$d_{\text{ff}}$，且MoE层有$N$个专家，每个专家具有中间维度$d_{\text{expert}}$，则总参数量大致与$N \times d_{\text{expert}}$成比例。您可能会选择$d_{\text{expert}}$使其小于$d_{\text{ff}}$，但总参数（$N \times d_{\text{expert}} \times d_{\text{model}}$）可能远大于原始密集层的参数（$d_{\text{ff}} \times d_{\text{model}}$）。
• 计算成本： 每个令牌的浮点运算次数取决于实际激活的专家大小。对于$k=1$的top-k路由，浮点运算次数与$d_{\text{expert}}$相关，而不是$N \times d_{\text{expert}}$。
• 与N的权衡： 对于固定的参数预算或计算目标，存在以下权衡：
  • 更少、更大的专家： 每个专家有更强的能力来建模复杂功能，可能需要较少细致的路由。负载平衡可能更容易。
  • 更多、更小的专家： 促进更高的专业化，但每个专家单独的能力较弱。更依赖路由器识别非常具体的场景。

这种选择通常取决于任务和数据的特性，以及关于实现复杂性和分布式训练性能的实际考量。

与Top-k路由的相互影响

大多数现代MoE实现使用top-k路由，其中门控网络为每个令牌选择前$k$个专家（通常$k=1$或$k=2$）。这种选择与容量考量有着很大的关联：

• k=1： 每个令牌路由到一个专家。批次/组中的令牌分配总数为$T$。总可用容量为$N \times C$。负载平衡旨在将$T$个分配大致均匀地分配给$N$个专家。
• k=2： 每个令牌路由到两个专家。令牌分配总数为2T。总可用容量仍为$N \times C$。这对系统施加了明显更大的压力。为避免过度令牌丢弃，您可能需要：
  • 更高的容量系数（CF）。
  • 通过辅助损失进行更良好的负载平衡。
  • 相对于令牌批次大小（T），更大的专家数量（N）。

使用$k=2$有时可以通过让令牌从多个专业功能中受益来改善模型质量，但这会增加计算量（每个令牌激活两个专家），并可能带来更高的通信和容量需求。

实用指导与监控

在开发和训练期间需要关注的重要指标包括：

1. 丢弃令牌百分比： 一个主要指标，表明容量是否充足。高值（>1-2%）通常预示着存在问题。
2. 专家使用率 / 负载平衡： 监控分配给每个令牌的专家数量的变异系数（CV）等指标，或专家分配的可视化。这有助于诊断路由器的异常行为并调整负载平衡损失（第三章）。
3. 计算成本（浮点运算次数）和训练吞吐量： 衡量不同配置的实际性能影响。
4. 整体模型性能： 跟踪困惑度（用于语言模型）或下游任务准确性等标准指标。确保架构选择能转化为更好的最终性能。
5. 内存使用： 在分布式环境中特别重要，确保所选配置符合设备内存限制。

通过系统地调整$N$、专家大小和CF，并观察这些指标，您可以为您的特定应用和硬件环境，找到一个能够平衡性能、专业化和计算可行性的MoE架构。