MoE训练论文解读之Megablocks：打破动态路由限制

NewBeeNLP

发布于 2023-11-08 15:18:20

1.1K0

发布于 2023-11-08 15:18:20

文章被收录于专栏：NewBeeNLP

作者 | 方佳瑞（已授权）整理 | NewBeeNLP https://zhuanlan.zhihu.com/p/653270049

GPT-4用了Mixture-of-Experts（MoE）架构，引起了广泛关注。然而，MoE训练并不是一项简单的任务，它面临着一些主要的挑战和难点：

1、动态路由限制：当前的框架对MoE层中的动态路由进行了限制，以满足现有软件和硬件的约束条件。用户必须在计算中选择drop tokens或zero-padding，前者影响模型效果，后者浪费计算资源。这种限制导致模型质量和硬件效率之间存在某种权衡，导致调参困难。

2、如果打破上述限制导致每个专家的负载动态变换，计算kernel和多卡通信实现都更加复杂。

3、Megatron-LM没有MoE官方实现，缺少一个能和Dense模型训练对比的Best Practice软件实现。

MLSys 23有两篇关于MoE训练的文章，它们都瞄准了传统MoE训练中上述问题。不过二者思路并不相同。尽管作者没有NVIDIA的人，这篇来自Standford，MSR和Google的作者们的产学研合作结果还是以Mega作为名字。

也许是因为Megablocks带着Megatron-LM品牌光环，自然容易引起更多关注，这点有点像茅台冰激凌，有了茅台两个字大家不太会关注冰淇淋了。

Tutel: Adaptive Mixture-of-Experts at Scale[1]
MegaBlocks: Efficient Sparse Training with Mixture-of-Experts[2]

本文带大家深度解读一下megablocks。先说结论，我觉得megablocks技术路线不适合作为MoE训练的方案，因为它的目标场景比较狭窄，拿锤子找钉子动机太明显。反而，Tutel更适合作为MoE大规模训练的技术路线参考。（2023.8.31增加Tutel文章解读见如下链接）

方佳瑞：MoE训练论文解读之Tutel: 动态切换并行策略实现动态路由[3]

另外，文章有很多弦外音，需要读者深挖才能意会。

1. MoE基础知识

要理解MoE训练面临哪些问题，我们需要一些背景铺垫。

1.1 MoE网络结构

MoE只对Transformers的FFN层进行改变。这里复习一下FFN的计算，其实就是两个MLP串联在一起。

FFN(x, W_1, W_2, b_1, b_2) = max(0, xW_1 + b_1)W_2+b_2

用MoE Layer替代FFN layer，from Switch Transformers

上图展示了在Transformer如何将FFN层取代为MoE，因为是从Switch Transformers论文扣的，MoE层这里叫Switch FFN Layer（浅蓝色）。该层独立地作用于MHA（Multihead-Attention）层输出Activation（batch, seqlen, hidden_dim）中的每个token (1, hiddeen_dim)。

图中显示了两个token（x1="More", x2="Parameters"）通过4个FFN专家进行路由(实线)。其中路由器独立地为每个token选择路由。Switch FFN Layer返回由选择的FFN计算并乘以路由器预知（虚线）后的输出。

图中只选择top-1路由到一个FFN，其实可以也路由到多个。关于top-k的k设置多少有过争议。Shazeer2017年时认为为了有non-trivial gradients传播到路由函数，每个token路由到k>1个专家是必要的。Switch Transfomer推翻了这个结论，认为仅路由到单个专家也可以，并证明这种简化策略可以保留模型质量，减少路由计算开销，且性能更优。

因为每个token被路由到哪个FNN是不确定的，因此叫动态路由。这里每个专家分配到的token不一样，计算量也就不平衡。为了强行把所有专家计算量拉平，有了专家容量（Expert Capacity）概念，即每个Expert计算的token数。它的计算方法是，将一个Batch的token总数除以专家数，然后再根据一个容量因子进行扩充。公式如下：

expert\ capacity=(\frac{tokens\ per\ batch}{number\ of \ experts}\times capacity\ factor)

之所所以有专家容量，当初是因为Switch Transformer用TensorFlow Mesh实现，因为用XLA实现，tensor shape必须静态指定，因此必须拉平所有专家计算量。为鼓励专家负载平衡，Switch Transformers添加 Load Balancing Loss 。对于每个Switch层，此辅助损失在训练期间添加到总模型损失中。

我们看一个top-1，expert capacity=1，专家数量为3的例子，超过容量的token直接丢弃。

A Mixture-of-Experts Layer. from Megablock.

(1) 生成每个token路由对应的专家编号和置信概率。

(2) 前一层Multihead Attention输出的特征向量，维度(batch, seq_len, hidden_dim)，按专家分配结果进行重排，将同一专家分配的token聚集在一起。如果分配给某个专家的token超过其容量，额外token将被丢弃。

(3) 专家的使用自己的tokens数据进行计算，计算方式和FFN一样，两个连续矩阵乘法。

(4) 专家计算结果经过反重排，再被token的概率scale一下作为输出。

1.2 传统MoE实现的问题

动态路由限制让每个专家计算量一致引入了模型质量和硬件效率之间的权衡，因为用户必须决定是丢弃token损失精度还是zero-padding浪费计算和内存资源。

这个决策通常是通过超参数调整来进行的，这增加了使用MoEs的复杂性。Megablocks就是想打破这种限制，让没有token drop且没有zero-padding开销，从而摈弃了loss balance和expert capacity限制。

1.3 矩阵视角看待MoE计算

我们从矩阵乘法操作角度来理解一下MoE的流程。

B = batch size; S = sequence length; H = hidden_dim; E = expert number; MLP的FFN的权重一般是(H, 4H)。

原始FFN，两个MLP的矩阵操作如下

A_2\left(B, S_0, H\right)=A_1(B, S, H) * W 1(H, 4 H) * W 2(4 H, H)

MoE，以两个Expert为例，矩阵操作如下

A_1(B, S, H) \text {-Permute->} A_1^0\left(B, S_0, H\right) \mid A_1^1\left(B, S_1, H\right)

专家1：

A_1^0\left(B, S_0, H\right) * W 1(H, 4 H) * W 2(4 H, H) \rightarrow A_2^0\left(B, S_0, H\right)

专家2：

A_1^1\left(B, S_1, H\right) * W 1(H, 4 H) * W 2(4 H, H) \rightarrow A_2^1\left(B, S_1, H\right)

A_2^0 \mid A_2^1

-Unpermute

\rightarrow(B, S, H)

如果在一个GPU上计算，如果

S_0

和

S_1

相等，其实就是一个Batched GEMM操作。每个矩阵维度信息如下

A_2=A_1(E, B * s, H) * W_1(E, H, 4 H) * W_2(4 H, H)

如果不相等，则是 Variable Sized Batched GEMM （名字我自己起的，不同论文叫法不一样，Megablock叫Block Sparse Matrix Multiplication，Byteformer叫Variable Sized Group GEMM），后面会讲到这是一个经典问题。

\left(E, B * s_i, H\right) *(E, H, 4 H) *(4 H, H) i=0 \ldots E

，

s_i

各不相同。

1.4 MoE和Dense结构的计算量和参数量

分析一下MoE计算量和参数量特性。

MoE MLP参数是原来E*倍，因为有E个专家。r是MoE层出现的频率，比如MoE一般是每隔一个FFN有一个，r=0.5。

MoE参数量

Dense参数量= 1+(8*h^2 + 5h)\times (E-1) / (12h^2+13h) * r

比如，对于一个h=768的125M参数MoE，r=0.5 参数量是同等配置Dense模型的5.99x。

使用top-k路由每个token在MoE层计算量是原来MLP的k倍。参数比例等价于计算的比例，训练相同tokens，计算量关系如下。

MoE计算量

Dense计量算= 1 + (8*h^2 + 5h)\times (E-1) / (12h^2+13h) * r * k

1.5 MoE的并行策略

Dense模型使用Tensor Parallelism（TP），Data Parallelism（DP）还有Pipeline Parallelism（PP)。对MoE增加了Expert Parallelism（EP）。

Expert parallelism把专家分配到不同的计算资源上。比如，一个专家分配1-N个GPU。如有下特点：

如果每个专家的tokens数目不同，则EP中不同GPU的计算负载不同。
在MoE层和MHA前后衔接地方需要额外通信操作：MHA如果使用DP，转化成MoE的EP需要all-to-all操作，

\mathrm{E}^*(\mathrm{~B}, \mathrm{~S}) \rightarrow(\mathrm{EB}, \mathrm{S}_i), \ldots,\left(E \mathrm{~B}, \mathrm{~S} \_\mathrm{e}\right)

，MoE算完也需要all-to-all操作。EP可以和TP和DP混合使用。如果不drop tokens，最优的混合策略需要根据输入activations的token分布决定，每个batch迭代的最优并行方案可能都不一样，这也是Tutel讨论的内容。

2. MegaBlocks

MegaBlocks解决 一个GPU上有多个专家 ，如何实现高效没有限制路由策略的MoE。

如果GPU上分配了E个专家，一个最简单的办法就是执行E次串行GEMM就可以了。但是这样GEMM之间没有并行，可能GPU利用率就达不到机制。所以我们硬要把E个维度不同的矩阵乘法一起算，也就是Grouped GEMM。这个操作在CUTLASS中有现成的实现。

(bsz, m, k) * (bsz, k, n) → (bsz, m, k)

Grouped GEMM要求所有矩阵乘法的维度都一致，而在MoE中，因为每个专家的token数不同。如果想硬凑成Grouped GEMM就需要zero-padding或者截断（drop tokens）。如果有一个m维度可以变化的Variable Sized Grouped GEMM方案可供调用，问题就可以完美解决了。

以下图为展示有三个Expert每个专家都要进行第一个MLP的GEMM。(A)控制expert capacity，让每个专家tokens数目相同，然后用三个GEMM操作实现。(B)将图(A)的三个GEMM拼在一起算，输出可以看成一个对角稀疏的大矩阵，Batched GEMM。(C)没有expert capacity约束，每个专家处理全部路由来的tokens。三个矩阵第0维表示token数的s0, s1, s2各不相同，因此输出对角矩阵的row number也不同。

我们从Batched GEMM角度描述整个MoE MLP前向后项计算特性。先做符号化定义：矩阵乘法需要的三个矩阵中的一个（两个输入和一个输出）是稀疏的，而其他的是稠密的。每个操作都用一个由三个字符组成的字符串来描述，其中每个字符可以是“S”表示稀疏或“D”表示密集。字符的顺序是输出，然后是左输入，最后是右输入。例如，两个稠密矩阵的稀疏输出乘积是“SDD”。上标“T”表示输入参数的转置。例如， SDD^T 表示SDD，其中右侧输入矩阵已经转置。

MoE中每个专家是一个两层的多层感知机（MLP）。前向传播需要进行SDD操作，然后是DSD操作。对于反向传播，我们分别计算第二层数据梯度的

SDD^T

和权重梯度的

DS^TD

，然后是第一层数据梯度的DSDT和权重梯度的

DD^TS

。

相似的问题在Transformers Encoder推理中也存在。MHA结构中也有batched GEMM，如果输入sequence长度变化，为了能调用Batch GEMM接口必须加Padding，都pad成最长的sequence。字节AML team实现一套接口，和本文类似也是基于CUTLASS库提供基本矩阵乘法运算操作，和本文工作不能说是殊途同归，只能说是不谋而合了。

ByteTransformer: A high-performance transformer boosted for variable-length inputs[4]

锤MHA的锤子我们这里可以再锤一下MoE。在本文里的方法叫”一个也不能少块稀疏法“（NO-TOKEN-LEFT-BEHIND WITH BLOCK SPARSITY）。

Variable Sized Batched GEMM优化性能的核心就是将多个GEMM操作分解成很多小矩阵块的GEMM操作，矩阵块矩阵乘法的尺寸是确定的，因为块比较小，它们可以同时计算来完成大矩阵的Variable Sized Batched GEMM。

这里小矩阵块乘法利用CULTLASS中的128*128矩阵乘。GPU同时发起很多这样的小块矩阵乘，以达到更高的计算效率。小矩阵乘的输出位置是不连续的，因此需要一个数据结构高效访问Blocked-Sparse Matrix。这里利用了blocked compressed sparse row（BCSR）数据结构，并对行列访问速度和转置操作做了优化。

这里Variable Sized Grouped GEMM实现具体细节我不赘述了，Megablocks项目里大家也看不到。因为这个实现封装在另一个叫 Sparse Toolkit （stk） 的项目中。stk是一个轻量级实现block-sparse 矩阵乘法的库。

https://github.com/stanford-futuredata/stk[5]

Megablocks有一个小节提了并行。对于并行策略它其实没做新的东西，反而只做了阉割。简单来书，MoE只支持EP，Transformer其他部分支持DP+TP。实验也都是设置64个专家在8 GPU上训练，为了能够凑出一个GPU多个专家的情况也够不容易的。因此，它无法避免不同GPU负载不均衡情况，而反而是Tutel要解决的问题之一。

3. 总结

3.1 文章优点

1、有开源代码，且基于Megatron-LM改的。鉴于Megatron-LM在大模型训练场景的统治地位，这点很容易引起好感。作者也把mega作为软件名字前缀，强调这一点。

2、stk项目工作量比较大，是作者花心血优化很久的项目。而且可以供其他Variable Sized Batched GEMM使用。

3.2 文章缺点

1、应用场景有限，甚至可能是伪需求。实际MoE训练中，GPU数目比Expert数目要大的，都是使用Expert Parallelism让多个GPU伺候一个专家，几乎没有多个专家在一个GPU上。论文里拿锤子找钉子痕迹过于明显。感觉就是单独用stk写论文缺乏创新点，所以套了一个MoE的壳。

2、没法大规模扩展。因为megablocks只解决在一个GPU里计算专家复杂不均衡问题，没解决GPU见不同专家复杂不均问题。

3、软件里megablocks实现的MoE只支持EP+TP，因为它直接把DP的process group挪给EP了。这其实和Tutel的观点冲突的，因为Tutel因为MoE最优并行方式应该cover EP+TP+DP。并行方式的缺陷，让Megablocks扩展到大规模会有局限，论文的实验也是在8 GPU上做的。尽管Megablock论文和Tutel进行了实验对比，但是硬件规模都太小不足以体现真实差异。

4、文章书写诚意不足。这里举几个我不满意的点：

Figue 3 (C)误导性太明显，文章只写了如何让Batched GEMM的一个维度是的变化的，而图中让两个维度是变化的了，后面正文澄清目前并不支持两个维度是变化的。这个图应该是文章最重要的了，读者会反复看的，在这里这个玩笑不太严肃。

整个Sec 5就是把如何实现Variable Sized Batched GEMM通用方法搬移过来了，对MoE问题没啥特殊处理。实现部分描述和配图可能直接搬运Variable Sized Batched GEMM技术报告，连MoE问题特点都不简单适配一下。Figure 6的矩阵都不是对角稀疏的，这种稀疏性根本不会再MoE出现。我其实更期待利用MoE的稀疏性可以比通用方法有什么提升点，在这篇文章里显然落空了。

总体来说，文章作为SysML文章还是合格的，也算是HPCer以子之矛攻MLSys之盾的一次成功尝试。如果想用它在生产环境真正去训练MoE就不太合适了。

最后，呼吁CUTLASS出个官方的Variable Sized Batched GEMM实现吧，别让孩子们自己折腾了！

本文参考资料

[1]

Tutel: Adaptive Mixture-of-Experts at Scale: https://proceedings.mlsys.org/paper_files/paper/2023/hash/9412531719be7ccf755c4ff98d0969dc-Abstract-mlsys2023.html

[2]

MegaBlocks: Efficient Sparse Training with Mixture-of-Experts: https://proceedings.mlsys.org/paper_files/paper/2023/hash/f9f4f0db4894f77240a95bde9df818e0-Abstract-mlsys2023.html

[3]

方佳瑞：MoE训练论文解读之Tutel: 动态切换并行策略实现动态路由: https://zhuanlan.zhihu.com/p/653518289

[4]

ByteTransformer: A high-performance transformer boosted for variable-length inputs: https://ieeexplore.ieee.org/abstract/document/10177488/

[5]

https://github.com/stanford-futuredata/stk: https://github.com/stanford-futuredata/stk

[6]

分析transformer模型的参数量、计算量、中间激活、KV cache，回旋托马斯x: https://zhuanlan.zhihu.com/p/624740065

[7]

MegaBlocks: Efficient Sparse Training with Mixture-of-Experts: https://proceedings.mlsys.org/paper_files/paper/2023/hash/f9f4f0db4894f77240a95bde9df818e0-Abstract-mlsys2023.html

[8]

Tutel: Adaptive Mixture-of-Experts at Scale: https://proceedings.mlsys.org/paper_files/paper/2023/hash/9412531719be7ccf755c4ff98d0969dc-Abstract-mlsys2023.html

[9]

Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity: https://www.jmlr.org/papers/volume23/21-0998/21-0998.pdf

本文参与腾讯云自媒体同步曝光计划，分享自微信公众号。

原始发表：2023-11-07，如有侵权请联系 cloudcommunity@tencent.com 删除

路由