MOE论文详解(4)-GLaM:Efficient Scaling of Language Models with Mixture-of-Experts
2022年google在GShard之后发表另一篇跟MoE相关的paper,
论文名为GLaM (Generalist Language Model),
最大的GLaM模型有1.2 trillion参数, 比GPT-3大7倍, 但成本只有GPT-3的1/3,
同时效果也超过GPT-3. 以下是两者的对比:
跟之前模型对比如下, 跟GShard和Switch-C相比, GLaM是第一个MoE在Decoder-only模型上的应用. (整体变化其实不大)
1. 训练数据
GLaM模型使用开源的1.6万亿(trillion)个token进行训练, 下图说明了训练数据的种类与混合比例:
2. 模型结构
整体模型结构跟GShard类似, 如下图所示, 总共有64个expert,
门控网络每次从中选择2个expert进行计算. 另外还有几点改动:
- 使用
relative positional bias(From Transformer-XL) 替换标准的position embedding - 在非MoE Transformer层,
使用
Gated Linear Unit(GLU)和Gaussian Error Linear Unit(GeLU)激活来替换第一个线性层和激活. - GLaM模型训练中权重和计算的切分使用2维的切分算法(From GSPMD).
- 把不同MoE layer上相同位置(具有相同的index值)的专家放到同一个device设备上, 使得不同的MoE layer得到相同的计算图
- 使用
while_loop控制语句wrap重复性的MoE layer - 要使得模型效果好, expert要足够大.
如果一个expert足够大时一个expert就对应多个core来分配,
也就是一个expert对应 \(\frac{N}{E}\)
个device. expert的权重shape为 \([E, M,
H]\), 按专家维度 \(E\)
和hidden维度 \(H\) 来切分;
input activation tensor的shape大小为 \([B, S, M]\), 按 batch维度 \(B\) 和 model维度 \(M\) 来进行切分. - 同时也依赖
GSPMD的编译pass过程来决定剩下tensor的切分策略
3. 结果
GPT-3与GLaM的比较如下, 分别在zero-shot, one-shot, few-shot等多个方面进行了比较, 效果比GPT-3要好
