梯度|北大校友“炼丹”分享：OpenAI如何训练千亿级模型？( 四 )

根据门控输出，并非每个专家都必须进行评估。当专家的数量太大时，可以考虑使用两层MoE。

文章插图
图10：专家混合（MoE）层的图示，门控网络只选择并激活了n个专家中的2个（来源：Shazeer等人，2017年）

G将输入与可训练权重矩阵Gg相乘，然后执行softmax：
由于这个过程会产生密集的门控制向量，不利于节省计算资源，而且时也不需要评估专家。所以，MoE层仅保留了顶部k值，并通过向G中添加高斯噪声改进负载平衡，这种机制被称为噪声top-k门。

文章插图

为了避免门控网络可能始终偏向少数强势专家的自我强化效应，Shazeer等人（2017）提出了通过额外重要损失的软约束，以鼓励所有专家拥有相同的权重。其数值相当于每个专家的分批平均值变异系数的平方：

其中，CV是变异系数，失重的waux是可调节的超参数。由于每个专家网络只能获得小部分训练样本（“收缩批次问题”），所以在MoE中应该尽可能使用大batch，但这又会受到GPU内存的限制。数据并行和模型并行的应用可以提高模型的吞吐量。

文章插图

图11：10亿单词的语言建模基准（左）模型从左到右包含4、32、256、256、1024和4096名专家（右）40亿参数的MoE模型在不同计算预算下的性能（来源：Shazeer等人，2017年）
GShard（Lepikhin等人，2020年）通过自动分片将MoE transformer 模型的参数扩展到了6000亿。MoE transformer 用MoE层取代其他每一个前馈网络层。需要说明的是，在多台机器上MoEtransformer 仅在MoE层分片，其他层只是复制。
GShard中的门控功能G有几种改进设计：

专家容量：通过一位专家的令牌数量不应超过“专家容量”的阈值。如果令牌被路由到已达到容量的专家，则令牌将被标记为“溢出”，并且门输出将更改为零向量。
本地组调度：令牌被均匀地划分为多个本地组，专家能力在组水平上得到加强。
辅助损失：与原始MoE aux损失相似，添加辅助损失可以最小化路由到每个专家的数据的均方。
随机路由：以与其权重成比例的概率选择第二位最佳专家；否则GShard遵循随机路由，增加随机性。

文章插图

图12：GShard中带辅助损失的组水平top-2门机制的伪代码（来源：Lepikhin等人，2020年）
Switch Transformer（Fedus et al.2021）用稀疏开关FFN层取代了密集前馈层（每个输入仅路由到一个专家网络），将模型规模扩展到数万亿个参数。负载平衡的辅助损失是，给定n个专家，fi是路由到第i个专家的令牌分数，pi是门控网络预测的专家i的路由概率。

文章插图
图13：Switch transformer，稀疏Switch FFN层位于蓝色框（来源：Fedus等人，2021年）
为提高训练稳定性，switch transformer采用以下设计：

选择精度：使用FP32精度以提高模型局部的稳定性，并降低FP32张量的通信成本。FP32精度仅在路由器功能主体内使用，结果将还原到FP16。较小的初始化：权重矩阵的初始化从平均μ=0且标准偏差的正态分布中采样，同时将Transformer初始化参数从s=1减小到s=0.1 。
使用更高的专家辍学率：微调通常适用于小数据集，增加每个专家的辍学率以避免过度拟合。他们发现，所有层中的辍学率增加会导致性能下降。在论文中，他们在非专家层中使用了0.1的辍学率，但在专家FF层中使用了0.4的辍学率。
- 上一页
- 1
- 2
- 3
- 4
- 5
- 6
- 下一页