★ FEATURED ARTICLE 从Transformer诞生起前馈网络FFN就是又笨重又绕不开的角色它占了模型近三分之二的参数量但每个token都得老老实实把所有参数过一遍。你扩大模型容量推理成本就跟着线性往上涨这让很多想做大规模模型的人非常头疼。MoE混合专家模型或者说稀疏前馈网络正是冲着这个矛盾来的——它把单个FFN拆成一组并行专家再加一道轻量门控路由让每个token只激活其中一小部分专家。参数总量可以扩展数倍到数十倍但单token的计算量基本维持原样容量和成本终于被解耦了。这篇文章我会从设计动机一路讲到训练和推理阶段的工程坑包括门控网络怎么工作、负载均衡为什么难搞、专家并行怎么通信以及一批可以直接上手的调参经验。不管你是想改造自己的Transformer还是想理解最近一系列MoE大模型的结构这篇都值得认真看完。1. 从 FFN 到 MoE改写前馈网络的真实动机1.1 一个容易被忽视的事实Transformer 的容量主要藏在 FFN 里很多人一提到Transformer下意识关注的是自注意力Self-Attention——tokens之间的交互全靠它。这话没错但你要是统计一下参数量注意力层其实只占很小一部分。一个标准Transformer Block里FFN通常包含两个线性变换先把隐藏维度比如4096升到中间维度比如11008再降回来。这两个矩阵相加参数量轻轻松松占总模型的2/3甚至更多。为什么设计成这个样子因为FFN负责的是逐位置的记忆与抽象每个token独立经过同样的非线性变换把注意力层已经汇集的上下文信息进一步加工成更高层的特征。可以这样理解注意力层负责商量FFN负责各自做事。商量只需要一个桌子做事却需要大量工具工具多了参数量自然就上去了。问题在于推理时即使工具再多每个token也得把这些工具全部用一遍。参数量翻倍单个前向的计算量就翻倍延迟也跟着上来。于是你面临一个非此即彼的选择要么保持计算量不变放弃扩展容量要么接受昂贵的推理成本。这个矛盾在模型规模做大之后尤其尖锐我当时在一个文本模型项目里试着把FFN中间维度从4倍扩到8倍效果确实有提升但显存和推理延迟的代价直接让我放弃了。1.2 稀疏激活的直觉让参数多、计算少变成现实MoE的思路一句话就能讲清楚既然每个token用不上所有工具那就把FFN拆成若干个可以独立工作的专家然后装一个分诊台门控网络让每个token只去它需要的几个专家那里其余专家本轮直接跳过。这就是稀疏激活的含义。传统FFN是稠密计算——所有参数对所有输入都参与计算MoE则是稀疏的——参数总量很大但单个token实际触发的参数很少。用专业一点的术语描述假设总共有E个专家每个token从路由分布里挑出K个专家K远小于E通常K取1或2那么这token的FFN计算量大约是原来单个FFN的K/E而不是E倍。直觉上可以类比全科医生和专科门诊的区别。普通FFN像一个能力全面的全科医生每个病人都由他一个人看诊所请不动更多医生病人再多也是他一个人忙。MoE则是一个开了很多专科诊室的医院每个专家只擅长某一类病分诊台根据病人症状快速分配科室每个病人只看一两个专家其他科室完全不打扰。医院规模可以很大但单个病人的看病时间并没有增加多少。模型容量来自有很多专科医生计算效率来自每个病人只看很少的医生。1.3 为什么改造的是 FFN而不是注意力层这里有个值得深思的问题为什么大家做MoE都动FFN而不直接把注意力层也拆成混合专家原因在于两种结构的Nature完全不同。注意力层在每层同时处理所有token通过Q、K、V矩阵建立跨token的交互每个token的结果都与其他token相关FFN则对每个token独立操作token之间没有任何信息交换。正因为FFN的逐token独立性我们才可以放心地把不同token送到不同专家专家之间不通信各自处理完再把结果放回原来的位置。如果对注意力层做路由比如让每个token只和一部分token交互那被谁注意到会因为不同token选了不同专家而变得极度碎片化上下文信息容易出现断裂训练难度也会直线上升。这不是说注意力MoE完全不可行而是在标准Transformer架构里FFN是成本收益最划算的改造点改动局部、逻辑清晰、收益最大。还有一点很实际FFN的权重是逐token独立存储的天然适合做参数并行——把不同专家放到不同GPU上互不干扰。注意力层则必须做完整的矩阵乘法和跨卡同步通信开销大得多。所以从工程角度FFN也远比注意力层更适合做稀疏化改造。2. 稀疏路由的完整工作链路门控网络、Top-K 选择与专家容量2.1 门控网络token 怎么决定自己去哪个专家门控网络Router也叫Gating Network是整个MoE最核心的组件。别看它决定着一大群专家的调用方式它本身的结构却非常简单通常就是一个线性层输入是token的隐藏表示h输出维度等于专家总数E。然后过一个Softmax得到该token对各专家的选择概率分布。用公式表示就是router_logits W_r hrouter_probs softmax(router_logits)这里的W_r是一个非常小的矩阵假设隐藏维度是4096专家数是8那W_r的参数量只有4096 × 8占整个模型的比重可以忽略不计。所以门控网络被称为轻量——它算得快、参数少却直接影响模型的容量利用率。一个关键的设计细节是Router的输入到底用哪一层表示。如果是用当前层的FFN输入那路由决策反映的是这一层之前的特征如果是用上一层的输出或者与attention输出拼接后的向量则可能包含更多上下文。不同实现各有取舍我个人的粗浅经验是在层数较深的地方用本层的输入就够如果训练不稳定可以尝试把上一层Router的logits也拼接喂进来给当前层更多先验。2.2 Top-K 选择与专家容量从分诊到预约挂号拿到全量专家概率后不能真的所有专家都参与计算那样又退回稠密了。做法是取概率最高的K个专家K通常取1或2。取到具体哪K个专家后对应专家计算得到输出再按路由概率或归一化后的权重加权求和得到当前token的FFN输出。当K1时每个token只走一个专家计算效率最高但每个token的感受域非常受限专家之间也完全没有冗余协作。K2是当前很多工作的经验默认值既保证了稀疏性又给了token一点选择的余地。我感觉K2相当于让分诊台不仅选最合适的科室还允许备选科室兜底对训练初期路由还不成熟时特别友好。接下来是专家容量Expert Capacity的问题。一个专家在某一层最多能处理多少个token如果不设上限某几个热门专家可能被分配大量token其他专家空转计算负载严重不均。于是实现时会为每个专家设置一个容量上限capacity capacity_factor * (T / E)。这个公式的含义是理论平均分配下每个专家应该处理T/E个tokencapacity_factor则允许一定的弹性冗余。举个例子假设T256个tokenE8个专家理论平均是每个专家32个tokencapacity_factor取1.25则容量为40。一旦某个专家被分配的token数超过了容量超出的token就只能被丢弃或者走一条残差连接直接跳过FFN的通路。很多人第一次看到丢弃token会吓一跳训练时丢信息模型还能学好吗实际上这是工程上的妥协——如果强行保留所有token计算图和同步逻辑会变得非常复杂还会导致不同专家之间的执行时间出现长尾。后续有人提出无丢弃drop-free的负载均衡方法核心是调整路由分配而不是限制容量效果也确实更好。但从工程复杂度考虑带容量的实现依然是目前最稳妥的基线方案。2.3 一个完整的 token 走向示例为了把上面这些概念串起来我走一遍一个token在MoE层的完整流程token输入为形状[B, H]的向量先复制一份作为残差为后续相加做准备。门控网络根据当前token的表示算出每个专家的概率得到形状[E]的分布。挑选概率最高的K2个专家记为专家3和专家7。token被送往专家3和专家7所在的设备如果是单卡就是同一设备上的不同矩阵。专家3和专家7分别对token做自己的FFN前向计算得到两个输出。两个输出按路由权重加权求和。加权结果与第一步的残差相加得到当前层的输出。这个流程里其他6个专家完全没有参与计算对应的参数也处于空闲状态。这就是稀疏两个字的真正含义——不是少设计了参数是每次只用到其中一小撮。3. 负载均衡与路由离散性训练阶段最头疼的两个工程问题3.1 离散路由不可微梯度怎么回到 RouterMoE听着完美但训练时第一个坑就摆在面前Top-K取的是索引是一个离散操作不可导。你没法把我选了专家3这个行为直接变成梯度信号来更新门控网络。怎么办经典解法是虽然选择是离散的但最终输出是对所选专家的输出按概率加权求和而概率本身是router logits经softmax得到的是连续可导的。于是梯度可以沿着概率权重这条路回传到路由网络。换句话说真正参与加权求和的权重是router_probs里选中那个位置的值而不是一个one-hot向量即使如此这仍然会产生一个什么问题呢问题在于概率权重在反向传播时数值通常比较小门控网络拿到的梯度信号偏弱。你可以理解为后端专家们学得很快但分诊台学得有点滞后。这带来一系列连锁问题后面会详细讲。还有一个细节容易被忽略训练时如果想用hard gradient估计算法类似Gumbel-Softmax的Straight-Through情况会稍有不同——forward时用one-hotbackward时用连续概率的近似梯度。但标准实现里大多数框架仍然采用加权求和式的软路由简单且稳定。我之前在某次实验里试过hard路由收敛速度明显变慢后来果断退回soft加权。3.2 负载坍缩少数专家过热、多数专家空转MoE最著名的训练困境就是负载坍缩load collapse。开始训练时Router的参数几乎随机各专家的分配还算均匀。但很快某些专家因为初始化稍好、梯度更新方向略微有利就会被分配更多的token形成正反馈热门专家学得更快、表现更好Router越来越倾向于选它其他专家则长期闲置梯度信号近乎为零逐渐变成死专家。最终整个模型退化成其实只用了两三个专家参数总量虚增容量优势形同虚设。解决负载坍缩的主要手段是加辅助负载均衡损失Auxiliary Load Balancing Loss让Router在每个batch里尽量把token均匀分配给专家。经典的实现大概长这样aux_loss E * sum(prob_per_expert * fraction_of_tokens_per_expert)其中prob_per_expert是所有token对本专家的平均选择概率fraction_of_tokens_per_expert是实际被分配到本专家的token占比。二者相乘再求和如果分配完全均匀这个值会趋近于1越不均匀值越大相当于对Router施加一个均匀化的惩罚。E作为缩放因子使得损失数值与专家数无关方便统一调系数。但这又引入新矛盾辅助损失权重设太大模型会倾向于为了均匀而均匀牺牲真正的专业分化整体表现下降设太小负载坍缩压不住。我在实际调参时观察到的情况是由于每个batch的分布都在变固定权重很难同时照顾训练前期和后期的需求。比较务实的方法是训练前几百步用小权重等专家分化初步形成后再逐渐加大或者观察负载的variance指标——如果某专家占比长期超过2倍平均值就临时加大权重拉一把。3.3 初始化与 Router 输入尺度的稳定性除了负载均衡Router训练的另一个隐性障碍是logits的数值波动。如果Router的输出一直非常大softmax出来的分布就会非常尖锐几乎等于固定选Top-1专家别的专家几乎没有机会。反过来如果logits太小分布太平滑Top-K的选中结果和未选中结果差异不大路由失去了区分能力。一个常用技巧是给Router加一个可学习的temperature项初始值设小一点让logits先温和一些。入口处做LayerNorm也是常见做法把不同token的表示尺度拉到同一水平避免某些高范数token支配Router的决策。我还用过一种更朴素的方法直接从已有Dense模型做热启动warm start。比如把稠密FFN拆成E个专家时克隆同一份参数作为每个专家的初始化这样所有专家在训练初期能力一致Router没有偏心的理由。这在多任务迁移场景下特别有效毕竟从零训练一个大规模MoE冷启动稳定性的难度是很多人低估的。4. 训练态与推理态的差异显存预分配、专家并行与通信开销4.1 训练时怎么摆放专家专家并行与 All-to-All 通信讨论完训练稳定性还得解决一个现实问题MoE参数总量大一张卡根本放不下。传统模型并行是切分每一层的矩阵MoE则有天然更优的方案——专家并行Expert Parallelism把不同专家分布在不同GPU上每张卡只保存一部分专家的权重Router所在位置则根据token的分布动态将token散播到对应专家所在的卡。这里有两个关键通信阶段。第一个阶段叫Dispatch某个token被路由到专家3而专家3在另一张卡上就必须把token的隐藏向量通过通信送到那张卡。第二个阶段叫Combine专家3算完输出后结果还要送回原来的卡供后续的残差相加和下一层使用。如果K2token实际上要做两次dispatch/combine通信量也随之倍增。这两个通信阶段使用的通常是All-to-All模式也就是每个节点都要向其他所有节点发送数据。这和传统All-Reduce所有节点只汇总一个结果完全不同。通信量大约正比于被路由的token数、隐藏维度和专家跨卡跳数。所以一个非常重要的工程教训是MoE的稀疏收益到一定规模后可能被通信开销吃掉。如果网络带宽不够你会发现训练速度不升反降加专家还不如加计算密度来得划算。我实际测试过一个典型配置8卡环境、每层8个专家、K2。在单机NVLink互联下通信损耗还能接受但一旦跨机走以太网训练step时间立刻涨了30%收益大打折扣。所以设计MoE规模时一定要先盘清楚自己的通信拓扑。4.2 推理态显存的真正瓶颈参数常驻而不是计算量很多人第一次接触MoE推理时都会有个错觉既然每个token只激活一小部分专家那显存占用是不是也按比例减小不是。稀疏只减少了计算量FLOPs参数依然全部常驻显存。无非是一个token实际只用到了门控网络和少数专家的权重但所有专家的参数都必须在内存里待命因为你不知道下一个token会访问哪个专家。这在推理侧会造成一个非常有趣的局面MoE模型的推理延迟可以接近同计算量的Dense模型但显存占用接近同参数量的Dense模型。换句话说你省的是算力省不了容量。正因如此大规模MoE推理基本都是多卡分布式方案每张卡放一部分专家。当请求到来时token先经过路由再通过通信调用其他卡的专家。这实际上把训练时的专家并行复用了只是batch更小、通信更多在卡间进行。显存优化通常还有两个方向一是对专家权重做更激进的低比特量化因为专家本身是逐token独立使用的量化误差不容易像attention参数那样跨token扩散二是预先把所有专家的参数做统一内存池管理配合显存碎片整理。后一种技巧偏工程化但收益很直观尤其在长期多请求服务的场景里。4.3 padding token 带来的路由损耗一个容易踩的隐藏成本训练时为了batch对齐我们会给短句做padding。在普通Transformer里padding位置虽然参与计算但因为注意力mask的存在它们对有效位置的影响基本可以被忽略于是人们很少专门去处理padding token的成本。但在MoE里情况变了每个padding token也要走Router、也要被dispatch到某个专家、也要占用专家容量。这带来两个直接后果。第一padding token占用专家容量可能导致有效token被挤掉超过capacity被丢弃。第二即使padding token没有实际信息它们依然占据通信通道白白消耗All-to-All带宽。好在我们找到的解决办法并不复杂在padding位置的路由概率上直接加一个非常大的负偏置让Router几乎不可能选中任何专家或者把padding token在进入MoE层之前直接从batch里摘除等所有专家计算完再放回原位。前者改动小后者更彻底但需要处理动态batch工程复杂度会增加一些。类似的浪费也出现在推理场景现在主流的推理框架普遍采用连续批处理continuous batching不同请求的长度不对齐时也会出现大量虚位token。所以做MoE推理服务时请求调度策略对性能的影响往往比模型结构本身还大。5. 实测调参清单与扩展思路从复现一个简单 MoE 层开始5.1 一个可直接改写的 MoE 层伪代码结构看完前面的原理你已经可以把一个标准的MoE层搭出来了。下面用离实现很接近的伪代码展示核心逻辑框架无关方便你对照着改写自己的代码。def moe_ffn(x, experts, router, K2, capacity_factor1.25): # x: [batch, hidden] B x.shape[0] E len(experts) T B # 这里假设每个token独立路由实际实现会考虑padding mask # 1. 计算路由概率 logits router(x) # [B, E] probs softmax(logits, dim-1) # [B, E] # 2. 取Top-K专家 top_probs, top_idx probs.topk(K, dim-1) # 各[B, K] top_probs top_probs / top_probs.sum(dim-1, keepdimTrue) # 归一化 # 3. 按专家容量筛选 capacity int(capacity_factor * B / E) dispatch_mask build_dispatch_mask(top_idx, E, capacity) # [B, E, K] 或管理token分配 # 4. 对每个选中的专家执行 FFN并按权重聚合 output zeros_like(x) for e in range(E): selected dispatch_mask[:, e] # 哪些token被分到专家e if selected.any(): expert_out experts[e](x[selected]) # [num_selected, hidden] output[selected] (top_probs[selected, e] * expert_out) # 5. 残差连接 return output residual_x当然真实工程里不会用for循环逐专家算而是通过矩阵乘法一次性完成所有专家的前向但核心逻辑就是上面的流程。为了效率专家计算通常会合并成一个大的张量运算把专家参数沿着一个轴堆叠起来然后用grouped GEMM一次算出所有选中专家的输出。这也是很多训练框架在MoE部分做的主要优化。5.2 调参顺序先稳定路由再追容量调MoE和调普通模型一个很大的区别是超参数之间存在明显的先后依赖关系。一上来就乱试专家数、路由权重很容易出现忘记调了什么的局面因为每个超参数都会间接影响Router的负载分布。我建议按下面的顺序做超参数常用范围影响效果调整优先级专家数 E8~64视参数量预算决定容量上限和并行度先定预算再选Top-K1~2计算量、冗余度、通信量建议固定为2capacity_factor1.0~1.5丢弃率、负载余量、性能优先调aux loss 系数0.001~0.1负载均衡强度次优先Router温度0.5~1.0可学习路由分布尖锐度小步微调专家Dropout0.0~0.1缓解过拟合、增强鲁棒性最后调具体来说先固定E和K把capacity_factor从1.25开始试观察训练时被丢弃token的比例。如果丢弃太多说明Router的选择太集中或容量太小可以适当调大capacity_factor如果几乎为零则逐步减小压出更多计算密度。然后才去调整aux loss系数从0.01起步每几百步看一眼各专家的token占比分布。占比的方差如果持续缩小到接近均匀说明权重偏大可以降如果某个专家占比持续超过2倍平均则要加大权重。一个反直觉的经验是对于中小规模的模型参数量在亿级以下MoE的优势并不明显。Router自身也需要数据爬坡假设总数据量没法让每个专家都充分训练那稀疏结构的收益就会打折扣。所以不要只盯着MoE大模型专利这个说法也要理解它依赖的规模前提。5.3 什么时候不应该用 MoE以及几个值得关注的扩展方向MoE不是银弹。如果你的使用场景是单机单卡、推理延迟极敏感、或者模型规模本来就很小那么MoE带来的通信开销和显存负担通常不值得。尤其是边缘端或者移动端场景参数量常驻和动态路由的不确定性会让部署变得非常不友好。宁可把精力花在蒸馏、量化或者知识蒸馏这类压缩手段上。反过来MoE真正发光的地方是训练语料极其丰富而单卡显存成为瓶颈、需要在不显著增加推理计算量的前提下扩展容量、以及多种异构任务希望共用一套底座但各任务又有自己的特性。多模态方向也可以大胆尝试——不同模态的信息处理模式差异大正好适合交由不同专家专人负责门控网络学到的模态路由有时会带来意想不到的可解释性。我在一次改造中把MoE从语言模型搬到了推荐系统的打分模型里不同用户群体走不同专家冷启动阶段用户走Generalist专家活跃用户则能访问专属专家后的组合结果效果提升比纯粹加深网络更明显。这说明MoE的分诊思想完全值得跨场景复用关键还是找到那些不同样本需要不同能力的问题。最后分享两个小技巧。如果你已经用上了MoE建议把Router的决策分布做成可视化日志每训练几千步看一眼比任何损失曲线都能更快发现问题另一个是在测试阶段记录专家热力分布如果模型收敛后大部分token的Top-1专家高度集中往往说明模型只学到了表面规律需要回头检查训练数据本身的多样性。MoE这条路技术在往前走但工程的坑说到底还是那几类通信、负载、稳定性。理解了底层逻辑剩下的就是耐心调参和踩坑。希望这些经验能帮你少走几步弯路。
阅读完成 · 觉得有帮助？
