专家模型不要专家并行！微软开源MoE新路径

近日，来自微软的研究人员开源了使用全新方法训练的MoE大模型，不走寻常路，且编码和数学表现出色。

继Phi家族之后，微软又开源了新的混合专家大模型——GRIN MoE。

与Phi-3.5同样的个头（16 * 3.8B），却采用了截然不同的训练方法。

这个「不走寻常路」如果写个太长不看版，那就是两句话：

1. 使用新一代SparseMixer来精确估计专家路由的梯度，解决传统方案中利用门控梯度代替路由梯度的问题。 2. 专家并行不要了，训练中改用数据、pipeline和张量并行，避免了传统方法丢弃token的问题。

论文地址：https://arxiv.org/abs/2409.12136

当然了，上面两句话是小编说的，多少有点糙，文中细节，还请诸君继续阅读~

这年头，新来一个LLM，当然要先刷分了——

参数要少，效果要好，所以要在左上角：

GRIN作为MoE架构，总参数量约42B，推理时激活的参数为6.6B，打同级别（7B）的非MoE模型是手拿把攥，甚至比14B的Phi-3还要略胜一筹。

在上面的这份成绩单中，GRIN MoE表现优异，尤其是在编码和数学测试中。

比如，在衡量数学问题解决能力的GSM-8K中，GRIN MoE得分为90.4，而在编码任务基准HumanEval上拿到了74.4分。

在MMLU（大规模多任务语言理解）基准测试中GRIN得分为79.4，超过了同为MoE架构的Mixtral（70.5分），以及自家的Phi-3.5（78.9分）。

如果对比流行的商用模型，GPT-3.5表示感受到时代的力量，默默退出群聊。

开放权重：https://huggingface.co/microsoft/GRIN-MoE

demo：https://github.com/microsoft/GRIN-MoE

MoE全新训练路径

GRIN MoE由常规的Transformer块构成，采用分组查询注意力（GQA）和滑动窗口注意力来提高计算效率。

采用RoPE进行位置编码，以便在预训练后实现长上下文能力。

在MoE架构中，模型通过路由网络为每个输入token挑选适合的专家模块。对于有n个专家的网络，一个用于推理的MoE模块的输出为：

其中z = Router（x，r），本文中Router采用线性网络，Gating是门控函数（通常为softmax），Expert是FNN层。

MoE通过TopK函数进行专家分配，这个专家路由的过程是不可微的，所以反向传播的时候没法求导。

对此，传统的MoE训练将TopK视为常数，仅通过Gating来反向传播计算路由权重梯度，相当于用门控的梯度代替了路由的梯度。

这多少有点糙。

不可导怎么办

恰好，本文一作之前有一篇工作（SparseMixer）：

论文地址：https://arxiv.org/pdf/2310.00811

受到直通梯度估计器的启发，作者扩展了前作，提出了SparseMixer-v2。

作者首先将TopK函数替换为模型训练中离散变量的随机采样，然后应用heun’s third order method来近似专家路由梯度，并构建一个改进的反向传播，为专家路由给出数学上合理的梯度估计。

前作中，SparseMixer的有效性在神经机器翻译任务和ELECTRA语言模型训练中得到了证明。

而在GRIN MoE的开发过程中，SparseMixer-v2终于有机会大规模应用于自回归语言模型训练。

作者用2.5T token训练了两个16×0.9B MoE。其中一个遵循GRIN MoE中使用的相同方案，另一个用传统的GShard方法替换 SparseMixer-v2。

如上图所示，将SparseMixer-v2的性能提升推广到16×0.9B尺度的自回归语言模型训练。

在前0.5T token上GShard表现更好，但SparseMixer-v2在训练后期取得了更强的性能。

专家模型不要专家并行

传统的MoE训练采用专家并行，简单理解就是把不同的专家分配到不同的显卡上。

一个明显的问题是负载不均衡，有的专家会分到更多的token，有的专家却很闲。

之前的做法是设定一个阈值，比如1000个token分给4个专家，每人应该是250，这时候每张卡就最多只算250个token，超过后直接丢弃（送到下一层）。

而在本文中，作者利用数据并行、pipeline并行和张量并行来训练GRIN MoE。

此外，对于没有专家并行性的MoE计算，作者发现Megablocks包非常有用，它的grouped_GEMM内核和包装器的性能更好。

应用这些新的工程化方法避免了专家并行，也就不用丢弃token了。

最终，与具有相同激活参数的密集模型相比，本文的方法实现了超过80%的训练效率提升。

上表中，作者将两种不同大小的MoE模型与具有相同激活参数量的密集模型进行了比较，使用相同的硬件测量了它们的训练吞吐量。

尽管MoE总的参数量是密集模型的六倍多，但在实验中达到了超过80%的相对吞吐量，证实了使用GRIN MoE方法的模型具有显著的计算扩展潜力。

（PS：密集模型的吞吐量是在与MoE模型相同的并行度设置下测量的，这里的比较是为了研究密集激活网络（非MoE）和稀疏激活网络（MoE）的GPU内核效率）

此外，在扩大模型大小时，密集模型和MoE模型显示出相似的减速模式，比如6.6B密集模型的训练吞吐量大约比1.6B密集模型的训练吞吐量慢4.19倍（后者的参数少4倍）。同样，42B MoE模型的训练吞吐量比10B MoE 模型的训练吞吐量慢约3.96倍（对应参数少4.2倍）。

并行实验

在只使用pipeline并行的情况下，通过在GPU之间进一步划分不同层，可以将最大专家数量从16个扩展到32个。但是，如果再增加专家数量，则会导致单个层的参数过多，一个GPU就放不下了。

所以下一个维度采用张量并行。

专家并行在前向和后向计算中有两个all-to-all通信开销，而张量并行在前向和后向计算中有两个all-reduce通信开销。

相比之下all-reduce操作的延迟更高一点，但可以通过精心排布前向和反向的计算来overlap掉一部分开销。

如上图所示，通过结合pipeline并行和张量并行，系统支持的最大专家数量扩展到52个（总共132B参数）。

这个数量是因为实验只用了64个GPU，最多能将模型划分为64个阶段，如果有更多的GPU，那么还能继续向上扩展。

不过作者也表示，使用更复杂的并行通常会导致计算吞吐量降低。

负载均衡

如前所述，本文没有采用专家并行，但是负载不均衡的事实依然存在。

作者在这里通过调整负载均衡损失来调节全局的负载均衡。常见的负载均衡损失定义为：

其中α是超参数，n是专家数量，fi是调度给专家的token比例。

传统方法在本地不同的GPU上计算fi，因此负载均衡损失将调节本地专家负载均衡并缓解token丢弃。

在本文中，作者通过计算全局的fi（比如数据并行过程中组内的all-reduce）来修改负载均衡损失，调节专家负载以达到全局平衡。

尽管这种调整会产生额外的通信开销，但类似于张量并行，这些通信也可以与计算overlap，从而在很大程度上减少额外的延迟。

最后，放一个测试结果来show一下GRIN MoE的数学推理能力：

作者注：我们对新发布的GAOKAO（即全国普通大学和学院入学统一考试）的数学问题进行案例研究，这是中国一年一度的全国本科入学考试。 该考试以其严格的安全协议而闻名，是评估AI模型回答数学问题的能力的理想测试平台。请注意，GRIN MoE的训练于太平洋标准时间6月3日结束，2024年GAOKAO于中国标准时间6月7日开始。

参考资料：

https://venturebeat.com/ai/microsofts-grin-moe-ai-model-takes-on-coding-and-math-beating-competitors-in-key-benchmarks/