MoE 中的负载均衡问题

有效训练专家混合（MoE）模型带来了标准密集架构通常不会遇到的挑战。其中首要的是负载均衡问题。此问题直接源于 MoE 的核心机制：由门控网络调节的条件计算。

回想一下，在 MoE 层中，门控网络决定哪个专家处理每个输入令牌。理想情况下，我们希望计算负载在训练批次过程中大致均匀地分配到层内所有可用专家。然而，没有内在保证门控网络（仅由最小化主要任务损失驱动，如交叉熵）会实现这种理想状态。

负载不均衡是什么？

当门控网络不成比例地将令牌分配给一部分专家，导致其他专家未被充分使用时，就会出现负载不均衡。对于一个拥有 $N$ 个专家的层，理想的平衡意味着每个专家处理在给定前向传播或一个训练批次中通过该层路由的大约 $1/N$ 的令牌。严重偏离这种均匀分布即构成不均衡。

假设一个 Transformer 块包含一个 MoE 层，有 $E=64$ 个专家和一个 $T$ 个令牌的批次。门控网络 $G$ 为每个令牌 $x$ 生成选择专家 $i$ 的概率 $p_i(x)$。例如，如果使用 top-k 门控，其中 $k=2$，则每个令牌被路由到两个专家。令 $C_i$ 为批次内分配给专家 $i$ 的令牌数量。当 $C_i$ 值在 $i=1, \dots, E$ 上的分布高度倾斜时，负载不均衡就会出现。

不均衡为何会造成问题？

专家使用不均导致几个重要问题，这些问题削弱了 MoE 的优势并使训练过程复杂化：

1. 计算效率低下： 稀疏 MoE 的主要动机是计算节省；我们仅激活模型参数的一小部分来处理每个输入。如果负载不均衡，一些专家（以及在分布式设置中分配给它们的硬件资源）会成为计算瓶颈，而其他专家则闲置。这抵消了潜在的吞吐量优势，因为整体处理时间由负载最重的专家决定。

2. 参数浪费和模型容量下降： 未被充分使用的专家未获得足够的输入信号来学习有意义的专业分工。它们的参数实际上被浪费了，对模型的整体表示能力贡献甚微。模型实际运行的活跃参数少于预期，限制了其容量。

3. 训练不稳定： 接收极少令牌的专家可能出现梯度消失，导致学习缓慢或停滞。反之，持续过载的专家可能经历大而嘈杂的梯度，可能导致不稳定或发散，特别是如果未通过梯度裁剪等技术进行仔细管理。

4. 专家专业化程度低： MoE 的目标是让专家学习针对不同输入类型的专用功能。如果门控网络持续偏向少数专家，模型未能形成这种多样性。这可能导致少数“通才”专家占据主导，而其他专家未能分化，这种现象有时被称为专家崩溃。

直观展示不均衡

想象一下在一个训练步骤中令牌在专家间的分配情况。不均衡的情况可能看起来像左侧的分布，而均衡的情况则显示在右侧。

批次中8个专家间的令牌分布。不均衡情况显示出显著倾斜，专家1和4处理了大部分令牌，而其他专家几乎闲置。均衡情况显示出更加均匀的分布。

不均衡的根本原因

负载不均衡在训练过程中通常自然出现：

• 初始化： 门控网络的初始随机权重可能内在偏向某些专家。
• 早期训练动态： 初始梯度可能会强化早期建立的路由模式，形成反馈循环，导致某些专家逐渐获得更多令牌，因为它们在初始阶段表现稍好。
• 数据分布： 某些特征或模式在训练数据中可能更为普遍，自然地促使优化良好的门控网络将它们路由到特定专家，如果这些模式很常见，可能会导致这些专家过载。
• 缺乏明确的压力： 如果没有专门鼓励平衡的机制，纯粹由任务损失驱动的优化过程没有动力均匀分配负载，如果一个不均衡配置能实现稍低的任务损失。

因此，解决这个负载均衡问题不仅仅是一个优化细节；它对于成功训练大型、高性能的 MoE 模型具有根本意义。接下来的章节将介绍常用技术，特别是辅助损失函数的使用，这些技术明确旨在抵消这些趋势并促进专家得到均衡使用。