MoE推理的批处理策略

标准批处理技术对密集模型有效，但在直接应用于混合专家（MoE）模型进行推理时，会遇到一些难题。主要问题源于MoE固有的条件计算：同一输入批次中不同的令牌根据门控网络的决策被路由到不同的专家。这种动态的令牌级路由打乱了传统批处理所依赖的计算一致性。

批处理中的稀疏激活问题

假设一个标准的Transformer推理场景。一批输入序列逐层处理。在每一层中，批次中的所有令牌都进行相同的计算（例如，自注意力，前馈网络）。这种同质性使得GPU等硬件能够进行高效的并行处理，从而提升吞吐量。

然而，在MoE层中，门控网络之后路径就变得不同了。对于包含$B$个长度为$L$的序列的批次，$B \times L$个令牌会通过门控网络。每个令牌随后被分配给一个或多个专家（通常是Top-k，推理时常为k=1或k=2）。如果有$N$个专家，最初按序列位置分组的令牌现在逻辑上分散到这$N$条计算路径上。

一种简单的批处理方法，即简单地将输入批次送入MoE层，会导致多方面的效率低下：

1. 专家利用不足： 某些专家可能从给定批次中接收到很少的令牌（甚至没有），而另一些专家则可能过载。这会导致硬件利用率低下，因为分配给空闲专家的处理核心处于闲置状态。
2. 负载不均衡： 即使专家并非完全空闲，每个专家处理的令牌数量也可能显著不同。这种不均衡意味着MoE层的总时间由负载最重的专家决定，从而抵消了潜在的加速效果。
3. 延迟增加： 等待最慢的专家路径会主导该层的执行时间。

当专家分布在多个设备上时（专家并行），这些问题尤为明显。简单地处理批次将需要低效、稀疏的通信模式，或者导致设备之间严重的负载不均衡。

高效MoE推理批处理的应对方法

为了解决这些困难，MoE的推理批处理需要能明确处理令牌动态路由的方法。主要目的是在门控决策之后但在专家计算之前重新组合令牌，确保每个专家处理一个分配给它的、密集且规模合适的令牌批次。

动态批处理（请求级）

动态批处理是服务系统中常用的一种通用技术，它会将传入的推理请求进行缓冲和分组，形成更大的批次，然后由模型处理。虽然通过提高硬件利用率有助于提升整体系统吞吐量，但它本身并不能解决MoE特有的批内路由分歧问题。它增加了共同处理的令牌总数，这与单请求处理相比，可以在统计上改善专家负载平衡，但不能保证动态形成的批次内的负载均匀分布。它通常与更针对MoE的特定方法结合使用。

令牌级分组和置换

这是高效MoE推理的基本方法。它涉及到根据令牌分配到的专家，主动重新排列批次内的令牌。工作流程通常如下所示：

1. 门控计算： 门控网络处理传入批次中的所有令牌（可能通过动态批处理形成），以确定每个令牌的专家分配。
2. 路由决策与排序： 确定每个令牌的目标专家。通常，为了简化和加速，推理时只使用Top-1专家，但也可能采用Top-2路由。令牌随后根据其分配的专家ID进行排序或索引。
3. 令牌置换（Gather操作）： 令牌在内存中进行物理重新排列（或逻辑寻址），使得所有分配给专家1的令牌连续排列，接着是分配给专家2的令牌，依此类推。如果专家分布在不同设备上，此步骤会涉及 All-to-All 通信模式，类似于训练过程，每个设备将发往远程专家的令牌发送出去，并接收分配给其本地专家的令牌。高效的实现通常使用优化的集体通信库（例如，NVIDIA GPU的NCCL）或专门的框架，如Tutel。
4. 批处理专家计算： 每个专家（或承载专家的设备）现在接收到一个专门分配给它的密集小批次令牌。它针对这组已分组的令牌高效地执行其计算（例如，前馈网络）。
5. 令牌反置换（Scatter操作）： 已处理的令牌表示必须重新排列回批次中的原始序列顺序。这与置换步骤相反，如果采用分布式设置，可能涉及另一次 All-to-All 通信。
6. 组合与继续： 专家的输出被组合起来（通常是基于路由器 logits 的加权和，即使计算时只使用了 Top-1），随后的批次进入模型的下一层。

以下图表展示了MoE推理中令牌置换的流程：

MoE层在推理过程中使用令牌级分组和置换的令牌处理流程。令牌经过路由、按专家分组、处理，然后重新组合。

专家容量管理

在训练期间，通常会定义一个expert_capacity（专家容量），通常capacity_factor > 1.0，以处理暂时的不均衡并留有余地。在推理时，此容量仍有作用。如果批次内分配给特定专家的令牌数量超出其定义的容量（令牌数量 / 专家数量 * capacity_factor），则令牌可能会被丢弃。

虽然在训练中（并通过辅助损失进行管理）有时允许丢弃令牌，但在推理时通常不希望这样做，因为它会导致信息丢失和输出质量下降。处理推理时潜在溢出的方法包括：

1. 足够的容量： 确保为推理配置的容量足够大，以处理典型批次大小的预期峰值负载。这可能需要对典型路由模式进行性能分析。
2. 填充： 如果某个专家接收到的令牌少于其容量，则可以填充其批次以保持计算一致性。这在许多MoE实现中是标准做法。
3. 不丢弃： 配置系统以避免丢弃令牌，这可能通过动态增加容量或在专家暂时过载时接受更高的延迟（尽管这会使实现复杂化）。使用 Top-1 路由比 Top-2 简化了容量管理。

权衡与考虑因素

令牌级分组通过最大化专家利用率并有效利用硬件并行性，显著提升了吞吐量。然而，它会带来额外开销：

• 延迟： 排序、置换和反置换步骤（尤其是在分布式设置中的 All-to-All 通信）会给每个MoE层增加延迟，与同等的密集模型相比。这是吞吐量与延迟的直接权衡。
• 实现复杂度： 需要仔细管理令牌索引、高效的置换核函数，并可能需要与专用库（例如 Tutel, DeepSpeed）进行集成。
• 内存： 在置换和反置换阶段，需要缓冲区来存储令牌。

选择合适的批处理方法需要根据特定应用的需求（例如，延迟敏感的实时推理与侧重吞吐量的批处理）和部署环境（单GPU、多GPU节点、多节点集群）来平衡这些因素。有效的批处理不仅仅是一种优化，它对于MoE模型实现实用的推理性能是必不可少的。