基于 Megatron-Core 的稀疏大模型训练工具：阿里云MoE大模型最佳实践

概述

随着大模型技术的不断发展，模型结构和参数量级快速演化。大模型技术的应用层出不穷。大模型展现惊人效果，但训练和推理成本高，一直是巨大挑战。模型稀疏化能降低计算和存储消耗。近期以Mixtral为代表的MoE（多专家混合）大模型证明了稀疏MoE技术能大幅降低计算量、提升推理速度，模型效果甚至超过同规模稠密模型。阿里云PAI和NVIDIA团队深入合作，基于Megatron-Core MoE框架，解决了MoE大模型训练落地时会遇到的可拓展性、易用性、功能性以及收敛精度等核心问题，在下游任务上取得了很好的模型效果。PAI团队将上述MoE训练框架和技术与阿里云AI平台产品深度整合，使云上大模型用户能方便地进行MoE大模型的训练和部署。

基于 Megatron-Core 的 MoE 训练工具链

这一章节中，我们从MoE算法介绍，Megatron-Core MoE训练框架、以及云平台工具三个方面了解支撑MoE落地背后的基础能力。

MoE 算法介绍

MoE的全称是Mixture of Experts，其中的Expert对应的是Transfomrer模型的MLP层，在训练的时候从多个MLP中选取一个MLP进行激活 [1]（如下图所示）。这意味着模型可以在不增加FLOPs的情况下，通过增加MLP模块的数量来增加模型参数量级，进而提升模型在下游任务上的效果。采用MoE后的稀疏Transformer模型和同等质量（验证集loss以及Zero-shot NLU下游任务性能）的稠密模型相比有较大幅度的训练和推理吞吐性能提升。MoE主要由以下两个关键部分组成：

• 稀疏MoE层: 这些层代替了传统 Transformer模型中的前馈网络 (FFN) 层。MoE层包含若干experts组成，每个experts本身是一个独立的神经网络。在实际应用中这些专家通常是前馈网络 (FFN)。
• Router路由: 这个部分用于决定哪些tokens被发送到哪个专家。例如在下图中，“More”这个tokens可能被发送到第二个专家，而“Parameters”这个tokens被发送到第一个专家。

接着我们讨论下MoE中的token负载均衡问题。如上图所示的Top-1 Routing算法，每个输入hidden_states被router分配了4个experts中的1个用于前向推理计算，然后其输出被加权求和汇总后送入Transformer的下一个处理单元。这种路由方式会出现少数受欢迎的experts需要处理很多token，而其他experts仅需处理极少数量的token的情况，这导致处理极少token的那些experts无法获得足够多的信息，从而导致训练效率大幅降低。

为了缓解这种token到expert负载不均衡(Load Imbalancing)的问题，可以引入了一个辅助损失函数，旨在鼓励给予所有专家相同的重要性。这个损失函数确保所有专家接收到大致相等数量的训练样本，从而平衡了专家之间的选择。另外也可以通过drop tokens的方式得到缓解。首先定义一个expert capacity，即为一个expert被分配到的token的容量。如果分配到一个expert的token数超出了一个固定容量，则多余的tokens会被丢掉。这些被丢掉的tokens不参与和experts的矩阵乘运算，直接通过一个残差连接进入到下一个处理单元。如果一个expert没有被分配到其容量上限内的足够多的tokens，则需要采用padding的方式填充到容量上限。

Mixtral-8x7b模型提供了另外一种思路来缓解负载不均衡的问题，其采用的Megablocks [2]论文中提出的dropless moe算法。如下图所示，首先tokens通过router被分配到具体的experts，token到expert的分配关系被存储到expert indices中，分配置信度被存储到Probabilities中。其次根据预先设置好的expert的容量将tokens分组放置，超出容量的tokens被丢掉。接着将分组中的tokens和experts进行批量矩阵乘运算。最后输出结果按照分配置信度进行加权求和（如下左边公式)，其分配置信度的计算采用的是normalized softmax（如下右边公式）。

Megablock论文认为这个drop token方案会导致效果下降，于是在此基础上引入块稀疏算子来实现dropless功能。如下图所示，dropable的moe算法在experts之间并行批量矩阵乘（如下图A所示）可以用一种块对角矩阵乘来等价表达（如下图B所示）。Megablock设计了一个新的算子允许变长的输入（如下图C所示），然后输出到一个稀疏对角矩阵上。引入了这个变长输入就意味着不需要drop的方式就能处理超出容量的tokens。类似的，Megatron-Core MoE使用GroupedGEMM解决了多 expert变长输入的问题。