MoE计算原理与执行过程：从Router打分到专家并行与负载均衡

MoEMixture of Experts在AI Infra圈子里已经不算新鲜词了但凡训过千亿级模型几乎绕不开它。继续之前的AI Infra系列这篇把重心完全放在MoE的计算原理与执行过程上从单个token怎么被Router打分到专家并行之后token怎么跨卡流转再到capacity、token drop、负载均衡这些训练里天天碰到的坑尽量用一条完整的计算链路把它们串起来。阅读对象是刚从单卡Dense训练转向分布式稀疏模型的新手以及在训练稳定性上反复碰壁的工程师。这篇文章不会停留在“MoE就是多个FFN加个路由”这种层面而是要讲清楚每一步张量是怎么变的、通信发生在哪个节点、容量因子设多少才合理以及为什么你的负载均衡loss不起作用。1. 为什么要折腾MoE稠密模型扩展的瓶颈1.1 参数量翻倍但FLOPs没翻倍的反直觉现象我们平时训练一个稠密Transformer把隐藏层从4096扩到8192参数量大概翻两到三倍FLOPs也几乎同比例上涨。这意味着“模型变聪明”的代价是“训练时间变长”两者强绑定。MoE打破了这层关系——它把原来一整块大的FFN拆成E个相对独立的专家FFN每次每个token只走其中K个。总参数量大得吓人但激活参数和计算量都被稀疏性锁住了。举个例子一个稠密模型FFN参数量是MMoE化之后如果有8个专家总FFN参数约8M甚至更高。但在K2的情况下每个token只激活2个专家实际参与计算的参数是2M比原来稠密的M只多了一倍而不是8倍。这就是“条件计算”的核心价值花的算力不跟着参数总量走只跟着激活量走。这带来一个很现实的好处你可以在同样的算力预算下把模型参数撑得更大期望获得更好的效果。前提是路由选择得足够聪明并且负载不失衡。如果路由塌缩——所有token全涌向同一个专家——那稀疏性就白瞎了热点专家的计算量甚至比稠密模型还大。1.2 条件计算MoE不光是把模型变宽很多人刚接触MoE时会下意识把它理解为“多个模型集成”。这种理解不太准确。集成学习是多个完整模型独立训练、最后做决策MoE则是在一个模型内部做动态子网络选择Router和所有Expert是端到端联合训练的。Router输出的不是最终答案而是决定“谁参与计算”的权重专家看到的是同一个模型中Attention层传下来的隐藏表示。从计算图上看MoE其实不是增加了一条并行分支而是把原来Transformer Block里的Feed-Forward层替换成了一个“路由分发多专家并行加权合并”的组合。Attention部分和LayerNorm等组件仍然是稠密的只有FFN部分被稀疏化了。这个定位很重要它决定了MoE并行方案的设计边界你的专家可以拆到多张卡但Attention必须和其他层一起保持某种一致性。1.3 计算原理的第一步稀疏激活到底激活了什么所谓稀疏激活指的是每个token在MoE层中并不会经过所有参数路径。它只经过Router选择出的K个专家。这个K就是稀疏程度的调节钮K1就是Switch Transformer那种极稀疏模式K2是GShard和早期MoE常用的配置K大于4就逐渐往稠密靠了。稀疏激活在数学上表达为从E个候选函数中选择K个把输入映射到输出输出是选中专家输出的加权求和。这个式子看似简单但落地到分布式系统里会触发一连串问题专家放在哪张卡、Token怎么送过去、结果又怎么送回来容量怎么定丢掉的token怎么处理。这些会在后面几节逐一拆开。2. 单卡视角下的MoE计算流从Router打分到加权合并在进入分布式之前先把单个设备上的计算流程捋清楚。MoE层在网络中的位置是输入向量经过Attention之后得到一个形状为[N, H]的张量N是当前batch内的token总数H是隐藏维度。接下来进入MoE层流程可以分成四个阶段。2.1 Router打分所有token对全部专家的一次线性变换Router本质上是参数矩阵W_r形状是[H, E]。它对输入的每个token都做一次线性变换得到logitslogits x W_r # [N, E] probs softmax(logits, dim-1)这里的E是专家数。也就是说每个token都会对每个专家产生一个匹配度分数所有token对所有专家都做了计算。这部分计算量本身不大但对显存带宽敏感因为它涉及一次[N,H]×[H,E]的矩阵乘法。有一些工程实现会把Router做成分组或低秩版本就是为了压这部分开销。Router打分还有一个容易被忽视的细节在早期Noisy Top-K Gating的论文里作者特地提出在logits上加一个噪声项来帮助负载均衡logits GaussNoise * softplus(noise_param)这样可以让不那么靠前的专家也有机会被选中避免路由过快地扎堆。正式训练时通常逐渐把噪声关掉或降低只在初始阶段帮助探索。2.2 Top-K选择离散决策与梯度回流拿到softmax概率后对每个token选出K个概率最大的专家下标同时保留对应概率值。关键是接下来要对选中的K个概率做一次重新归一化让它们的和为1。这一点不同框架的处理有差别——有些直接使用softmax概率原值有些做归一化。多数实现会选择归一化因为这可以让门控权重在语义上更接近“混合比例”。这里有个有趣的问题TopK是离散的不可导操作Router怎么拿到梯度梯度的秘密在于门控权重本身就是Router的softmax输出所以损失函数通过最后的加权求和结果传播到这些权重再通过softmax反向传播到Router参数上。专家的选择路径是离散的但选出来的权重是连续的梯度正是沿着这个连续通道回流的。这也是为什么Router能学出来而不是靠强化学习硬搜。2.3 专家前向与加权合并选出了K个专家之后每个专家网络通常是标准的FFN结构一个升维线性层、一个激活函数、一个降维线性层。上面的示例可以扩展成更常见的SwiGLU结构三个参数矩阵输入先经过gate_proj和up_proj两个分支再用silu(gate) * up做逐元素相乘最后由down_proj投影回隐藏维度。合并时每个token的输出是y sum_i( gate_i * expert_i(x) ) # i 遍历被选中的K个专家单纯从数学上看整个过程就是一个“根据Router选出的稀疏路径做加权求和”。但到了多卡环境这一步会和通信强耦合。单卡上看不出来多卡上你才会发现瓶颈基本不在专家计算本身而在token的搬运。2.4 一个简化的全流程示例代码下面给一个偏概念性的PyTorch风格示例帮助把流程串起来。这不是完整框架代码但足以反映MoE层的主干逻辑import torch import torch.nn.functional as F E 4 K 2 N 32 H 128 x torch.randn(N, H) router_weight torch.randn(H, E) * 0.02 logits x router_weight probs F.softmax(logits, dim-1) # 选出每个token的topK专家 scores, indices torch.topk(probs, K, dim-1) # 对入选的K个权重重新归一化 scores scores / scores.sum(dim-1, keepdimTrue) # 4个专家每个专家用简单两段FFN代替 experts [ torch.nn.Sequential( torch.nn.Linear(H, 4 * H), torch.nn.GELU(), torch.nn.Linear(4 * H, H), ) for _ in range(E) ] out torch.zeros_like(x) for e in range(E): # 找出哪些token选了专家e mask (indices e).any(dim-1) # [N] if mask.any(): h x[mask] expert_out experts[e](h) # 只对被选中的token做计算 weight scores[mask].mean(dim-1) # 简化处理实际要按每个token的选中次数 out[mask] expert_out * weight[:, None] # 残差连接在层外统一处理这段代码最大的意义是展示了mask这条逻辑不同专家收到的token数量不一样所以不能简单地用N x H的大矩阵并行过所有专家而必须动态地“分桶”。到了分布式场景每个桶就是一组要发给同类专家的token这会引出一对通信原语dispatch和combine。3. 专家并行后的真实成本Dispatch/Combine与AlltoAll3.1 为什么单卡装不下专家的摆放与局部性如果E比较大比如64个专家每个专家500M参数光专家参数就32B这还没有算Attention和其他层。单卡装不下就必须做专家并行。最直白的方案是把专家按设备切成E_P组每个设备负责一组专家。Attention层和Router层仍然是稠密的一般配合张量并行或序列并行来切分。这样摆放会引起一个严重问题每个设备上的token是随机分布的但只有本设备上的专家才能本地计算。一个token选中了专家3但专家3在设备2上它的隐藏向量就得从当前设备传到设备2。这本质上是一个稀疏通信任务而非简单的全量广播。3.2 dispatch阶段把token按路由结果送到对应设备在专家前向开始之前需要做一次数据交换。过程是这样的每个设备先把本地token的Router打分表整理出来得出“哪些token要去哪个专家”。接着根据目标设备把token打包通过通信把包发给目标设备。常规做法是使用AlltoAll通信原语每个进程都有数据要发给其他所有进程也需要从其他所有进程接收数据。执行时每个进程维护一个长度等于进程数的交换列表。真正做的时候要小心一个细节AlltoAll要求每个进程发往目标进程的数据长度事先已知这样才能预分配接收缓冲区。因此必须先用一次通信把每个专家对应的token数量汇总或者用定长capacity下一节会详细说来预先锁定缓冲区大小。否则就需要使用可变长通信会增加复杂度。3.3 combine阶段结果回收时的不对称性专家前向完成后每个设备上存着“被本设备专家处理的token的输出”但输出需要回到token原本所在的设备因为后续的残差连接、归一化和Attention都在原设备上。这个回收过程就是combine一般也要一次AlltoAll把结果按原路径送回。combine和dispatch的不对称点在于dispatch时你要搬运的是原始隐藏向量combine时你搬运的是专家输出而且回来之后还要做一次加权合并。如果Router选择的权重没有伴随token一块儿搬走就要在发送侧缓存好对应关系或者多携带一份门控权重。工程上常见做法是把权重和时间下标都打包发送这样接收方可以重建原始顺序。3.4 Capacity Factor、Token Drop与显存/通信的水位控制为了不让通信消息长度来回变化几乎每个MoE框架都会定义一个专家容量capacity概念。简单说每个专家的Token处理上限是capacity ceil(N / E * capacity_factor)N是当前样本中的所有token总数E是专家数。capacity_factor就是控制水位线的旋钮。设为1意味着完全均匀的理想情况设为1.25表示允许每个专家最多处理理论平均量的1.25倍设为2则是两倍。为什么不能设成1因为路由几乎不可能做到绝对均匀。如果设成1只要某个专家多分配了一个token就会有token被溢出处理损失信息。实践中常见的策略是训练时capacity_factor取1.25到2.0之间并在配置里写明“超出capacity的token会被截断”同时用下面的负载均衡loss去拉低溢出概率。如果一个专家收到的token数量超过capacity多出来的token就叫overflow。处理方式有两种流派一种是直接丢弃Switch Transformer训练时就是丢弃这些token不参与该层计算另一种是发送给下一个专家或者直接跳过MoE层效果因任务而异。但无论如何这些token在该MoE层的梯度任务被中断模型不会根据它们去更新对应的专家参数。实际训练中需要监控overflow rate这个指标。如果你发现溢出率偏高但负载均衡loss已经很低观察capacity_factor的设定是否过小。附带一个显存问题专家层的batch量天然是不均匀的某个专家收到特别多token时这个专家的升维矩阵batch可能非常大显存峰值被瞬间拉高。这也是为什么很多系统会做“同专家token重新group”并把几个专家合并成一次大GEMM执行以降低启动开销和显存碎片。4. 负载均衡MoE训练的“血压计”4.1 路由塌缩专家饿死与被冲垮的两种极端如果不加任何约束Router很快会发现一个捷径把所有token都路由到少数几个输出层偏好“更好”的专家。因为梯度会让那些被选中的专家继续变好它们更容易被选中于是形成正反馈。这会导致大部分专家几乎不参与训练“饿死”的专家等于没训练整个模型实际容量缩水。反方向的问题同样存在少数热点专家的token可能超过capacity大量token被丢弃整个层的信息损失严重。在这个意义上负载均衡不只是“公平性”问题而是直接影响模型质量的上限。不是“大家平均用”就好而是需要一套机制让路由的流量分布保持在一个合理区间又不要过度均衡到让每个专家完全同质化。4.2 辅助负载均衡Loss公式与直觉理解几乎所有主流MoE方案都会在分类或生成损失之外加一项辅助loss来拉负载均衡。经典的Switch Transformer负载均衡loss可以写成L_aux alpha * E * sum_i(f_i * P_i)解释一下f_i是分配给专家i的token比例P_i是Router对所有输入分配给专家i的平均softmax概率。直观上如果专家i既吸引了大量tokenRouter又在大部分输入上对它给出高置信度那f_i * P_i就大loss就高从而惩罚这种扎堆现象。反过来如果专家i虽然吸引了一部分token但Router对它的平均置信度不高惩罚就变小。alpha是权重系数常见取值在0.01到0.1之间。一些实现还会在f和P之外增加“方差惩罚”或者换成基于熵的正则项。不同框架差异很大但核心思路一致让Router的选择分布趋向均匀同时对“均匀”的强迫程度要适当。如果你发现负载已经平均得很漂亮但模型效果反而下降那很可能是alpha太大把所有专家训练成同一个平均模型了。稀疏性被过度压制。4.3 Top-1与Top-2对均衡性的影响为什么K越大越容易均衡Top-1是最极端的稀疏选择每个token只选一个专家。这种情况下如果某个专家稍强一点就会形成赢者通吃的局面负载均衡压力非常大。开关式的硬分配也让梯度更稀疏。Top-2时router权重分布在两个专家上即使第一选择偏科第二选择仍可能让另一个专家拿到梯度天然缓解了塌缩。K越大路由越平滑但激活参数量线性上升稀疏带来的节省就消耗掉了。另一个和K直接相关的工程点scores重新归一化后如果K2但两个专家权重差距很大例如0.9和0.1那么后者的梯度贡献是很小的。这也是为什么有些实现会对TopK分数施加温度缩放。不要小看这个缩放它会在计算图上成为一个可调的超参数影响路由熵。4.4 工程上怎么观察和调整均衡训练MoE我强烈建议你固定一下这些监控指标每个专家的token接受计数、每个专家的平均score、overflow rate、以及设备间AlltoAll的通信时长。如果某专家长期占用过多token除了调大alpha以外还有一个工程手段是专家复制——把热点专家复制到多个设备上。这样同一个专家逻辑上可以在多卡同时处理不同token降低单卡热点。如果发现某个专家长期没有token不要直接用“增加随机噪声”了事先排查是不是初始化把该专家的bias设得过大或者对应设备的通信超时导致该专家的token被偷偷丢弃了。有些隐藏在通信层的坑表面看是负载不均实际是某个进程的接收缓冲区没收到数据把本该给专家B的token默默丢掉了。5. 训练与推理的分岔路从显存规划到解码吞吐5.1 训练时的显存/计算模式为什么要AllToAll而不是逐卡拷贝训练阶段每步都要在全部专家参数上算梯度所以专家权重必须驻留在显存中。AllToAll方式的核心优势是每张卡只发一次收一次就能完成所有专家的数据交换通信复杂度与设备数和专家数都有关但不像逐对收发那样膨胀。在实际训练中通信量和batch size直接相关N越大发送的数据量越大这是训练MoE最大的扩展性瓶颈之一。你有时会听到“通信和计算重叠”的做法当某个专家在计算时下个micro batch的token正在传输这样能把AlltoAll的延迟藏起来。这就是流水线并行思想在MoE层的应用。我自己在做大规模MoE训练时最先优化的反而不是专家内GEMM的kernel效率而是把AlltoAll的通信切得足够细让它和算得最慢的那个设备对齐这个收益远大于微调GEMM的tile大小。5.2 推理时的显存与带宽困境参数总归要读一遍推理阶段和训练的区别非常大。训练时你在乎FLOPs和通信延迟推理时你在乎的是每一个token的延迟和整体吞吐。MoE在推理时的理论优势是激活参数只有K/E的比例所以计算量明显比稠密低。但现实是显存带宽往往成为瓶颈——即使只激活两个专家你也要把这两个专家的参数从显存读到计算核心如果专家在别的卡上还要加上网络延迟。稀疏性的收益在batch size很小比如decode阶段时会被显著削弱。为什么decode阶段更严重因为生成是串行的每个step处理的token数通常等于当前batch大小。batch小的时候每token激活的专家数量少但每个专家kernel的batch size也小GEMM的算力利用率低。更麻烦的是如果你把all-to-all通信安排在每次decode step中那么频繁的小包通信会让延迟上涨非常明显。业界常见的优化手段是prefill阶段使用大的batch提高吞吐decode阶段尽量减少跨卡通信甚至在某些请求量不高的时候只在一张卡上顺序加载专家。这会增加磁盘或内存拷贝但对时延敏感场景可能更划算。5.3 从配置参数看训练稳定性E、K、capacity_factor和alpha的搭配给一个比较常用的起点配置方便新人上手参数常见范围说明E专家数8 ~ 128专家数越多单个专家容量越小路由压力越大KTop-K1 ~ 4K1极稀疏但负载均衡压力大K2平衡性较好capacity_factor1.1 ~ 2.0训练时建议1.25起步显存紧张再降到1.1alphaaux loss权重0.01 ~ 0.1初始用0.01观察负载不均再逐步提高Router隐藏层无或少量有些实现在Router前加一个非线性的小net不建议过度复杂如果你的模型是8个专家、K2、capacity_factor1.25那每个专家最多处理N/8*1.25个token。当你把batch size调大N变大所有分段容量也变大显存方差更陡。经验是加大全局限流时优先调整capacity_factor而不是盲目加大alpha因为后者牺牲的是模型容量前者牺牲的只是极端情况的容错。5.4 几个常见的坑和排查思路训练MoE时会遇到很多过滤不掉的怪问题我把我实际排查链路里印象最深的三类列一下。第一类loss突然起飞。通常是某个时刻有大规模token drop喂给后续层的分布发生突变。先看overflow率曲线如果在一个step内爆升再回看是否batch size变大导致capacity没跟着scale。第二类某专家梯度显著高于其他专家。这不一定是它学得最好有可能是它的token分配数量大梯度累积多不妨比较单位token的平均梯度。如果某个专家的平均梯度异常检查该专家的输入分布是否有异常埋点。第三类通信等待时间长但是NVIDIA的带宽监控显示网卡利用率不高。这多半是AlltoAll的分段不合理某个小包阻塞了队列。改成按专家数均分消息块或在低流量阶段调整拓扑亲和性都会有效。6. 一些容易被误解的点和我的实操体会最后写几个我在给别人讲MoE时经常遇到的理解偏差。第一很多人认为MoE里Router是“学出来的注意力”。实际上Router的参数远小于注意力它并不是让token去搜索专家内部信息的机制而只是一个粗粒度的分发器。Router打分的依据是当前整层的隐藏状态它看不到专家内部更细的特征。真正学习精细分工的是专家本身。第二MoE不是“每层都要最好有几个专家”。有时只在部分层的FFN上做MoE就够其他层保持稠密可以显著省显存和通信。比如40层模型只在最后16层启用MoE前24层全部稠密训练成本能省下近一半。不要机械地认为“全MoE一定比部分MoE好”。第三推理的batch size之所以那么重要是因为它决定了专家kernel的几何形状。如果在线推理场景一直是小batchMoE的显存优势很难兑现。反过来离线批量打分的场景才是MoE最能发挥的舞台因为可以把大量token聚在一起一次性触发大batch的专家计算。如果让我给一句话总结那就是MoE把“计算规模”的问题转化成了“调度与通信”的问题。理解它的计算原理本质上就是理解一个动态、稀疏、多目标调度的系统如何优雅地塞进分布式训练框架里。下一期我大概会接着写专家并行和流水线并行夹在一起时那些调度细节是怎么纠缠的。先说这么多眼睛已经看花了。