先这样答
MoE 的负载均衡解决的是门控网络把 token 全分给少数专家的专家塌缩问题,专家并行则是为了解决单卡存不下海量专家参数的问题,将不同专家分置于多张 GPU 上。
传统的负载均衡依赖辅助损失函数。Switch Transformer 通过计算各专家分配 token 比例和路由概率的乘积和,分布越不均惩罚越重,推动路由均匀分发。ST-MoE 为了稳定训练还引入了 Router z-loss 来惩罚过大的 logits,曾帮助 ST-MoE-32B 在 SuperGLUE 上取得 91.07 的分数。但辅助损失会干扰主损失梯度,影响模型质量。目前的无辅助损失方案是在路由打分上加一个动态调整的偏置,这个偏置只影响 top-k 选择,不参与门控权重计算,从而实现无损均衡。
专家并行的切法是把不同专家放到不同 GPU 上,每张卡只存部分专家。前向传播时,token 经过路由计算后,通过一次 all-to-all dispatch 通信发送到目标专家所在的卡上,专家计算完成后,再通过一次 all-to-all combine 通信把结果收回原卡。专家并行天然和数据并行耦合,非专家参数照常走数据并行,实际部署常和张量并行混合使用。
整体来看,负载均衡保证了专家参数能被有效训练,而专家并行配合专用通信库解决了这些参数的物理存放与计算调度问题。
面试官会怎么追问
-
「如果不做任何负载均衡,模型训练会出现什么现象?」 会出现专家塌缩。门控网络会倾向于把绝大多数 token 路由给少数几个明星专家,导致这些专家过载,而其余专家几乎拿不到 token 和梯度,越训越废。这等价于退化成一个浪费了大量参数的普通稀疏网络。
-
「辅助损失函数有什么代价?现在去掉它的主流做法是什么?」 辅助损失函数会干扰主任务的梯度,为了拉平专家负载而牺牲了模型最终的生成质量。去掉它的做法是引入动态偏置机制,给每个专家分配一个根据负载动态调整的偏置值。这个偏置只参与决定 token 去哪个专家的 top-k 选择过程,不进入最终的门控权重相乘环节。
-
「专家并行和张量并行有什么区别?为什么MoE推理要用它?」 张量并行是切分单层算子的矩阵乘法,解决单层参数过大的问题,通信用 all-reduce。而 MoE 的专家参数总量非常大,张量并行和数据并行都覆盖不了这个维度。专家并行按专家粒度切分到不同卡,代价是引入了对网络带宽敏感的两次跨卡通信,因此通常需要专用内核来降低通信延迟。
回答的坑
- 把专家并行等同于数据并行或张量并行。要明确专家并行是 MoE 专属的并行维度,它切的是专家个数,且必须伴随两次跨卡分发与合并通信,非专家参数依然走数据并行。
- 认为负载均衡只能靠修改损失函数。要补充说明除了辅助损失函数,还可以通过在路由打分阶段加动态偏置来解决,这是目前避免损失模型质量的重要考点。
同系列的题