当前位置：首页 > ops >正文

【大模型推理加速】MOE加速比与batchsize 关系

ops 2025/6/27 3:35:53

https://zhuanlan.zhihu.com/p/18788701543
在这里插入图片描述

您的查询描述了Mixture of Experts (MoE) 模型在推理过程中不同批次大小（Batch Size, BS）下的性能特征，特别是在小批次、中间区域和大批次三种情况下的行为。我将针对您的问题“为什么大批次时MoE吞吐量更高”进行详细解释。首先，简要回顾MoE模型的核心特性：

MoE模型原理：MoE模型由多个专家（experts）组成（如前馈神经网络层），每个输入token通过路由机制被分配到少数专家（通常k=1或2）。相比之下，稠密模型（dense model）对所有token使用相同的参数。MoE的优势在于，它可以通过增加专家数量来扩展模型规模（总参数量更大），但每个token只激活部分专家，从而减少每个token的计算量（FLOPs）。
关键性能指标：
- 延迟（Latency）：处理一个批次所需的时间。
- 吞吐量（Throughput）：单位时间内处理的token数量（例如，tokens/秒）。
- 瓶颈：推理性能受限于内存带宽（memory bandwidth）或计算量（compute-bound）。内存带宽限制时，时间主要花在从内存（如DRAM）读取参数；计算量限制时，时间主要花在执行浮点运算（FLOPs）。

现在，针对您描述的三种情况，我先快速总结，然后重点解释大批次时MoE吞吐量更高的原因。

行为：当批次大小较小时，推理延迟主要取决于“读入所有激活参数的时间”。这是因为每个token可能被路由到不同专家，导致参数访问分散，无法有效利用缓存。内存带宽成为瓶颈（即硬件等待数据加载，而非忙于计算）。
原因：小批次时，计算单元（如GPU核心）经常空闲，等待参数从内存加载。MoE模型可能比稠密模型更差，因为路由增加了额外的参数访问开销（不同专家的参数位置不连续）。

行为：随着批次大小增加，所有专家被激活（即每个专家都处理至少一个token），但相比同等规模的稠密模型（例如，总FLOPs相当的稠密模型），MoE没有性能收益（延迟或吞吐量无改善）。
原因：此时，计算量接近稠密模型（所有专家都活跃），但引入了路由开销（如选择专家的softmax计算）、负载不平衡（某些专家处理更多token，导致等待）或同步成本。这些开销抵消了MoE的潜在优势，而稠密模型由于参数共享更均匀，效率更高。

行为：当批次大小继续增加，瓶颈从内存带宽转为“单token激活计算量”（即计算成为主要限制）。此时，MoE的吞吐量显著高于同等规模的稠密模型。
关键问题：为什么大批次时MoE吞吐量更高？
原因在于大批次下，硬件（如GPU）进入计算绑定（compute-bound）状态，而MoE模型每个token的计算量（FLOPs per token）更低，且参数访问开销被有效摊销。具体机制如下：

MoE模型的核心优势是每个token只激活少数专家（例如，k=2），而稠密模型每个token必须通过全部参数。假设一个MoE模型有E个专家，每个专家的计算量较小（例如，专家FFN层的隐藏尺寸d_ff_expert < 稠密模型的d_ff_dense），那么：
- 稠密模型的每token计算量：~ ( O(d_model \times d_ff_dense) )（例如，一次大型矩阵乘）。
- MoE模型的每token计算量：~ ( O(k \times d_model \times d_ff_expert) )（例如，k个小型矩阵乘）。
  通常，MoE设计为 ( k \times d_ff_expert \ll d_ff_dense )，因此每token FLOPs更低（可能减少2-4倍）。
在计算绑定区域，硬件计算能力（如GPU的TFLOPS）是固定的。吞吐量公式为：
[
\text{Throughput} = \frac{\text{Hardware Compute Capacity (FLOPS)}}{\text{FLOPs per Token}}
]
由于MoE的FLOPs per token更低，它能处理更多token/秒。例如：
- 硬件计算能力：100 TFLOPS。
- 稠密模型：每token需10 GFLOPs → 吞吐量 = 100e12 / 10e9 = 10,000 tokens/秒。
- MoE模型：每token需4 GFLOPs → 吞吐量 = 100e12 / 4e9 = 25,000 tokens/秒（提升2.5倍）。

大批次时，许多token被路由到相同的专家（例如，一个专家处理批次中的多个token），允许专家参数在缓存（如GPU共享内存或L2缓存）中重用：
- 参数加载一次后，可多次用于计算，减少平均每个token的内存访问次数。
- 而小批次时，token路由分散，参数无法有效重用，内存访问频繁。
稠密模型在大批次下也有参数重用（所有token共享参数），但由于其每token计算量高，计算瓶颈更早出现，吞吐量上限较低。
因此，大批次时MoE的内存访问开销被大幅摊销，瓶颈从内存带宽转向计算，MoE的低FLOPs优势得以发挥。