AI Infra学习笔记2MoE模块分析——MoE的计算原理和过程最近一直在整理AI Infra相关的学习笔记这篇讲的是MoE架构也就是Mixture of Experts、混合专家模型。我先把结论放在前面MoE几乎是当前大模型突破密集模型规模瓶颈最务实的路线之一它不改变模型结构的基本骨架而是通过“条件计算”的思路把模型容量和计算代价解耦。很多刚接触AI Infra的朋友一听到MoE第一反应是“又是个新模型结构”其实它的算子基本还是Transformer那套东西真正复杂的是数据流、并行方式和通信模式。这篇笔记我会从MoE的计算原理讲起重点放在计算过程、路由机制、负载均衡、并行实现和训练推理差异上。适合正在做大模型训练、推理加速、或者想深入理解模型底层运行逻辑的工程师参考。内容以我的理解和实测记录为主不保证适合所有人但至少能帮你少走一些弯路。我一直觉得搞懂MoE的计算原理才能真正看懂为什么它在AI Infra层面会带来那么大的通信和显存压力。所以这一篇我们先把计算这件事掰开揉碎。1. MoE是怎么做到“模型更大、计算反而更省”的1.1 从Dense到MoE为什么要把“一个模型”拆成“一组专家”传统Transformer模型前馈网络FFN这一层是稠密的不管输入什么token都要跑一遍完整的矩阵乘法。模型参数从70B涨到200B、500B计算量基本是线性增长的。换句话说你想让模型记住更多知识、拥有更强的推理能力就得把参数量堆上去随之而来的是训练成本和推理延迟的同步上涨这就是经典的“密集模型之墙”。MoE的思路是完全不同的我不需要让每个token都经过全部参数而是学一个“路由”能力让不同token只激活一小部分参数。比如一个模型总参数量200B但每个token实际只激活20B参数那么单次前向的计算代价就接近一个20B的Dense模型同时模型容量依然是200B级别的。这个思路在业界不是新东西早年NLP里就有混合专家模型真正让它爆发的是Scaling Law时代大家发现它能在不无限堆算力的情况下继续提升模型能力。从AI Infra的角度看MoE本质上是“把一系列小模型拼成一个大模型再训练一个调度器”。这个调度器就是路由网络它决定了每个token去哪些专家那里“办事”。所以MoE的计算过程至少包含两个核心部分专家网络本身的FFN计算以及路由网络的调度决策。后者虽然计算量小却牵动整个数据流的走向是系统优化的关键。1.2 条件计算MoE省计算量的底层逻辑MoE的底层逻辑是条件计算Conditional Computation。意思是模型对不同的输入走不同的计算路径。传统模型是“不管你是谁统统过一遍”MoE是“看一眼你是谁只带你去该去的地方”。这和推荐系统里的外环粗排、内环精排有点像也和人脑中“你不需要动用全部神经元才能识别一个苹果”的原理类似。这个机制的收益体现在两方面。一是训练效率。在相同的训练FLOPs预算下MoE模型可以做得比Dense大很多最终性能往往也更优。行业里普遍的经验是MoE在同等FLOPs下效果优于Dense这也是为什么很多超大模型最后都走MoE路线。二是推理效率。虽然MoE把参数摊大了但每次推理只需要激活一小部分参数那么单次请求的延迟理论上接近小模型这给了“超大模型、可控延迟”的工程窗口。不过天下没有免费的午餐。条件计算带来两个麻烦第一路由需要处理“token去哪个专家”的分派如果分配不均部分专家过载、部分专家闲置整体的算力利用率会大幅下跌第二专家分布在多张卡上token被路由到别的卡上的专家时就会产生跨卡通信训练和推理都会因此背上额外的通信开销。后面我会分别展开讲。2. MoE的计算过程从输入到输出的每一步2.1 总体计算流Transformer里多出来的两段逻辑先看一个标准的MoE Transformer层定位。绝大多数大模型MoE都把专家网络放在FFN位置替代原来的稠密FFN。Attention部分保持不变。所以MoE层的前向过程可以拆成下面的步骤输入token经过LayerNorm得到隐藏表示h_i。门控网络Gate/Router接收h_i输出所有专家上的打分logits。对打分取Top-K选出K个专家同时对打分做一次Softmax得到归一化的权重。把token的表示发给被选中的K个专家每个专家做一次FFN计算得到输出。将K个专家的输出按照权重加权求和得到MoE层的最终输出。输出经过残差连接和LayerNorm进入下一层。其中步骤1、2、3是MoE特有的步骤4是专家本身的计算步骤5是合并。这里注意第3步的Top-K和Softmax的顺序是有讲究的。一种做法是对所有专家先Softmax再取Top-K另一种是先Top-K再对选中专家做Softmax。工程上一般先用Top-K确定候选再在候选内做Softmax因为门控的logits往往存在绝对值较大的负数直接全量Softmax会稀释选中专家的权重。如果用伪代码来描述流程大致是h layer_norm(x) # 输入 logits gate(h) # 形状 (B*T, E)E为专家数 top_k_values, top_k_idx topk(logits, K) # 得到权重和索引 probs softmax(top_k_values) # 对选中专家做归一化 for i in range(K): expert_out expert[top_k_idx[i]](h) # 每个专家做FFN out probs[i] * expert_out out x out描述并不难但工程实现里每一步都有很多变体。下面我拆开讲。2.2 门控网络与Top-K路由打分、筛选、归一化门控网络通常就是一个Linear层输入维度是hidden_size输出维度是专家数量E。它做的事情非常直接给每个专家打个分。打分之后我们一般取Top-1或Top-2。Top-1是每个token只去一个专家最省计算但专家利用率可能不稳所以很多模型用Top-2来缓解单一专家过拟合的问题同时保持稀疏激活的比例不大。Top-K的选择是MoE架构最关键的超参数之一。K越大每个token涉及的专家越多模型的表现越接近Dense但稀疏性越低通信和计算开销越大。主流大模型通常K1或2也有部分模型做Top-6之类的设置但那是为了配合更复杂的专家分组策略。从AI Infra的视角来看K值影响的是激活参数量和跨卡通信量K每增加1通信量几乎翻倍。权重归一化的细节也值得注意。路由权重最终要作为专家输出的线性组合系数所以必须在选中的K个专家之间做一个归一化。如果没有归一化模型整体输出的尺度会漂移收敛会不稳定。大多数实现用的是软max不过也有工作指出直接用原始logits除以温度也能收到相似效果这里不再展开。另外门控网络的训练是端到端的。它和专家网络一起通过反向传播优化目标是学会把token分到能够最小化损失的专家。这个“目标”是隐式的模型没有显式标注“哪个专家该处理什么”一切都靠梯度信号自己涌现。这导致一个很经典的ai infra问题如果没有约束门控网络很容易把所有token都发给同一个少数专家出现路由坍缩专家的多样性就没了MoE也退化成一个小Dense模型。2.3 专家FFN的计算细节与参数量分布专家本身就是一个标准的FFN结构包含两个线性变换和一个激活函数。设输入维度是d_model专家中间维度是d_ff那么专家的参数量大约是2 * d_model * d_ff两个矩阵加上bias。MoE总参数量就是专家数乘以单专家参数量再加Attention和门控网络。举个例子一个拥有64个专家的MoE模型d_model4096d_ff11008那么单个专家参数量约为9000万64个专家就是约57.6亿甚至可能超过Attention加Embedding的总和。这解释了为什么MoE模型的参数量普遍很大但计算量却没有等比例上涨。还有一个时常被忽略的关键点专家的维度设计。如果把d_ff减小、专家数量增多理论上模型容量可以做得更大但每个专家的“能力“会下降反之专家少而大就接近多个大的Dense子网络稀疏性下降。这个调节关系是模型架构设计的核心杠杆。针对AI Infra一个常识是专家数量越多路由矩阵的并行和通信压力越大训练框架的调度开销也越高。所以不能单纯追求专家数量多要结合集群拓扑来做取舍。3. 路由策略与负载均衡MoE里藏在暗处的性能杀手3.1 负载不均衡是怎么产生的MoE训练得多了几乎每个人都会碰到专家负载不均的问题。直观原因是训练初期门控网络随机初始化会把token随机分布但训练一段时间后某些专家在特定类型数据上表现好门控网络就会倾向把相同类型token发给这几个专家形成“富者愈富”的马太效应。在数据并行专家并行的典型训练框架里负载不均衡带来的问题非常直接。全局计算是同步的所有GPU要等最慢的专家算完才能进入下一步。如果一个专家被分配了2倍的token整个step的耗时就会被拖长到接近2倍其他专家都在空转等它。这就是负载均衡问题在AI Infra层面最直接的表现。此外负载不均还带来显存利用率问题。每个专家在训练时需要存储该批次所有token的激活值以便反向传播使用。如果某个专家的token特别多这块显存放不下即使整个集群的显存总量还有富余这个专家所在的那张卡也会OOM。所以MoE的负载平衡不只是优化流畅度问题而是直接决定任务能不能跑起来的问题。3.2 辅助损失、专家容量和随机路由为了对抗路由坍缩主流方案是给门控网络加一个负载均衡辅助损失。最经典的实现是Switch Transformer里提出的aux loss设E为专家总数T为当前batch的token总数f_i是分给专家i的token比例的平均值P_i是门控网络对专家i的平均路由概率辅助损失定义为 E * sum(f_i * P_i)。当token均匀分布时f_i接近1/EP_i也接近1/E这个值趋于1越不均衡值越大。最终公式是总loss 原始任务loss alpha * aux_lossalpha一般取0.01这个量级。不过辅助损失是软约束实践中还是会出现部分专家过载。所以很多训练框架引入了“专家容量”的硬约束。专家容量 token总数 / 专家数 * capacity_factor其中capacity_factor通常设为1.25附近意思是允许每个专家最多接收1.25倍平均份额的token。超出容量的token不会排队等待而是被当作残差直接跳过专家走残差连接输出。这种token叫dropped token虽然理论上丢掉了一部分信息但能阻止单个专家拖垮整个训练step。随机路由是另一种思路训练时给路由的概率加一点噪声比如乘以一个随机噪声项让token偶尔流向非最优专家提高专家的覆盖率和多样性。但随机性在推理时会去掉保证结果可复现。这些方法可以组合使用实际部署时通常先把aux_loss加上再设一个合适的专家容量而不是只依赖某一种手段。3.3 路由策略的工程实现与调参经验在实现层面路由过程涉及两个独立的任务打分和派发。打分在PyTorch里就是一次矩阵乘法没什么复杂度。派发则涉及数据处理需要把token按专家ID重新组合然后分批送到对应的专家那里。训练框架里这个过程往往是分组矩阵乘GroupedGEMM或者通过张量拼接来做而在分布式环境下就是All-to-All通信。我踩过的一个坑是专家容量设得太小。为了追求极致负载均衡我把capacity_factor压到1.0结果大量token被dropped模型收敛速度肉眼可见地变慢loss周期性震荡。后来发现在训练的初期阶段路由完全没学到规律drop比例高得吓人把capacity_factor调大到1.5前期的loss曲线才稳定下来。我的建议是训练初期用偏大的容量训练后期再逐步收紧甚至可以动态调整让负载均衡辅助损失去接管细调。另一个调参经验aux_loss系数不要一味加大。系数太大门控网络会优先把token均匀分出去而不是按真实需求去分专业分工被破坏模型效果反而变差。常规律是alpha在0.001到0.01之间如果负载仍然不均优先查是不是数据分布本身太偏而不是直接加大惩罚。4. AI Infra视角MoE的并行策略、通信开销与显存分析4.1 专家并行与All-to-All通信的基本逻辑MoE在分布式训练里的核心并行策略是专家并行。简单说就是不同的专家放在不同的GPU上token被发往对应的GPU算完之后再把结果传回来。具体做法是先做常规的数据并行或者张量并行到MoE层时把路由结果汇总成“哪个token需要哪个专家”然后通过All-to-All通信把token跨卡交换过去每个GPU算完自己管辖的专家后再次做All-to-All把结果送回原卡。这种通信模式和常规流水线并行里的点对点通信不同它是全局性的每张卡都要和其他卡交换数据。想象一下64个专家分布在64张卡上路由Top-2那么一次All-to-All就把约两倍于token量的数据打到全集群里去了。通信量正比于token数、专家并行规模、隐藏维度也正比于K。通信本身是纯开销不产生任何计算收益所以在AI Infra优化里减少通信、隐藏通信是MoE训练的重要命题。我们常说的“所有专家都在所有GPU上复制一份”那是不叫专家并行那叫纯数据并行加MoE每个GPU都存全部专家参数显存开销巨大。专家并行是为了分摊参数存储和计算压力代价就是通信。在工程上专家并行程度越高单卡显存越低通信占比越高反之亦然。4.2 训练框架里的通信量估算我们做一个简单的量化分析。假设每个token的hidden size是d_model每step的token数是N即batch_size * seq_len专家数为ETop-K为K专家并行数是E一个专家一张卡。训练时每个token要发给K个专家K个专家的输出也要回收。所以每层MoE通信量的量级是 2 * N * K * d_model单位为字节按float32算再乘4。举个例子N4096个tokend_model4096K2则单层MoE通信量为 2 * 4096 * 2 * 4096 * 4字节约537MB。假设模型有8个MoE层那么一个step的通信总量大约4GB。这不是一个很大的数字但问题在于它是跨卡的全互联通信单条链路的带宽有限多卡同时通信会产生严重的网络争用。在千卡规模集群上通信时间完全可以占到一个step总时间的30%到50%这个比例会直接影响训练吞吐。针对这里有一些常见的手段一是计算和通信重叠在专家计算的同时预取下一批token二是梯度压缩用bf16通信替代fp32三是做本地路由优化尽量让token落在本卡或本机减少跨机通信。需要注意的是All-to-All的通信均衡也会被路由影响如果某些专家集中接收大量token与之相关的卡的网络负载也会集中出现峰值。4.3 显存占用分析参数、梯度与激活显存占用同样需要分开算。参数方面MoE总参数量变大但分布在更多GPU上单卡占用的参数反而可能比同规模Dense小。梯度存储与参数成正比一般占用参数的一半到两倍取决于优化器。激活显存可能是MoE最特殊的地方每个专家在处理token时都会产生中间激活这些激活是分散存储在专家所在卡上的。训练时token被分派到专家所在卡激活也跟着过去了反向传播时需要重新使用这些激活。这带来两个影响第一路由分配不均衡会导致某张卡激活显存过高第二即使是均衡分布专家并行下每张卡上的激活总量相比数据并行下的Dense模型反而可能更低因为在数据并行中一份激活可能在全卡复制而MoE把token均匀打散了。所以MoE并不天然更吃显存它吃的是通信和调度复杂度。推理场景则相反。推理时一般不保存激活做反向传播显存大头是权重和KV Cache。MoE模型的KV Cache部分和Dense没有本质区别因为Attention结构不变但权重部分因为专家参数众多即使稀疏激活全部专家参数依然要加载到内存中。这就引出了MoE推理的一个经典现象显存占用看全部专家计算耗时看激活专家。4.4 推理侧的特殊性预填充与解码阶段的差异MoE在推理阶段的难点和训练完全不同。训练阶段我们关心吞吐可以在batch里同时处理大量token做高效的分组矩阵乘推理阶段尤其是解码阶段每一步只生成一个token即使batch里有多个请求每个请求也是逐个生成路由和专家的计算都面临严重的批大小波动。预填充Prefill阶段一个请求的整个prompt一次性进来token数量多MoE可以像训练一样把token派发到各个专家批量矩阵乘的效率很高并行度充足。解码Decode阶段则完全不同每个step只新增一个token虽然可以batch多个请求但稀疏激活意味着每个请求路由到不同专家各个专家收到的token数很少矩阵乘变成小矩阵甚至单样本GPU算力利用率大幅下降。更关键的是MoE推理时所有专家权重都必须待在显存里但每次只激活一小部分。如果把全部专家放在一张卡上权重加载的开销会吃满显存带宽。我曾在一个8卡推理部署里实测过随着batch增大MoE推理的瓶颈根本不是算力而是需要不停访问专家权重带来的显存带宽压力。很多推理框架因此采用“在线专家加载”策略把不活跃的专家暂时卸载到CPU内存或NVMe磁盘上但这又引入主存/磁盘带宽瓶颈。所以MoE的推理优化本质上是一场内存带宽和延迟的博弈需要结合部署场景选择合适的专家放置方式。5. 实操过程中的常见问题与排查技巧5.1 训练时掉点、不收敛先查路由我处理过不少MoE训练异常第一个排查点永远是路由分布。如果某个专家被疯狂选中或者大部分token集中在少数几个专家上模型几乎必然出问题。最直接的排查办法是周期性打印路由统计信息统计每个专家接收的token数量占比、平均门控概率、被Drop比例。这些信息在训练日志里应该被记录为自定义指标。如果均衡性没问题但loss还是异常再看辅助损失的系数和专家容量设置。经验是先跑一个小规模的实验batch固定随机种子对比不同capacity_factor下的token drop比例。drop比例高于2%就要注意了超过5%基本可以断定路由压力太大。这个经验来自我调MoE的一次实际经历一开始drop比例接近12%loss怎么都压不下去把capacity_factor从1.15提到1.5之后训练曲线立刻恢复平滑。5.2 跨卡通信变慢怎么定位网络瓶颈如果你在大规模集群上跑MoE训练发现step时间比预估长很多大概率是All-to-All通信出了问题。先确认通信时间和计算时间的占比很多profiler工具可以统计。如果通信占比过高下一步检查是不是网络拓扑的问题专家被分配到跨机跨交换机的位置通信路径比本机内长得多。一个有效的做法是调整专家放置策略让专家尽量分布在同一台机器或同一个pod内。比如8卡一台机器就可以把专家数设为8的倍数并尽可能让一个MoE层里的专家落在同一台机器的8张卡上这样All-to-All会变成机内通信带宽高两个数量级。这个优化在部署推理服务时同样适用。另一个常见问题是每张卡的batch大小不均匀导致All-to-All出现热点解决办法是在每次路由后做一次微均衡把过多token重新调度给空闲专家。5.3 推理阶段OOM与KV Cache的取舍MoE推理部署时最常见的OOM场景是使用全部专家权重加KV Cache总量超过单卡显存。这时很多朋友会选择减少batch size或者降低KV Cache精度但有另一个思路MoE推理并不需要所有专家同时驻留显存可以把最不常被路由到的专家放在低速介质上作为“冷专家”按需加载。我实测过冷热专家比例一般可以在7比3左右对延迟影响很小但显存压力大幅缓解。不过要小心推理框架对动态加载专家的支持。如果框架不支持自己写的话需要锁粒度、并发控制和预取复杂度不低。我的建议是按业务需求来如果并发低、不需要硬性满足SLA用冷热专家切换如果并发很高还是买显存或者考虑减少专家数量、增大d_ff的架构方案牺牲一点稀疏性来换稳定性和开发效率。5.4 一个快速自查清单路由分布专家间token占比标准差是否过大Drop比例是否在可接受范围。辅助损失aux_loss的数值是否持续上升或者从训练一开始就特别大。通信耗时profiler中All-to-All时间是否被非均衡路由拉长网络拓扑是否为跨机瓶颈。激活显存是否有单卡显存明显高于其他卡观察是否由某几个过载专家导致。推理步长生成阶段主要耗时是否花在专家权重读取而非矩阵乘上由此判断是否应启用冷热专家策略。6. 一点实操心得MoE工程化的核心权衡这篇笔记写到这里最大的感受是MoE的工程难点不在“计算”本身而在权衡。参数做大的同时控制激活规模专家增多的同时控制通信开销负载均衡约束和模型效果之间存在张力训练吞吐和推理延迟的需求又未必一致。每次改动一个超参数影响的往往不是单点而是整条链路上的显存、通信、收敛行为。我建议所有做MoE相关工作的工程师把路由机制的细节彻底搞清楚再去碰训练框架的并行配置。很多人一上来就调专家并行度、调通信buffer结果问题出在门控网络收敛不稳或者专家容量设置不合理越调越乱。先通过日志和可视化理解token和专家的对应关系再动手改系统参数才能做出真正有效的优化。如果你也在做MoE的AI Infra工作欢迎把这篇文章当作一份来自一线操作的参考。我的经验未必适用于所有规模、所有框架但这些权衡和排查思路大概率能帮你省去不少调试时间。下一篇笔记我会写MoE在训练框架里的具体实现包括PyTorch原生实现和主流训练框架对MoE的抽象方式到时候再详细展开。