DeepSpeed-MoE:训练更大及更复杂的混合专家网络

这是微软发布在2022 ICML的论文，MoE可以降低训练成本，但是快速的MoE模型推理仍然是一个未解决的问题。所以论文提出了一个端到端的MoE训练和推理解决方案DeepSpeed-MoE：它包括新颖的MoE架构设计和模型压缩技术，可将MoE模型大小减少3.7倍；通过高度优化的推理系统，减少了7.3倍的延迟和成本;与同等质量的密集模型相比，推理速度提高4.5倍，成本降低9倍。

MoE

MoE架构使用了类似GPT-3的NLG模型。包括350M/1.3B/6.7B(24/24/32层，1024/2048/4096隐藏尺寸，16/16/32注意头)。例如:“350M+MoE-128”是指以350M密集模型为基础模型，128位专家组成的MoE模型。

训练使用了128个A100 gpu。令牌数量为300B。

下图显示了MoE模型的验证损失明显优于其密集对应部分(例如，1.3B+MoE-128 vs 1.3B dense)。

并且MoE模型与基数大4 - 5倍的dense模型的验证损失相当(例如，1.3B+MoE-128 vs . 6.7B dense)。

在零样本评估方面，模型质量也不相上下：1.3 b +MoE-128模型所需的训练计算量与13 b密集模型大致相同，但模型质量要好得多。1.3 b +MoE-128模型可以用13B密集模型的训练成本达到6.7B密集模型的模型质量，而训练计算量减少5倍。

Proposed Pyramid-Residual-MoE (PR-MoE)

1、所有的层都学习相同的表示吗?

这个问题在计算机视觉(CV)中已经有了很好的研究:浅层(接近输入)学习一般表示，深层(接近输出)学习更客观的特定表示。

论文研究了两种不同的Half-MoE结构。

First-Half-MoE：MoE层在模型的前半层

Second-Half-MoE：MoE层在模型的后半层

结果可以看到，“Second-Half”的表现好。这也可以证明并非所有的MoE层都学习相同级别的表示。更深层的人从大量的专家中获益更多：或者说浅层的一般表示都是通用的。

2、是否有一种方法可以在获得泛化性能增益的同时保持训练/推理效率?

为了提高MoE模型的泛化性能，通常有两种方法:(1)增加专家数量，增加内存;(2)采用Top-2专家选择，这样计算量会多33%。

论文对比了两种moe

1、Top2-MoE使用Top2专家将参数提高一倍，2、Residual-MoE固定一个专家，然后并在不同的专家之间改变第二个专家，也就是说将来自MoE模块的专家视为误差校正项（或者叫一个主要专家，一个辅助专家）。

这两种模型的泛化性能相当。但是因为计算方式不同，Residual-MoE比Top2-MoE快10%以上。

3、PR-MoE

根据上面的研究结果，论文的新体系结构在最后几层使用了更多的专家。并且使用Residual-MoE架构，每个令牌分别传递一个固定的MLP模块和一个选定的专家。

PR-MoE使用更少的参数，但达到与标准moe模型相当的精度。

Proposed Mixture-of-Student (MoS)

对于损失函数，使用KD损失的一般公式来迫使学生模型模仿教师的输出:

损失函数测量预测与给定硬标签之间的交叉熵损失的加权和，以及预测与教师软标签之间的KL散度损失。

左图显示虽然KD损失最初提高了验证的准确性，但它在训练结束时开始损害准确性。由于PR-MoE与标准MoE相比已经降低了容量，这回导致学生模型进入欠拟合状态。

所以作者建议在训练过程中逐渐减弱KD的影响或尽早停止KD。

右图在400K步停止KD后，学生模型现在具有与教师相似的验证曲线。也就是说，在训练的后期可以不需要再让学生模型模仿教师模型，而是直接让学生模型自学。

如上上图（表格）所示，通过分段KD的MoS平均准确率为42.87和47.96，尽管层数减少了12.5%，但仍保持了350M(43.08)和13b教师模型(48.37)的99.5%和99.1%的性能。

DeepSpeed-MoE

DeepSpeed主要只在MT-NLG 530B等Transformer上进行研究，之前没有在MoE上进行研究。而MoE推理性能取决于两个主要因素:总体模型大小和总体可实现的内存带宽。

1、数据并行

张量切片(用于非专家参数)，专家切片(用于专家参数)是将单个参数拆分到多个gpu上，可以充分利用gpu上的内存带宽。

为了将非专家计算扩展到相同数量的gpu，可以在没有通信开销的情况下使用数据并行。

2、将能分层的都进行分层

与all-reduce、broadcast等通信集合一起使用，可以将通信跳数从O(p)减少到O(G + p/G)，其中G为节点中GPU的数量，p为GPU设备的总数，也就是说，GPU越多，减少的效果越明显

上图显示了该实现的设计概述。尽管通信量增加了2倍，但这种分层可以实现对小批量进行更好的扩展。

3、 并行通信和内核优化

这两部分比较深奥，并且更加专业，所以我这里只简单说明

由于张量并行度中的每个GPU都包含相同的数据，因此可以将所有对所有的通信限制在具有相同张量并行度等级的GPU之间。

将门控函数融合到单个内核中，并使用密集的令牌到专家映射表来表示令牌到专家的分配，大大降低了内核的启动开销。