MoE训练的难题MoE的训练关键是路由层但关键问题是路由层是不可微分的无法反向传播有三种解决方法RL学习路由加上随机扰动启发式的负载均衡lossRL路由RL理论上有效但实验中和另外两种baseline比没有明显的优势随机扰动另一个做法是给路由时加上随机扰动这样可以让负载更均衡一点不容易出现一个专家永远被选或者永远不被选。实验表明训练时加入随机扰动确实可以减少训练崩溃图中下面的表第二列是正常训练的比例加入扰动后从4/6提升到了3/3.但是模型能力变弱了见表中第三列模型最终收敛时的表现不如不加扰动的baseline因此这个方法最终还是被废弃了。启发式负载均衡loss第三种方法也是最主流的。加一个启发式的loss来实现负载均衡。这里loss的公式是(4)可以看到对每个专家的f,p两个值乘积求和。其中f是这批token中被路由到这个专家的比例p是这批token每个token路由到这个专家的原始概率的均值。这两个值都刻画了这个专家被选中的概率。这个loss是可导的并且可以反应负载均衡情况可以让模型学习到路由规则。专家/设备负载均衡原始的公式是对于每个专家计算负载均衡的loss更进一步还可以对并行训练时每个设备做负载均衡f,p改成每个设备的路由token个数以及设备上的所有专家的路由概率均值。毕竟除了专家均衡提高训练稳定性让设备均衡提高设备利用率也是重要的这可以提高训练效率。无辅助损失均衡更进一步在Deepseek V3中提出了新的负载均衡方法不改变loss项。而是给每个专家加一个偏置项b计算每个token和专家的得分s后临时加上专家的偏置b来做topk排序如果在前k大则路由到这个专家但路由后的输出里不包含偏置项b。每一轮统计每个专家被路由到的次数如果大了就降低b小了就增大b。这样可以在不碰loss的情况下实现负载均衡。这样做的意义是修改loss来实现负载均衡可能强迫模型把不适合给某个专家的token为了负载均衡的考虑塞给专家导致输出质量低。消融实验去掉负载均衡后模型的loss下降更慢如上图粉色是带负载均衡loss的蓝色不带。去掉负载均衡确实也会导致专家负载不均如下图八根线表示八个专家被路由到的次数左边是没有负载均衡有两个专家主要被路由到剩下的专家几乎没用。右侧是带负载均衡loss八种专家的路由次数都比较接近。并行对于一个设备上的多个专家朴素的做法是每个专家和路由到他的token组成的矩阵分别做矩阵乘法如左图。但既然在一个设备上一个自然的优化是把专家和token都拼起来得到两个大矩阵做一次大矩阵乘法这样对GPU更友好GPU喜欢大量数据流大量小数据导致的同步开销会造成吞吐浪费。但这样的问题是只有对角线上的块是有效的结果其他部分都浪费了如中图。多次计算融合成一次GEMM是好的为了解决了一次GEMM的冗余计算使用结构化稀疏这是英伟达硬件就支持的功能可以传入一些元数据表示最终的结果矩阵中有哪些位置是需要计算的然后利用硬件的稀疏算力只算这些稀疏位置的结果一般来讲要求每四个位置有且只有两个位置需要计算然后可以把吞吐提升一倍。现在有很多现成的库可以帮我们把数据重新排布满足每四个位置有且只有两个有效然后调用显卡的结构化稀疏GEMM接口。下采样另外还有一些优化思路比如英伟达的Nemotron3在进入MoE计算前对激活值做一个下采样压缩到更小的维度这样可以减少MoE计算量降低MoE的专家大小从而还可以支持更多专家。最后MoE层计算完毕后再上采样回原来的大小。MoE的随机性MoE的推理结果具有更大随机性对比Dense模型。这是因为在多用户推理系统中多个用户的请求会被打包成一个batch发送给MoE层先做路由。由于每个专家一般会设置接受的token上限如果一个专家能处理的token满了剩余原本被路由到这个专家的token会被重新路由或者丢弃等待下一轮路由。这导致其他人的推理请求可能会影响你的请求的路由情况如果你的请求在一个batch的靠前的位置更可能路由到想要的专家如果在一个batch靠后的位置更可能被丢弃路由不到想要的专家。而你的请求在一个batch的靠前还是靠后这是纯随机的可能和系统的随机策略甚至网速有关。训练稳定性路由结果非常影响MoE模型的训练稳定性而bf16位数有限舍入误差可能导致实际得分不一样的两个专家在bf16下得到相同的路由分数从而进行错误的调度导致loss下降缓慢甚至上升。解决方法是在计算路由分数的地方使用fp32高精度另一个解法是下方的公式(5)增加一个loss项被称为z-loss把所有结果求以e为底数的幂求平均。这能有效惩罚数值过大的结果。数值过大也是浮点溢出的主要原因。消融实验验证了zloss确实有效粉色是有zloss蓝色是无zloss可以看到蓝色在训练过程中明显多很多尖峰这都是训练中不稳定的时候。过拟合MoE更容易过拟合。蓝绿分别是MoE稠密模型的训练集表现MoE收敛得更快。但MoE在验证集上表现更差明显出现了过拟合橙红分别是MoE和稠密模型的验证集表现。一种解决方法是微调是冻结MoE层微调其他层。例如上图分别微调不同层可以发现微调其他层benchmark得分都很高只有微调MoE时得分明显低。另一种简单但有效的方法是增加微调数据过拟合的一个重要原因就是数据集太小Deepseek的解法是SFT阶段增加大量数据。Upcycling翻译成中文是回收再利用。意思是我们可以用一个已经训练好的稠密模型作为MoE的基础模型。见上图具体架构是注意力层和归一化层都直接使用稠密模型的参数。MoE层的多个专家都使用稠密模型的相同线性层。也就是开始初始化了多个相同的专家后续在训练中再让这些专家学习分工。可以看到在相同的训练轮数后橙色的Upcycling明显比蓝色的稠密模型表现好。甚至large级别的Upcycling MoE模型表现反超了更大规模的XL级别稠密模型。举例MiniCPM采用这个路线的经典例子是面壁智能的MiniCPM这个模型主打小参数高性价比定位是端侧手机车载部署。采用Upcycling MoE的13.6B版本性能远超激活参数大小接近的其他稠密模型。回顾Deepseek演进MoE这节大部分计数都是在讨论Deepseek这里全面回顾一下DS的演进路线V1就采用了MoE使用了topk路由负载均衡loss共享专家和细粒度专家。V2参数量增大专家划分更细160个专家激活10个。此外还有很多系统级别的优化比如topm的设备路由对每个设备的得分就是这个设备上的所有专家得分之和据此选出得分前m的设备。对这些设备上的专家再做topk专家路由。这样的好处是最后选出来的topk的专家一定是在不超过m个设备上的MoE的all-to-all通信量是和设备数成二次方关系的限制设备数可以优化通信提高吞吐。设备负载均衡不以专家为单位而是以设备为单位计算负载均衡loss这强迫模型学习如何在设备间负载均衡。V3专家力粒度进一步细分258个专家激活8个。在topk阶段把专家打分的softmax换成了sigmoid因为softmax中每个人的得分都是和其他人分数相关的给一个专家分数提高会挤压其他专家的分数容易导致赢家通吃负载不均。而sigmoid每个专家的分数都是独立计算的。另外去掉了aux-loss辅助损失也就是loss里的负载均衡项改为一个偏置项b如果一个专家这轮负载过重了就减小b这也可以实现负载均衡但不动loss项不会让模型为了负载均衡强行把token塞给不合适的专家。当然还保留了一点负载均衡loss不过不是在token级别的了而是在序列级别的也就是处理一个序列各个专家之间需要负载均衡。MLA另外V3的另一个优化是MLA多头潜在注意力这和MoE没啥关系但也是个重要的优化。核心思想是推理过程中一个经典优化是KV Cache但是cache很占显存想要压缩cache大小于是考虑把kv cache做一个下采样压缩到一个潜空间(latent space)存入显存使用时再上采样解码。这能带来显存优化从而可以跑更长的上下文更大的模型同时减少搬运量还能优化吞吐。另外惊喜的是这样似乎还能优化模型表现这可能是因为kv cache有冗余信息多加了两层下采样下采样让模型学会了提取最重要的信息。MTP另一个优化点是MTP在主要的大部分注意力block结束后最后接几个简单MTP block每层基于前面的结果预测一个token这样一次完整的前向传播可以预测多个token