从dense到MoE -- sparse upcycling

发表于 2024-07-19 分类于 CS ， NLP ， LLM 本文字数： 4k 阅读时长 ≈ 7 分钟

【本文已在同名微信公众号 / 知乎 / 个人博客linsight.cn 上线】

目前已经有很多优秀的dense大模型，那么要通过MoE获得更强的模型，用已有的dense模型进行初始化是一个自然的想法。Google的sparse upcycling对此做了一些实验，由于实验是在2022年做的，模型用的是T5系列语言模型和Vision Transformer系列视觉模型。

文中给出两个适合使用sparse upcycling的场景：
- 已有dense模型，想在有限的计算资源下提升模型效果。
- 要训一个模型，不知道dense模型和MoE哪个会有更好的效果（虽然通常来说MoE更好，但是训练难度和结果不确定也更大），那么就可以先训练一个dense模型保底，然后再在dense模型的基础上扩展成MoE结构继续优化。

下面具体看下一些实验细节。

设置

对于transformer模型，sparse upcycling的操作如下图

除了原模型的MLP层替换成MoE层外，其他组件包括layernorm、attention都直接从原dense模型copy到MoE模型。

实验上，一些具体的基础设置如下：
- 在原模型基础上，每2层替换一个MoE层，从第二层开始替换
- MoE模型的总层数的dense模型层数相同
- 每个MoE层专家数为32个；虽然使用更多的专家不会明显增大训练的FLOPS，但是更多的专家会带来larger initial quality drop relative to baseline dense model，而需要更多的计算资源来恢复这个quality drop；后续会有实验探索expert数量的影响
- 每个expert都用原模型的MLP层参数初始化
- router使用standard deviation=0.02的zero-mean normal distribution随机初始化
- 在encoder使用expert choice routing，基础的设置是capacity factor C = 2，后面也做了关于capacity factor的消融实验
- 在decoder使用token choice routing（top-k routing），k=2，并加上auxiliary loss帮助负载均衡，权重为0.01；在decoder使用top-k routing的原因是"to avoid train time (full batch teacher forcing) versus inference time (single token auto-regressive decoding) discrepancies"（和expert choice routing的设计相关）

MoE模型训练时使用和dense模型一致的batch size、learning rate schedule和weight decay等。

其中learning rate schedule用的是inverse square root learning rate schedule，因此MoE的训练可以接着dense模型的schedule接着进行。

实验中所用的一些模型参数如下表

实验

CORE RESULTS

1、dense模型继续训练 vs upcycling

随着训练量的增加，upcycling相比dense模型继续预训练的优势逐渐扩大，如下图所示

2、下游任务模型效果

上面对比的是预训练模型。这些预训练模型经过微调后的效果对比如下。

相比预训练模型，微调模型表现出相对更大的震荡，不过大致趋势还是可以看出MoE模型更有优势。

3、MoE from scratch vs upcycling

从零开始训练的MoE和upcycling方法的对比如下

从零开始预训练的MoE模型效果提升得更快，这可能得益于多样化的专家初始化和更大的lr。
只要给的计算资源足够多，从零开始训练的模型最终会赶上甚至超过upcycling的模型。
在有限的训练资源下，upcycling的训练效率更高，从零开始训练的模型大约需要相当于原dense模型1.2倍的训练资源才能达到upcycling模型的效果。如果现在的训练资源<=训练dense模型的资源，那么sparse upcycling是更划算的。

4、sparse upcycling vs dense upcycling

对比《Scaling language models: Methods, analysis & insights from training gopher》中的depth tiling（dense upcycling）和 sparse upcycling的预训练效果，结果当然是sparse upcycling效率更高点，如下图所示

（不过这里没有提及depth tiling之后的模型规模）

消融实验

1、Amount of dense pretraining

upcycling的效果可能受用于初始化的dense模型的收敛情况影响，因此取了不同step的dense模型checkpoint作为upcycling的初始化，并且都继续训练了200k个step，结果如下图

结论是基本上无论从哪个checkpoint初始化MoE模型，收益都比较稳定。

2、Router type

使用不同的router（expert choice和token choice）对比结果如下

结论是，在相同的step下，expert choice和token choice的效果基本一样，但是如果从时间上来看，使用expert choice routing的模型训练更快。

3、Expert capacity factor

每个专家处理的token越多，计算量就越大，理论上效果也越好。

使用不同的capacity factor，模型效果对比如下

结论是，虽然理论上增加专家容量可以提升效果，但时间上，通常C = 2的效率比较好，即一定的时间内提升的效果最多（注意计算资源是有限的）。

4、Number of MoE layers

在视觉模型上对MoE层数的效果进行了式样。

如下图右边两个小图，是使用不同的MoE层的效果，比如1表示只把最后一层MLP层替换为MoE层，以此类推

结论是，更多的MoE层并不总是更好，大概是把5_{6层替换成MoE层的时候效果最好（40%}50%的层数）。

5、Initialization of experts

对比了使用dense模型的MLP层初始化专家，和随机初始化专家，结果如下

结果上看，使用dense模型的参数初始化专家效果更好。

6、Number of experts

如前面提到的，增加专家数并不会增大计算量，下图实验了2~128个专家下的效果

结果上来看，效果是随着专家的增加而提升的，虽然最后表现出了收益递减的情况。

其他

1、optimizer

在vision模型上，还尝试了使用dense模型的optimizer状态来训练MoE模型，但是并没有带来任何收益。

2、router normalization

另外，为了减少从dense到MoE初始化的performace drop，尝试了对router的输出进行normalization，以保持每个token得到的weight总和是1。

这个做法直觉上应该是有益的，不过会有一个小问题，那就是对于只被一个expert选中的token，会有vanishing routing gradients。

实践上，router normalization在视觉模型上基本和不进行normalization的效果差不多，但是在语言模型上，会使得MoE模型效果变差。这二者的表现差异可能是因为语言模型上部分router使用了token choice routing。

实际上目前大部分最新的MoE模型都没有开router normalization，但这里的原因感觉还有待深入验证。

小结

如果资源无限，那么从零初始化MoE模型更好，这点和天工MoE的实验是相同的；
从结果上看，MoE层和专家数不是越多越好，但是多少是好，感觉至今其实没有很明确；
目前大部分的MoE都使用token choice routing了，这里可能需要重新实验一下效果。

读到这了，来一发点赞收藏关注吧~

博客：http://www.linsight.cn/
知乎：Linsight
微信公众号：Linsight

【往期文章】
MoE模型的前世今生
 DeepSeek-V2和MLA
昆仑万维-SkyworkMoE
成本10w刀的JetMoE
MoE的top-p routing
对MoE模型的一些观察
 从loss视角理解大模型涌现能力
 LLM长上下文的问题
 解锁大模型长上下文能力
 大模型推理窗口-从有限到无限大
 理解Attention:从起源到MHA,MQA和GQA
大模型推理加速-投机解码
 大模型推理加速-MEDUSA
LLM的重复生成和ICL
大模型偏好对齐-DPO
大模型偏好对齐-ODPO
大模型偏好对齐-simPO
大模型偏好对齐-IPO
Qwen2技术报告
 Yi技术报告-划重点看细节
 MiniCPM
GLM4报告的一些技术点
 Gemma2
苹果的OpenELM
从Yuan2.0到Yuan2.0-M32
bilibili的index-1.9B
transformer中normalization的二三事
 从代码实现看normalization-到底做了什么
 稀疏注意力计算:sliding window attention
理解LLM位置编码:RoPE
RoPE的远距离衰减
 大模型算法题(1)
大模型算法题(2)
大模型算法题(3)
大模型算法题(4)
大模型算法题(5)
大模型算法题(6)
大模型算法题(7)
大模型算法题(8)

Reference

【1】SPARSE UPCYCLING: TRAINING MIXTURE-OF-EXPERTS FROM DENSE CHECKPOINTS https://arxiv.org/abs/2212.05055