MOE论文详解(4)-GLaM

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要好

4. 参考

• GLaM
• MOE论文详解(4)-GLaM