DS-MoE：让 MoE 模型更高效、内存占用更少

预计阅读时间：4 分钟

专家混合 (MoE) 语言模型以其在不牺牲性能的情况下将计算需求比传统密集模型降低 2 到 4 倍的能力而闻名。这使得它们在计算资源有限的情况下特别有用。然而，MoE 模型通常需要 2 到 4 倍的参数才能达到与密集模型相同的性能。例如，像 DeepSeekMoE-16B 和 Qwen1.5-MoE-A2.7B 这样拥有 16B 参数的模型旨在与 7B 模型的性能相匹配。MoE 模型的大量参数导致更大的 GPU 内存需求，这使得它们在 I/O 密集型场景（如自回归生成）中效率较低。

MoE 模型必须如此庞大才能实现高性能吗？我们能否创建一个既能保持性能又使用更少参数和计算能力的 MoE 模型？这就是 DS-MoE。该模型实现了与密集模型相似的性能，但使用的计算资源约为其他 MoE 模型的参数的三分之一，参数量只有一半。

DS-MoE 的概念涉及密集训练专家，并强制模型的路由器逐渐忽略给定 token 的不必要专家。我们在训练过程中采用互信息 (MI) 损失，它平衡了整个批次中每个专家的负载，但也鼓励每个输入 token 将其门控概率集中到更少的专家。

MI 损失定义为 $$L_{MI}=-H(e)+\frac{1}{\mathrm{\mid }X\mathrm{\mid }}\sum _{x\in X}H(e\mathrm{\mid }x),H(e)=-\sum _{i=1}^{N}p(e)\log (p(e)),L\_\{MI\}\; =\; -\; H(e)\; +\; \{1\; \backslash over\; |X|\}\backslash sum\_\{x\backslash in\; X\}\; H(e|x),\; \backslash quad\; H(e)\; =\; -\; \backslash sum\_\{i=1\}^\{N\}p(e)\backslash log(p(e)),$$ 其中 X 表示小批量中的 token，e 表示专家。直观地说，最大化 H(e) 可以平衡整个批次中每个专家的负载，而最小化 H(e|x) 则鼓励每个输入 x 将其门控概率集中到更少的专家。

在推理过程中，DS-MoE 仅根据其分数选择前 K 个专家。K 的数量是根据预定义值或自适应方法确定的，具体取决于得分超过某个阈值的专家数量。因此，DS-MoE 可以像类似大小的密集模型一样表现良好，同时使用的活动参数要少得多，如表中所示。

我们还使用 vLLM 对 DS-MoE 进行了测试，以了解其在 7B 性能级别与其他模型在处理速度和内存使用方面的比较。我们考察了它每秒能够处理的请求和 token 数量，设置是每个输入和输出由 1,000 个 token 组成，GPU 内存使用率上限为 90%。

测试表明，DS-MoE 在计算成本方面优于密集模型，在模型内存方面优于稀疏训练的 MoE 模型，从而在计算密集型和 I/O 密集型场景中都能实现更快的处理。请注意，DS-MoE-6B 在下游性能方面尚无法与其他模型相媲美，因为它仅在 1000 亿个 token 上进行了训练（而其他模型则在数万亿个 token 上进行了训练）。然而，DS-MoE 已经展示出在可比训练数据量下达到密集模型性能水平的巨大潜力。

阅读论文了解更多