6.1 用 Router z-loss 稳定模型训练 在论文ST-MOE: Designing Stable and Transferable Sparse Expert Models中,作者提出了一种新的辅助损失函数,称为Router z-loss,用于提高稀疏模型的训练稳定性,同时保持或稍微提高模型质量。这个损失函数是针对稀疏专家模型中的路由器(router)部分设计的,路由器负责将输入的 token ...
相反地,z-loss自然地鼓励模型产生较小的对数值,因此可以更精确地建模。 加入到模型训练的总损失中,如下式所示。 ST-MoE经过实验,选择了。 是auxiliary load balance loss负载均衡损失,ST-MoE这里使用了和GShard/Switch Transformer所用的相同的损失计算,这里回顾一下: 是专家数, 是包含 个token的batch。 表示被分...
GShard 引领了 top-2 gating 和容量约束;Switch Transformer 用 top-1 gating 打开了更大规模的可能性;GLaM 让我们见识了训练能耗可以大幅下降;DeepSpeed-MoE 则深入挖掘了训练和推理层面的负载平衡;ST-MoE 用 z-loss 改善了稳定性;Mixtral 看到了专家分配的时间局部性;等等……最后到一些更灵活的思路,如 DeepSe...
每个 MoE 层中的专家更少,但规模更大。以前,这种高FLOP MoE 在训练过程中被发现不稳定,但增加路由器 z-loss 使 ST-MoE-32B 模型能够成功训练。由于与之前的 MoE 模型相比,ST-MoE-32B 的参数更少,因此更易于微调和部署。 “我们在 FLOP 和参数之间寻求平衡。高 FLOP 稀疏模型以前在我们的设置中不稳定(即编...
用Router z-loss 稳定模型训练 之前讨论的平衡损失可能会导致稳定性问题。我们可以使用许多方法来稳定稀疏模型的训练,但这可能会牺牲模型质量。例如,引入 dropout 可以提高稳定性,但会导致模型质量下降。另一方面,增加更多的乘法分量可以提高质量,但会降低模型稳定性。
在负载均衡算法方面,支持Sinkhorn(S-BASE)、Z-Loss,以及Load balancing Loss。 Grouped GEMM Megatron-Core MoE开发了GroupedGEMM来解决多Experts变长输入这一问题。当每个Rank有多个专家时,Megatron-Core MoE利用自CUTLASS 2.8引入的Grouped GEMM特性,将多个局部(可能是较小的)GEMM操作合并为单个GroupedGEMM kernel,...
稳定训练基于 Z-loss 损失函数的路由网络 上文讨论的平衡损失可能会导致不稳定问题。但是可以使用许多方法以牺牲质量为代价来稳定稀疏模型。例如,引入 dropout 可以提高稳定性,但有损于模型的性能质量。另一方面,添加更多的乘法分量可以提高模型的性能质量,但会降低稳定性。
可以通过混合精度训练、更小的参数初始化,以及 Router z-loss 提升训练的稳定性。 第三个问题:MoE 如何解决 Fine-Tuning 过程中的过拟合问题? 可以通过更大的 dropout (主要针对 expert)、更大的学习率、更小的 batch size。目前看到的主要是预训练的优化,针对 Fine-Tuning 的优化主要是一些常规的手段。
Switch Transformer(2021)改进了Load Balance的设计,同时提出了Capacity Factor的概念 ST-MoE提出了router-z loss的概念,总loss是load balance loss和router-z loss的加权
用Router z-loss 稳定模型训练 之前讨论的平衡损失可能会导致稳定性问题。我们可以使用许多方法来稳定稀疏模型的训练,但这可能会牺牲模型质量。例如,引入 dropout 可以提高稳定性,但会导致模型质量下降。另一方面,增加更多的乘法分量可以提高质量,但会降低模型稳定性。