做 NLP 训练这几年，MoE（混合专家）已经从一个圈子里的工程黑话，变成了大模型设计的标配选项。刚接触这个概念时，我也觉得“专家”“路由”这些词有点玄乎，但真正把一个带稀疏激活的前馈层写进 Transformer 之后，才理解它解决的问题其实特别朴素：想让模型更强，就得扩大参数量；参数量上去了，又不希望每个 token 都跑一遍全部参数，于是只能让 token 只激活一小部分参数路径，剩余专家保持静默。这篇文章就把“稀疏激活”这条路径从数学推导到代码实现都拆开讲一遍，适合正在训练大模型、跑微调，或者准备把 MoE 层搬进自己项目的同学直接参考。1. 为什么会出现 MoE：从密集到稀疏的前馈网络1.1 密集模型的成本瓶颈传统 Transformer 里最吃算力的部分，往往不是注意力，而是前馈网络（FFN）。每个 token 经过多头注意力之后，都会被同一个 FFN 完整过一遍：升维、激活、降维，参数多、计算量固定，token 数量再乘上去，FLOPs 一下子就爆炸了。我们平时用 FLOPs 衡量模型成本，但真正上线部署时还要考虑显存和带宽。FFN 的参数量往往占到整个模型的一半以上，ViT、BERT、GPT 这些结构都躲不开。于是遇到一个很现实的问题：想要模型“认知能力”更强，通常要加大 FFN 中间层维度；维度一大，FLOPs 和显存占用跟着涨。密集模型在“能力”和“开销”之间几乎是绑定的，想提升一方面，就必须付出两方面的代价。我最早踩到这个问题是在训练一个小规模多语言分类器的时候。模型参数从两千万涨到一亿，分类效果提升不错，但显存和单步训练时间同步翻了四五倍，这种性价比对很多实际业务是接受不了的。后来开始接触 MoE，才意识到 MoE 的设计初衷，就是要在这个绑定关系里撕开一道口子。1.2 稀疏激活的核心思想MoE 的思路是：不把 FFN 做成一个巨大的稠密网络，而是把 FFN 拆成多个结构相似但参数独立的“专家”子网络。前面跟着一个路由模块，每个 token 进来先计算它和所有专家的匹配度，选出最合适的 k 个专家，然后只让这 k 个专家处理 token，最终结果按权重加权输出。关键点在于“稀疏激活”。假设总共有 16 个专家，每个 token 只激活其中 2 个，那么理论上每个 token 的计算量大约只有满激活的 1/8，但参数量依然是 16 个专家的总和。模型容量变大，单 token 计算量没有同步变大，这正是 MoE 最核心的卖点。需要注意，“稀疏激活”是指激活路径稀疏，不代表权重矩阵稀疏。MoE 层里每个专家都是一组完整权重，路由模块也是一组权重矩阵，整体还是稠密存储。真正收益是单次前向只走其中一小部分，所以 FLOPs 受限，显存驻留空间却没有成比例下降。做工程时不要把这两者混淆，否则后面设计显存策略会带偏。一句话总结：MoE 是用“参数总量换能力”，用“稀疏计算换效率”。这条路并不难理解，真正的难点在于如何让每个专家都被均匀利用，以及如何在训练中保持路由稳定，这就要把数学定义先捋清楚。2. 稀疏激活的数学表达2.1 门控函数：从 softmax 到 top-k 约定MoE 的数学起点是假设有 $E$ 个专家，输入 token 特征为 $x \in \mathbb{R}^d$，路由模块学习一个映射：$$ p_i(x) = \frac{\exp\left(x \cdot W_{g,i}\right)}{\sum_{j=1}^{E} \exp\left(x \cdot W_{g,j}\right)} $$这里的 $W_g \in \mathbb{R}^{d \times E}$ 就是路由矩阵。$p_i$ 表示 $x$ 被分配到第 $i$ 个专家的概率。严格来说，最终真的参与计算的并不是全部 $E$ 个专家，而是根据概率值排序后取出的 top-k 个。用数学表达这个“只选少数专家”的过程：$$ y = \sum_{i \in \mathcal{K}(x)} g_i(x) \cdot E_i(x) $$其中 $\mathcal{K}(x)$ 是通过 top-k 筛选出的专家集合，$g_i(x)$ 是归一化后的路由权重，$E_i(x)$ 是第 $i$ 个专家对输入 $x$ 的输出。这个公式看起来很简单，但绝大部分实现细节都藏在 $g_i(x)$ 的归一化方式上。通常做法是先算所有专家的 softmax 概率，再取 top-k 对应的概率值，最后对这 k 个概率再做一次归一化，确保它们和为 1。这个“二次归一化”很重要。如果不做，前面分数稍微大一点，输出幅度就会跳变；做了之后，无论怎么选，最终输出都有稳定的量纲。我早期实现时偷懒直接拿 softmax 概率当权重，结果模型输出方差一路变大，损失曲线怎么调都抖。后来查实现，发现大家几乎都做这一步，才把它补上。另一个容易忽略的约定是，top-k 选择的对象是 logits 还是 softmax 概率。这两者在数学上顺序一致，因为 softmax 是单调递增函数，不会改变排序结果。但实际计算时，直接对 logits 做 top-k 能避免概率极端小时出现下溢，数值上更稳。我在代码里会先用 logits 做选择，再用 softmax 概率取权重，这样可读性和数值稳定性兼顾。2.2 负载均衡与辅助损失稀疏激活有一个天然的毛病：如果路由学歪了，所有 token 都涌向少数几个专家，其他专家几乎不参与计算。不仅算力浪费，专家之间的梯度分布也会极不均衡，整个 MoE 层会退化成“大 FFN 加一堆装饰”。所以要加辅助负载均衡损失。最常用的有两种：一种基于 token 分配比例，一种基于路由概率。实践中经常把两者结合起来。假设第 $i$ 个专家被 top-k 选中的比例是 $f_i$，路由模块给出第 $i$ 个专家的平均概率是 $P_i$：$$ L_{balance} = E \cdot \sum_{i=1}^{E} f_i \cdot P_i $$这个损失在理想情况下，当每个专家被选中的频率接近 $1/E$，且路由概率对每个专家也是均匀 $1/E$，那么 $f_i \cdot P_i$ 就会收敛到 $1/E^2$，求和后得到 $1$，这是一个可以预测量级的基准。实际训练时乘上一个很小的系数，比如 0.01，作为辅助项加到总 loss 里。注意 $f_i$ 是用“离散选择”得到的比例，本身不可导；$P_i$ 是连续的、可导的。两者相乘之后，梯度能通过 $P_i$ 回流到路由参数上，所以这个损失是可训练的。我最初看论文时对这一点困惑了很久，总觉得离散选择没法反向传播，其实关键是只让连续项承担梯度。这一点后面代码里会直观看到。辅助损失的系数不能拍脑袋乱定。太大会让路由过度均匀化，就是明明某些 expert 就是擅长处理某些模式，却被迫各占 1/16 的流量；太小又拉不回失衡。常见习惯是从 0.01 开始调，观察训练早期 500 步里的专家分配直方图，再决定往上加还是往下减。2.3 专家容量与 token 丢弃除了负载均衡损失，工程上还有一个更硬的约束叫专家容量（capacity）。MoE 层在某个 batch 里，总 token 数为 $T$，每个 token 激活 $k$ 个专家，那么平均每个专家理论上会收到 $k \cdot T / E$ 个 token。为了避免单个专家计算量波动太大，我们设定容量：$$ C = \left\lceil c \cdot \frac{k \cdot T}{E} \right\rceil $$其中 $c$ 是容量系数。$c$ 大于 1，意味着允许专家最多处理理论上平均数的一点几倍；$c$ 等于 1，则严格按平均容量处理。如果一个专家收到的 token 超过容量，超出的 token 会被直接丢弃，不参与该 expert 计算。容量设置会直接影响训练质量。$c$ 太大，负载均衡失效，快的专家要等慢的专家，计算资源空转；$c$ 太小，token 被丢弃的概率增加，信息丢失严重，loss 会变得不稳定。实际训练中，$c$ 常用 1.0 到 2.0 之间，但具体要看 batch 大小和 token 分布。batch 小、每个 batch 的 topic 差异大时，分布波动也大，$c$ 就要放宽一些。有一个容易被忽略的点：被丢弃的 token 并不是完全不产生任何影响。它们依然会通过路由的连续路径产生梯度，只是专家计算路径的梯度为零。这就是为什么丢弃完成后，还要在代码里用 mask 把对应位置的“输出贡献”清零，否则丢弃 token 会带着部分概率梯度误导路由。3. 从零实现一个 MoE 层3.1 搭建路由器和专家模块理论说完，直接上代码。我用 PyTorch 写一个最小可运行的 MoE 层，核心组件只有两个：路由器和若干专家。路由器用一层线性层把输入映射到 $E$ 个 score，专家用两层 MLP 表示。为了教学清晰，代码里暂时不做特别复杂的优化，先把逻辑讲通。import torch import torch.nn as nn import torch.nn.functional as F class Expert(nn.Module): def __init__(self, d_model, hidden_dim): super().__init__() self.net = nn.Sequential( nn.Linear(d_model, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, d_model) ) def forward(self, x): return self.net(x)专家内部的维度可以比 d_model 大很多，这就是模型容量的主要来源。实际工程里，专家之间通常结构相同，但参数独立；如果为了省显存，也可以共享中间激活，但初始实现不需要考虑这个。3.2 前向传播中的排列与聚合路由器的核心任务是产生 logits 和选取 top-k 索引。我在这一步顺便把辅助损失的必要信息算好。下面给出路由模块：class TopKRouter(nn.Module): def __init__(self, d_model, num_experts, top_k): super().__init__() self.gate = nn.Linear(d_model, num_experts, bias=False) self.num_experts = num_experts self.top_k = top_k def forward(self, x): logits = self.gate(x) probs = F.softmax(logits, dim=-1) top_probs, top_indices = torch.topk(probs, self.top_k, dim=-1) top_probs = top_probs / (top_probs.sum(dim=-1, keepdim=True) + 1e-6) return top_probs, top_indices, probstopk返回的索引在多卡和多 batch 场景下是随机的吗？不完全是，但为了稳定复现，最好对排序操作加上stable=True。我在实现中处理 token 到 expert 分组时，会把top_indices展平，然后做一次稳定排序，让每个专家拿到的 token 顺序可预测。接下来是 MoE 层的主干。这里有一个很容易绕晕的地方：每个 token 被分配给多个专家，要先复制 token 到多个槽位，再按专家编号排序，再把排好序的 token 分批送入专家网络，最后把结果通过散射加回每个 token 的对应位置。看代码可能更直接：class MoELayer(nn.Module): def __init__(self, d_model, num_experts, top_k, hidden_dim, capacity_factor=1.25): super().__init__() self.num_experts = num_experts self.top_k = top_k self.router = TopKRouter(d_model, num_experts, top_k) self.experts = nn.ModuleList([ Expert(d_model, hidden_dim) for _ in range(num_experts) ]) self.capacity_factor = capacity_factor def forward(self, x): original_shape = x.shape if x.dim() == 3: bsz, seq_len, d_model = x.shape x = x.reshape(-1, d_model) T = x.shape[0] top_w, top_idx, probs = self.router(x) # 每个 token 复制 top_k 次，形成扁平化 token 槽位 flat_idx = top_idx.reshape(-1) flat_w = top_w.reshape(-1) token_idx = torch.arange(T, device=x.device) \ .repeat_interleave(self.top_k) # 按专家编号排序，方便批量送入 expert sorted_idx, order = torch.sort(flat_idx, stable=True) token_idx_sorted = token_idx[order] flat_w_sorted = flat_w[order] tokens_expanded = x.repeat_interleave(self.top_k, dim=0)[order] # 根据专家计数切块 counts = torch.bincount(sorted_idx, minlength=self.num_experts) boundaries = counts.cumsum(0) capacity = int(self.capacity_factor * (T * self.top_k) / self.num_experts) + 1 expert_out = torch.zeros_like(tokens_expanded) for e in range(self.num_experts): start = 0 if e == 0 else boundaries[e - 1].item() end = boundaries[e].item() if end - start > capacity: end = start + capacity if start >= end: continue expert_out[start:end] = self.experts[e](tokens_expanded[start:end]) # 将各专家输出乘以对应权重并散射回原 token out = torch.zeros_like(x) contribution = expert_out * flat_w_sorted.unsqueeze(-1) idx_2d = token_idx_sorted.view(-1, 1).expand(-1, x.size(-1)) out.scatter_add_(0, idx_2d, contribution) return out.view(original_shape)这里有几个细节值得解释。tokens_expanded看起来把每个 token 复制了 $k$ 次，但复制之后又经过order排序，所以 token 的排列顺序已经打乱，等维度上每个槽位对应一个“token + 被选中的专家”。每个专家处理它拿到的连续一段 token，再通过scatter_add_把多个 slot 的结果累加回原始 token。有同学可能会问，为什么不用index_add或直接矩阵乘法批量并行。这里的循环是按专家粒度循环，如果专家数量不大，开销可以接受；如果专家数达到几十上百，就需要在 CUDA 层面做 fused kernel，或者在专家内部用更粗粒度的 batch 计算。但作为学习示例，逐专家循环是最容易理解和调试的版本。3.3 引入路由噪声与定制反向传播数学上，如果完全不做随机扰动，路由很可能过早收敛到一个次优点：某个专家刚开始表现好一点，后续所有 token 都倾向它，它继续获得更多梯度，其他专家越来越难被选中。为了解决这个问题，通常会在路由 logits 上增加噪声：if self.training and self.noise_std > 0.0: logits = logits + torch.randn_like(logits) * self.noise_std这种噪声只作用于训练，推理时关闭。噪声强度不能太大，否则路由信号会被淹没；也不能太小，否则起不到探索作用。经验上常见的取值范围是 0.01 到 0.1，具体要对路由 logits 的量级做归一化之后再看。真正在工程实现里，更精致的做法是给噪声乘上标准差，而不是加一个固定尺度。也就是说，噪声强度应该是基于当前 logits 的标准差，这样不管训练到哪个阶段，探索幅度都相对稳定。我这里先用固定尺度，简单直观，也已经能解决大部分早期收敛问题。反向传播方面，topk这个操作本身是不可导的，但我们在torch.topk之后取的top_probs是连续的，它带着最终的权重梯度回流到 softmax 上，所以路由参数能正常训练。这一点刚才在数学部分强调过，写代码时会更加明显：只要你把可导的权重赋值给flat_w_sorted，反传就自动成立，不需要手动 hack。4. 训练 MoE 的稳定性问题4.1 路由崩溃：为什么一切都被路由到一个专家训练中最常见的失败模式是路由崩溃。表现为某个专家承接了超过半数的 token，其他专家要么闲置，要么只处理零星 token。从 loss 曲线看，前期正常下降，某一步之后突然出现剧烈的抖动，或者收敛速度明显变慢。路由崩溃的原因，可以从梯度信号理解。token 被分到专家 A 后，专家 A 的梯度立刻更新，它处理该 token 的 loss 可能会下降一点。由于最终输出又依赖路由权重，路由模块也会收到反馈，让它更倾向于选专家 A。这个正反馈循环在训练初期特别危险，因为初始权重是随机的，哪怕只是偶然优势，也可能被放大。我实际调试时见过一个挺讽刺的现象：路由 logits 的范数在训练中不断变大，softmax 输出变得非常尖锐，专家利用率的直方图从均匀分布慢慢变成一根柱子。碰到这种情况，首先要确认辅助负载均衡损失是否真的在更新，其次要重新初始化路由矩阵。把路由矩阵的初始化方差调低一点，能显著缓解早期“赢家通吃”的状态。4.2 负载均衡损失尺度怎么调辅助损失的系数 $alpha$，是我在 MoE 训练里花费最多时间调的超参数。设置太小，专家分桶可能倾斜；设置太大，路由会变得过分均匀。有个典型的诊断指标：看 load balance loss 的量级和主任务 loss 的量级差异。如果主任务 loss 在 2 附近，辅助 loss 在 0.5 附近，乘以 0.01 之后才 0.005，对总 loss 影响不大。如果主任务 loss 已经降到 0.1，而辅助 loss 还在 0.5 附近，即使乘以 0.01 也会产生明显干扰，需要进一步降低系数。另一个更直接的监控方式是记录每个专家获得的 token 数量，计算它们的标准差。标准差过高，调大系数；标准差过低，调小系数。初期可以固定 $alpha=0.01$ 跑 2000 步，画出专家 bean count 直方图。如果能看到大多数专家都有非零流量，就不一定要调；如果柱子形状严重偏斜，再把系数提高到 0.02 或 0.05。反之，如果明显出现“过于均匀”，比如所有专家都要挤到同一个使用比例，说明 $alpha$ 太强，降到 0.005 再观察。4.3 专家 dropout 与初始化MoE 的专家参数通常不会复用，所以每个专家都要独立初始化。最怕的是所有专家用同一套随机种子，导致它们一开始完全相同，路由退化成一个随机数游戏。正确做法是让每个专家模块使用不同的随机状态，PyTorch 里每个nn.Module默认就是不同的随机初始化，只要别手动共享参数就行。专家层里的 Dropout 位置也很有讲究。如果放在专家输出之后，再经过路由权重加权，等于对最终 token 高度施加随机扰动，这个扰动通常合理。如果放在专家内部的激活处，可能影响专家自身的学习信号，但并不均匀作用于所有 token。经验上，普通 FFN 怎么设置，专家内部就怎么设置，不需要搞特殊。梯度裁剪对 MoE 也很重要。由于每个 token 只会激活一部分专家，某些 batch 里某个专家可能被极端赋值，梯度范数突然放大。不做裁剪，一次大更新就可能把路由权重推入饱和区，再想拉回来就得很长时间。我通常把全局梯度范数裁剪到 1.0，训练稳定性会明显好一截。5. 部署推理里的稀疏激活收益5.1 收获的算力和付出的显存很多人第一次算 MoE 的 FLOPs 时会很乐观，因为激活的专家数少，计算量低。但真正部署时发显存不够，或者 batch 吞吐表现不如预期，原因在于参数驻留和计算是两回事。模型容量来自所有专家的参数总和，所以每个 GPU 的显存最好能容纳全部专家权重。如果放不下，就必须把专家分片放到多台卡上。此时还会引入一个很本质的代价：通信。不同于普通 Transformer 层里张量按 batch 并行切分，MoE 的 expert 是“按功能”切分的，某个 token 可能被送到另一张卡上的专家，这个跨卡交换叫 all-to-all 通信。如果路由分布很均匀，all-to-all 通信量是稳定的；一旦某个卡上的专家被大量命中，通信热点就会出现，形成等待，吞吐反而下降。稀疏激活在算力上的收益，往往会被慢速通信吃回去一部分。所以推理部署时，不只是看 FLOPs，还要看通信带宽和 batch 形状。5.2 跨设备放置与通信开销工程上通常把 expert 分布在多卡上，每个卡上放部分专家，同时保留完整的自注意力和 embedding。每层 MoE forward 就要做一次 token 和专家之间的排列组合。token 从本卡发到目标卡，专家算完，再拿回来，过程中有两轮网络通信。减少通信的办法之一是“专家本地化”。如果某些 token 的总数不多，可以把一批 token 分到同一批卡的专家集合，降低跨卡比例。更激进的方案是让每张卡上的专家数量变多，每张卡处理它自己的 token，只在不得已时跨卡。这本质上是容量和负载均衡之间的又一层 trade-off。从实践中看，小 batch 推理时，稀疏激活的加速比例常常不如大 batch。小 batch 下 token 数量少，专家分布随机性强，通信次数没减少，但每次通信携带的数据量下降，带宽利用率很差。所以 MoE 更适合高并发、大 batch 的场景；在线小请求，一般要加 batch 聚合或者把多个请求组合起来一起路由。5.3 面向推理的瘦身：蒸馏与层融合如果发现推理吞吐达不到预期，一个思路是先把 MoE 转回密集模型。具体做法是把多个专家的输出按照路由频率加权合并，变成一个普通 FFN，这就是“MoE 蒸馏”。蒸馏之后模型结构和普通 Transformer 完全一样，可以直接享受成熟的推理优化，但参数量依然很大，模型文件大小不会变小。另一个思路是只保留高频专家，把低频专家直接删除。有些稀疏模型经过训练后，大量专家使用率极低。做一次统计：训练集上每个专家被选中的次数，低于阈值的专家可以被砍掉，然后微调一小段时间恢复性能。这个方法在线上模型瘦身时很实用，但要注意：专家少不等于质量线性下降，具体要看路由分布和任务多样性。还有一个工程技巧是融合路由和专家输入。如果路由是线性层，那么“算 logits”这一步可以并入上一层 attention 的线性输出中，减少一次独立的矩阵乘。类似的可融合点还包括 softmax 与 topk 的 kernel 优化。这些细节在模型很大时能带来百分之十几的延迟收益，值得做。6. 常见问题与排查实录6.1 训练中 loss 出现 NaN 怎么办NaN 的常见来源，一个是 softmax 的 logits 太大，另一个是topk之后路由权重出现极小值，除法爆炸。先检查路由矩阵的梯度范数，如果路由梯度明显比其他层大，就加梯度裁剪或降低学习率。另一个隐蔽原因是 fp16 混合精度下，repeat_interleave和scatter_add_这类操作的数值范围容易溢出。MoE 里的 token 复制会放大显存占用，同时放大激活值的尺度。如果模型在 MoE 层附近出现异常，可以先把该项切到 fp32 做基准对比，再决定哪里保持 fp16。从我的排查经验看，直接把辅助损失的输出乘上系数后参与 loss 计算，一般不会造成 NaN。最容易造成 NaN 的是手工实现里对丢弃 token 位置填了一个inf或者nan的 markers，看似无害，实际经过scatter_add_之后会污染所有 token。这个坑我踩过，最好用浮点零而不是-inf来标记无效位置。6.2 容量设置与 token 丢弃比例观测容量设置不当时，训练日志里最明显的是 token 丢弃率上升。MoE 层可以在 forward 里记录一个计数，统计每个专家被超容量的 token 数，输出到日志里。如果丢弃率超过百分之五，就该提高容量系数。有一个我觉得很实用的脚本：训练前 500 步正常跑，每 50 步输出一次“专家使用率直方图”。不要等训练完了再看统计，那样已经来不及调整。在训练早期，专家使用率就应该比较接近均匀分布；如果前 500 步就已经严重倾斜，说明负载均衡损失没起作用，先检查损失计算是否漏了 mask。容量过大虽然会降低丢弃率，但也会模糊稀疏性的优点。因为容量乘上去，实际计算量没有严格减少，甚至出现大量 padding 和空转。所以容量不是越高越好，而是要让容量和负载均衡损失协同工作，一个管硬上限，一个做软引导。6.3 评估指标异常排查我遇到过一种情况：训练 loss 下降正常，但验证集上的指标波动很大。排查后发现是路由在评估阶段对 token 顺序敏感，不同 batch 下专家的选取结果不一致。这个现象在 batch 很小的时候尤其明显，因为 top-k 排序结果对随机 seed 和 batch 组成非常敏感。解决方案之一是评估时固定 batch 顺序和随机种子，让每次评估的输入排序完全一致，至少保证指标可复现。另一个方案是在验证集上多做几次推理取平均，降低路由随机性带来的波动。若验证内容是 token-level 的任务，还可以把每个 token 的 top-k 索引提交 log 出来，检查是否合理。还有一个容易忽略的点：评估阶段如果关闭了路由噪声，但训练阶段使用了噪声，模型的分布会发生轻微偏移。有些论文会把推理噪声保留一个极小的值，以平滑评估指标。但如果你用固定 seed 复现，还是会发现：同一份 checkpoint 在不同 batch size 下，路由结果可能变化，模型输出也因此变化。这不是 bug，而是 MoE 的固有特性，设计实验时需要把 batch 统一。收尾我个人在实际项目中体会最深的一件事是：MoE 的数学并不难，难的是它把“容量、算力、内存、通信”这四个变量同时拉进来，让每个决策都变成多目标权衡。刚开始调 MoE 层那段时间，我几乎每天都在改容量系数、改专家数量、改负载均衡损失权重，模型本身的结构反而一直没动。后来发现，真正让训练稳定下来的往往是三个不起眼的细节：路由 logits 的初始化方差、噪声强度、辅助损失的系数。把这三个调顺了，很多玄学问题会自己消失。如果让我给刚上手 MoE 的人一个最直白的建议，我会说：先把单层 MoE 用一个小模型跑通，加好资源开销统计，再谈扩展到大模型。不要一上来就追求几十个专家、多卡并行，先让路由“学得动、均衡得了、丢得起”，每一步都有日志可以回查，剩下的事情反而水到渠成。祝大家都能在自己的项目里，把稀疏激活这条路走顺。
做 NLP 训练这几年，MoE（混合专家）已经从一个圈子里的工程黑话，变成了大模型设计的标配选项。刚接触这个概念时，我也觉得“专家”“路由”这些词有点玄乎，但真正把一个带稀疏激活的前馈层写进 Transformer 之后，才理解它解决的问题其实特别朴素：想让模型更强，就得扩大参数量；参数量上去了，又不希望每个 token 都跑一遍全部参数，于是只能让 token 只激活一小部分参数路径，剩余专家保持静默。这篇文章就把“稀疏激活”这条路径从数学推导到代码实现都拆开讲一遍，适合正在训练大模型、跑微调，或者准备把 MoE 层搬进自己项目的同学直接参考。1. 为什么会出现 MoE：从密集到稀疏的前馈网络1.1 密集模型的成本瓶颈传统 Transformer 里最吃算力的部分，往往不是注意力，而是前馈网络（FFN）。每个 token 经过多头注意力之后，都会被同一个 FFN 完整过一遍：升维、激活、降维，参数多、计算量固定，token 数量再乘上去，FLOPs 一下子就爆炸了。我们平时用 FLOPs 衡量模型成本，但真正上线部署时还要考虑显存和带宽。FFN 的参数量往往占到整个模型的一半以上，ViT、BERT、GPT 这些结构都躲不开。于是遇到一个很现实的问题：想要模型“认知能力”更强，通常要加大 FFN 中间层维度；维度一大，FLOPs 和显存占用跟着涨。密集模型在“能力”和“开销”之间几乎是绑定的，想提升一方面，就必须付出两方面的代价。我最早踩到这个问题是在训练一个小规模多语言分类器的时候。模型参数从两千万涨到一亿，分类效果提升不错，但显存和单步训练时间同步翻了四五倍，这种性价比对很多实际业务是接受不了的。后来开始接触 MoE，才意识到 MoE 的设计初衷，就是要在这个绑定关系里撕开一道口子。1.2 稀疏激活的核心思想MoE 的思路是：不把 FFN 做成一个巨大的稠密网络，而是把 FFN 拆成多个结构相似但参数独立的“专家”子网络。前面跟着一个路由模块，每个 token 进来先计算它和所有专家的匹配度，选出最合适的 k 个专家，然后只让这 k 个专家处理 token，最终结果按权重加权输出。关键点在于“稀疏激活”。假设总共有 16 个专家，每个 token 只激活其中 2 个，那么理论上每个 token 的计算量大约只有满激活的 1/8，但参数量依然是 16 个专家的总和。模型容量变大，单 token 计算量没有同步变大，这正是 MoE 最核心的卖点。需要注意，“稀疏激活”是指激活路径稀疏，不代表权重矩阵稀疏。MoE 层里每个专家都是一组完整权重，路由模块也是一组权重矩阵，整体还是稠密存储。真正收益是单次前向只走其中一小部分，所以 FLOPs 受限，显存驻留空间却没有成比例下降。做工程时不要把这两者混淆，否则后面设计显存策略会带偏。一句话总结：MoE 是用“参数总量换能力”，用“稀疏计算换效率”。这条路并不难理解，真正的难点在于如何让每个专家都被均匀利用，以及如何在训练中保持路由稳定，这就要把数学定义先捋清楚。2. 稀疏激活的数学表达2.1 门控函数：从 softmax 到 top-k 约定MoE 的数学起点是假设有 $E$ 个专家，输入 token 特征为 $x \in \mathbb{R}^d$，路由模块学习一个映射：$$ p_i(x) = \frac{\exp\left(x \cdot W_{g,i}\right)}{\sum_{j=1}^{E} \exp\left(x \cdot W_{g,j}\right)} $$这里的 $W_g \in \mathbb{R}^{d \times E}$ 就是路由矩阵。$p_i$ 表示 $x$ 被分配到第 $i$ 个专家的概率。严格来说，最终真的参与计算的并不是全部 $E$ 个专家，而是根据概率值排序后取出的 top-k 个。用数学表达这个“只选少数专家”的过程：$$ y = \sum_{i \in \mathcal{K}(x)} g_i(x) \cdot E_i(x) $$其中 $\mathcal{K}(x)$ 是通过 top-k 筛选出的专家集合，$g_i(x)$ 是归一化后的路由权重，$E_i(x)$ 是第 $i$ 个专家对输入 $x$ 的输出。这个公式看起来很简单，但绝大部分实现细节都藏在 $g_i(x)$ 的归一化方式上。通常做法是先算所有专家的 softmax 概率，再取 top-k 对应的概率值，最后对这 k 个概率再做一次归一化，确保它们和为 1。这个“二次归一化”很重要。如果不做，前面分数稍微大一点，输出幅度就会跳变；做了之后，无论怎么选，最终输出都有稳定的量纲。我早期实现时偷懒直接拿 softmax 概率当权重，结果模型输出方差一路变大，损失曲线怎么调都抖。后来查实现，发现大家几乎都做这一步，才把它补上。另一个容易忽略的约定是，top-k 选择的对象是 logits 还是 softmax 概率。这两者在数学上顺序一致，因为 softmax 是单调递增函数，不会改变排序结果。但实际计算时，直接对 logits 做 top-k 能避免概率极端小时出现下溢，数值上更稳。我在代码里会先用 logits 做选择，再用 softmax 概率取权重，这样可读性和数值稳定性兼顾。2.2 负载均衡与辅助损失稀疏激活有一个天然的毛病：如果路由学歪了，所有 token 都涌向少数几个专家，其他专家几乎不参与计算。不仅算力浪费，专家之间的梯度分布也会极不均衡，整个 MoE 层会退化成“大 FFN 加一堆装饰”。所以要加辅助负载均衡损失。最常用的有两种：一种基于 token 分配比例，一种基于路由概率。实践中经常把两者结合起来。假设第 $i$ 个专家被 top-k 选中的比例是 $f_i$，路由模块给出第 $i$ 个专家的平均概率是 $P_i$：$$ L_{balance} = E \cdot \sum_{i=1}^{E} f_i \cdot P_i $$这个损失在理想情况下，当每个专家被选中的频率接近 $1/E$，且路由概率对每个专家也是均匀 $1/E$，那么 $f_i \cdot P_i$ 就会收敛到 $1/E^2$，求和后得到 $1$，这是一个可以预测量级的基准。实际训练时乘上一个很小的系数，比如 0.01，作为辅助项加到总 loss 里。注意 $f_i$ 是用“离散选择”得到的比例，本身不可导；$P_i$ 是连续的、可导的。两者相乘之后，梯度能通过 $P_i$ 回流到路由参数上，所以这个损失是可训练的。我最初看论文时对这一点困惑了很久，总觉得离散选择没法反向传播，其实关键是只让连续项承担梯度。这一点后面代码里会直观看到。辅助损失的系数不能拍脑袋乱定。太大会让路由过度均匀化，就是明明某些 expert 就是擅长处理某些模式，却被迫各占 1/16 的流量；太小又拉不回失衡。常见习惯是从 0.01 开始调，观察训练早期 500 步里的专家分配直方图，再决定往上加还是往下减。2.3 专家容量与 token 丢弃除了负载均衡损失，工程上还有一个更硬的约束叫专家容量（capacity）。MoE 层在某个 batch 里，总 token 数为 $T$，每个 token 激活 $k$ 个专家，那么平均每个专家理论上会收到 $k \cdot T / E$ 个 token。为了避免单个专家计算量波动太大，我们设定容量：$$ C = \left\lceil c \cdot \frac{k \cdot T}{E} \right\rceil $$其中 $c$ 是容量系数。$c$ 大于 1，意味着允许专家最多处理理论上平均数的一点几倍；$c$ 等于 1，则严格按平均容量处理。如果一个专家收到的 token 超过容量，超出的 token 会被直接丢弃，不参与该 expert 计算。容量设置会直接影响训练质量。$c$ 太大，负载均衡失效，快的专家要等慢的专家，计算资源空转；$c$ 太小，token 被丢弃的概率增加，信息丢失严重，loss 会变得不稳定。实际训练中，$c$ 常用 1.0 到 2.0 之间，但具体要看 batch 大小和 token 分布。batch 小、每个 batch 的 topic 差异大时，分布波动也大，$c$ 就要放宽一些。有一个容易被忽略的点：被丢弃的 token 并不是完全不产生任何影响。它们依然会通过路由的连续路径产生梯度，只是专家计算路径的梯度为零。这就是为什么丢弃完成后，还要在代码里用 mask 把对应位置的“输出贡献”清零，否则丢弃 token 会带着部分概率梯度误导路由。3. 从零实现一个 MoE 层3.1 搭建路由器和专家模块理论说完，直接上代码。我用 PyTorch 写一个最小可运行的 MoE 层，核心组件只有两个：路由器和若干专家。路由器用一层线性层把输入映射到 $E$ 个 score，专家用两层 MLP 表示。为了教学清晰，代码里暂时不做特别复杂的优化，先把逻辑讲通。import torch import torch.nn as nn import torch.nn.functional as F class Expert(nn.Module): def __init__(self, d_model, hidden_dim): super().__init__() self.net = nn.Sequential( nn.Linear(d_model, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, d_model) ) def forward(self, x): return self.net(x)专家内部的维度可以比 d_model 大很多，这就是模型容量的主要来源。实际工程里，专家之间通常结构相同，但参数独立；如果为了省显存，也可以共享中间激活，但初始实现不需要考虑这个。3.2 前向传播中的排列与聚合路由器的核心任务是产生 logits 和选取 top-k 索引。我在这一步顺便把辅助损失的必要信息算好。下面给出路由模块：class TopKRouter(nn.Module): def __init__(self, d_model, num_experts, top_k): super().__init__() self.gate = nn.Linear(d_model, num_experts, bias=False) self.num_experts = num_experts self.top_k = top_k def forward(self, x): logits = self.gate(x) probs = F.softmax(logits, dim=-1) top_probs, top_indices = torch.topk(probs, self.top_k, dim=-1) top_probs = top_probs / (top_probs.sum(dim=-1, keepdim=True) + 1e-6) return top_probs, top_indices, probstopk返回的索引在多卡和多 batch 场景下是随机的吗？不完全是，但为了稳定复现，最好对排序操作加上stable=True。我在实现中处理 token 到 expert 分组时，会把top_indices展平，然后做一次稳定排序，让每个专家拿到的 token 顺序可预测。接下来是 MoE 层的主干。这里有一个很容易绕晕的地方：每个 token 被分配给多个专家，要先复制 token 到多个槽位，再按专家编号排序，再把排好序的 token 分批送入专家网络，最后把结果通过散射加回每个 token 的对应位置。看代码可能更直接：class MoELayer(nn.Module): def __init__(self, d_model, num_experts, top_k, hidden_dim, capacity_factor=1.25): super().__init__() self.num_experts = num_experts self.top_k = top_k self.router = TopKRouter(d_model, num_experts, top_k) self.experts = nn.ModuleList([ Expert(d_model, hidden_dim) for _ in range(num_experts) ]) self.capacity_factor = capacity_factor def forward(self, x): original_shape = x.shape if x.dim() == 3: bsz, seq_len, d_model = x.shape x = x.reshape(-1, d_model) T = x.shape[0] top_w, top_idx, probs = self.router(x) # 每个 token 复制 top_k 次，形成扁平化 token 槽位 flat_idx = top_idx.reshape(-1) flat_w = top_w.reshape(-1) token_idx = torch.arange(T, device=x.device) \ .repeat_interleave(self.top_k) # 按专家编号排序，方便批量送入 expert sorted_idx, order = torch.sort(flat_idx, stable=True) token_idx_sorted = token_idx[order] flat_w_sorted = flat_w[order] tokens_expanded = x.repeat_interleave(self.top_k, dim=0)[order] # 根据专家计数切块 counts = torch.bincount(sorted_idx, minlength=self.num_experts) boundaries = counts.cumsum(0) capacity = int(self.capacity_factor * (T * self.top_k) / self.num_experts) + 1 expert_out = torch.zeros_like(tokens_expanded) for e in range(self.num_experts): start = 0 if e == 0 else boundaries[e - 1].item() end = boundaries[e].item() if end - start > capacity: end = start + capacity if start >= end: continue expert_out[start:end] = self.experts[e](tokens_expanded[start:end]) # 将各专家输出乘以对应权重并散射回原 token out = torch.zeros_like(x) contribution = expert_out * flat_w_sorted.unsqueeze(-1) idx_2d = token_idx_sorted.view(-1, 1).expand(-1, x.size(-1)) out.scatter_add_(0, idx_2d, contribution) return out.view(original_shape)这里有几个细节值得解释。tokens_expanded看起来把每个 token 复制了 $k$ 次，但复制之后又经过order排序，所以 token 的排列顺序已经打乱，等维度上每个槽位对应一个“token + 被选中的专家”。每个专家处理它拿到的连续一段 token，再通过scatter_add_把多个 slot 的结果累加回原始 token。有同学可能会问，为什么不用index_add或直接矩阵乘法批量并行。这里的循环是按专家粒度循环，如果专家数量不大，开销可以接受；如果专家数达到几十上百，就需要在 CUDA 层面做 fused kernel，或者在专家内部用更粗粒度的 batch 计算。但作为学习示例，逐专家循环是最容易理解和调试的版本。3.3 引入路由噪声与定制反向传播数学上，如果完全不做随机扰动，路由很可能过早收敛到一个次优点：某个专家刚开始表现好一点，后续所有 token 都倾向它，它继续获得更多梯度，其他专家越来越难被选中。为了解决这个问题，通常会在路由 logits 上增加噪声：if self.training and self.noise_std > 0.0: logits = logits + torch.randn_like(logits) * self.noise_std这种噪声只作用于训练，推理时关闭。噪声强度不能太大，否则路由信号会被淹没；也不能太小，否则起不到探索作用。经验上常见的取值范围是 0.01 到 0.1，具体要对路由 logits 的量级做归一化之后再看。真正在工程实现里，更精致的做法是给噪声乘上标准差，而不是加一个固定尺度。也就是说，噪声强度应该是基于当前 logits 的标准差，这样不管训练到哪个阶段，探索幅度都相对稳定。我这里先用固定尺度，简单直观，也已经能解决大部分早期收敛问题。反向传播方面，topk这个操作本身是不可导的，但我们在torch.topk之后取的top_probs是连续的，它带着最终的权重梯度回流到 softmax 上，所以路由参数能正常训练。这一点刚才在数学部分强调过，写代码时会更加明显：只要你把可导的权重赋值给flat_w_sorted，反传就自动成立，不需要手动 hack。4. 训练 MoE 的稳定性问题4.1 路由崩溃：为什么一切都被路由到一个专家训练中最常见的失败模式是路由崩溃。表现为某个专家承接了超过半数的 token，其他专家要么闲置，要么只处理零星 token。从 loss 曲线看，前期正常下降，某一步之后突然出现剧烈的抖动，或者收敛速度明显变慢。路由崩溃的原因，可以从梯度信号理解。token 被分到专家 A 后，专家 A 的梯度立刻更新，它处理该 token 的 loss 可能会下降一点。由于最终输出又依赖路由权重，路由模块也会收到反馈，让它更倾向于选专家 A。这个正反馈循环在训练初期特别危险，因为初始权重是随机的，哪怕只是偶然优势，也可能被放大。我实际调试时见过一个挺讽刺的现象：路由 logits 的范数在训练中不断变大，softmax 输出变得非常尖锐，专家利用率的直方图从均匀分布慢慢变成一根柱子。碰到这种情况，首先要确认辅助负载均衡损失是否真的在更新，其次要重新初始化路由矩阵。把路由矩阵的初始化方差调低一点，能显著缓解早期“赢家通吃”的状态。4.2 负载均衡损失尺度怎么调辅助损失的系数 $alpha$，是我在 MoE 训练里花费最多时间调的超参数。设置太小，专家分桶可能倾斜；设置太大，路由会变得过分均匀。有个典型的诊断指标：看 load balance loss 的量级和主任务 loss 的量级差异。如果主任务 loss 在 2 附近，辅助 loss 在 0.5 附近，乘以 0.01 之后才 0.005，对总 loss 影响不大。如果主任务 loss 已经降到 0.1，而辅助 loss 还在 0.5 附近，即使乘以 0.01 也会产生明显干扰，需要进一步降低系数。另一个更直接的监控方式是记录每个专家获得的 token 数量，计算它们的标准差。标准差过高，调大系数；标准差过低，调小系数。初期可以固定 $alpha=0.01$ 跑 2000 步，画出专家 bean count 直方图。如果能看到大多数专家都有非零流量，就不一定要调；如果柱子形状严重偏斜，再把系数提高到 0.02 或 0.05。反之，如果明显出现“过于均匀”，比如所有专家都要挤到同一个使用比例，说明 $alpha$ 太强，降到 0.005 再观察。4.3 专家 dropout 与初始化MoE 的专家参数通常不会复用，所以每个专家都要独立初始化。最怕的是所有专家用同一套随机种子，导致它们一开始完全相同，路由退化成一个随机数游戏。正确做法是让每个专家模块使用不同的随机状态，PyTorch 里每个nn.Module默认就是不同的随机初始化，只要别手动共享参数就行。专家层里的 Dropout 位置也很有讲究。如果放在专家输出之后，再经过路由权重加权，等于对最终 token 高度施加随机扰动，这个扰动通常合理。如果放在专家内部的激活处，可能影响专家自身的学习信号，但并不均匀作用于所有 token。经验上，普通 FFN 怎么设置，专家内部就怎么设置，不需要搞特殊。梯度裁剪对 MoE 也很重要。由于每个 token 只会激活一部分专家，某些 batch 里某个专家可能被极端赋值，梯度范数突然放大。不做裁剪，一次大更新就可能把路由权重推入饱和区，再想拉回来就得很长时间。我通常把全局梯度范数裁剪到 1.0，训练稳定性会明显好一截。5. 部署推理里的稀疏激活收益5.1 收获的算力和付出的显存很多人第一次算 MoE 的 FLOPs 时会很乐观，因为激活的专家数少，计算量低。但真正部署时发显存不够，或者 batch 吞吐表现不如预期，原因在于参数驻留和计算是两回事。模型容量来自所有专家的参数总和，所以每个 GPU 的显存最好能容纳全部专家权重。如果放不下，就必须把专家分片放到多台卡上。此时还会引入一个很本质的代价：通信。不同于普通 Transformer 层里张量按 batch 并行切分，MoE 的 expert 是“按功能”切分的，某个 token 可能被送到另一张卡上的专家，这个跨卡交换叫 all-to-all 通信。如果路由分布很均匀，all-to-all 通信量是稳定的；一旦某个卡上的专家被大量命中，通信热点就会出现，形成等待，吞吐反而下降。稀疏激活在算力上的收益，往往会被慢速通信吃回去一部分。所以推理部署时，不只是看 FLOPs，还要看通信带宽和 batch 形状。5.2 跨设备放置与通信开销工程上通常把 expert 分布在多卡上，每个卡上放部分专家，同时保留完整的自注意力和 embedding。每层 MoE forward 就要做一次 token 和专家之间的排列组合。token 从本卡发到目标卡，专家算完，再拿回来，过程中有两轮网络通信。减少通信的办法之一是“专家本地化”。如果某些 token 的总数不多，可以把一批 token 分到同一批卡的专家集合，降低跨卡比例。更激进的方案是让每张卡上的专家数量变多，每张卡处理它自己的 token，只在不得已时跨卡。这本质上是容量和负载均衡之间的又一层 trade-off。从实践中看，小 batch 推理时，稀疏激活的加速比例常常不如大 batch。小 batch 下 token 数量少，专家分布随机性强，通信次数没减少，但每次通信携带的数据量下降，带宽利用率很差。所以 MoE 更适合高并发、大 batch 的场景；在线小请求，一般要加 batch 聚合或者把多个请求组合起来一起路由。5.3 面向推理的瘦身：蒸馏与层融合如果发现推理吞吐达不到预期，一个思路是先把 MoE 转回密集模型。具体做法是把多个专家的输出按照路由频率加权合并，变成一个普通 FFN，这就是“MoE 蒸馏”。蒸馏之后模型结构和普通 Transformer 完全一样，可以直接享受成熟的推理优化，但参数量依然很大，模型文件大小不会变小。另一个思路是只保留高频专家，把低频专家直接删除。有些稀疏模型经过训练后，大量专家使用率极低。做一次统计：训练集上每个专家被选中的次数，低于阈值的专家可以被砍掉，然后微调一小段时间恢复性能。这个方法在线上模型瘦身时很实用，但要注意：专家少不等于质量线性下降，具体要看路由分布和任务多样性。还有一个工程技巧是融合路由和专家输入。如果路由是线性层，那么“算 logits”这一步可以并入上一层 attention 的线性输出中，减少一次独立的矩阵乘。类似的可融合点还包括 softmax 与 topk 的 kernel 优化。这些细节在模型很大时能带来百分之十几的延迟收益，值得做。6. 常见问题与排查实录6.1 训练中 loss 出现 NaN 怎么办NaN 的常见来源，一个是 softmax 的 logits 太大，另一个是topk之后路由权重出现极小值，除法爆炸。先检查路由矩阵的梯度范数，如果路由梯度明显比其他层大，就加梯度裁剪或降低学习率。另一个隐蔽原因是 fp16 混合精度下，repeat_interleave和scatter_add_这类操作的数值范围容易溢出。MoE 里的 token 复制会放大显存占用，同时放大激活值的尺度。如果模型在 MoE 层附近出现异常，可以先把该项切到 fp32 做基准对比，再决定哪里保持 fp16。从我的排查经验看，直接把辅助损失的输出乘上系数后参与 loss 计算，一般不会造成 NaN。最容易造成 NaN 的是手工实现里对丢弃 token 位置填了一个inf或者nan的 markers，看似无害，实际经过scatter_add_之后会污染所有 token。这个坑我踩过，最好用浮点零而不是-inf来标记无效位置。6.2 容量设置与 token 丢弃比例观测容量设置不当时，训练日志里最明显的是 token 丢弃率上升。MoE 层可以在 forward 里记录一个计数，统计每个专家被超容量的 token 数，输出到日志里。如果丢弃率超过百分之五，就该提高容量系数。有一个我觉得很实用的脚本：训练前 500 步正常跑，每 50 步输出一次“专家使用率直方图”。不要等训练完了再看统计，那样已经来不及调整。在训练早期，专家使用率就应该比较接近均匀分布；如果前 500 步就已经严重倾斜，说明负载均衡损失没起作用，先检查损失计算是否漏了 mask。容量过大虽然会降低丢弃率，但也会模糊稀疏性的优点。因为容量乘上去，实际计算量没有严格减少，甚至出现大量 padding 和空转。所以容量不是越高越好，而是要让容量和负载均衡损失协同工作，一个管硬上限，一个做软引导。6.3 评估指标异常排查我遇到过一种情况：训练 loss 下降正常，但验证集上的指标波动很大。排查后发现是路由在评估阶段对 token 顺序敏感，不同 batch 下专家的选取结果不一致。这个现象在 batch 很小的时候尤其明显，因为 top-k 排序结果对随机 seed 和 batch 组成非常敏感。解决方案之一是评估时固定 batch 顺序和随机种子，让每次评估的输入排序完全一致，至少保证指标可复现。另一个方案是在验证集上多做几次推理取平均，降低路由随机性带来的波动。若验证内容是 token-level 的任务，还可以把每个 token 的 top-k 索引提交 log 出来，检查是否合理。还有一个容易忽略的点：评估阶段如果关闭了路由噪声，但训练阶段使用了噪声，模型的分布会发生轻微偏移。有些论文会把推理噪声保留一个极小的值，以平滑评估指标。但如果你用固定 seed 复现，还是会发现：同一份 checkpoint 在不同 batch size 下，路由结果可能变化，模型输出也因此变化。这不是 bug，而是 MoE 的固有特性，设计实验时需要把 batch 统一。收尾我个人在实际项目中体会最深的一件事是：MoE 的数学并不难，难的是它把“容量、算力、内存、通信”这四个变量同时拉进来，让每个决策都变成多目标权衡。刚开始调 MoE 层那段时间，我几乎每天都在改容量系数、改专家数量、改负载均衡损失权重，模型本身的结构反而一直没动。后来发现，真正让训练稳定下来的往往是三个不起眼的细节：路由 logits 的初始化方差、噪声强度、辅助损失的系数。把这三个调顺了，很多玄学问题会自己消失。如果让我给刚上手 MoE 的人一个最直白的建议，我会说：先把单层 MoE 用一个小模型跑通，加好资源开销统计，再谈扩展到大模型。不要一上来就追求几十个专家、多卡并行，先让路由“学得动、均衡得了、丢得起”，每一步都有日志可以回查，剩下的事情反而水到渠成。祝大家都能在自己的项目里，把稀疏激活这条路走顺。