MoE（Mixture of Experts，混合专家）最近几乎成了大模型基础设施里绕不开的关键词。它的核心思路不复杂：把Transformer中原先一个又大又笨重的FFN层，拆成一组独立的小专家网络，再让一个门控网络为每个token只挑少数几个专家来算。这样总参数量可以做得很大，但单个token的实际计算量不会跟着参数规模线性膨胀。这篇笔记从计算原理和过程出发，把Router打分、Top-k选择、dispatch/combine通信、容量因子、负载均衡损失一次讲清楚，同时落到AI Infra视角下最关心的显存、通信和并行策略上。适合正在做分布式训练、推理优化或模型结构分析的人，也适合刚接触MoE、想搞明白它内部到底在干什么的同学。

1. MoE是什么：把FFN从“单通道”改成“多专家分诊”

1.1 为什么MoE盯上的是FFN而不是Attention

Transformer block里，Attention负责token之间的信息交换，FFN负责对每个token做一次比较重的非线性变换。标准FFN通常是一个“两层线性+激活函数”的结构，参数量很大，而且不管输入是什么，每个token都必须完整过一遍相同的计算路径。

MoE的想法很直接：既然不是所有token都需要同样粒度的细节处理，那就把FFN复制成E份，让不同专家分别擅长不同类型的输入。路由模块先看每个token的特征，再决定该token去哪些专家。整个过程有点像医院的问询台：不是所有病人都需要所有科室，先分诊，再对症下药。

这里面有一个非常关键的性质：模型的总参数量等于E个专家参数之和，但每个token只会激活其中k个专家（通常k=1或2）。参数量和计算量被解耦了。想要更大容量的模型，可以增加专家数；想要控制计算成本，可以保持较小的k。AI Infra里经常说“MoE是参数高效而不是计算高效”，指的就是这个逻辑——参数还是那么多，甚至更多，但每次推理只走一小部分。

1.2 MoE层里的标准组件

一个典型的MoE层，至少包含四部分：

• Router（门控网络）：通常是一个线性层，输入token的hidden state，输出E个logits，再经过Softmax变成每个专家的概率，然后取Top-k。
• Expert列表：E个并行的专家网络，一般是标准的FFN/MLP结构。每个专家共享相同的输入输出维度，但参数互相独立。
• 辅助负载均衡损失：防止几乎所有token都涌向同一个专家。它不改变前向计算路径，但会在训练时通过梯度修正Router的分配倾向。
• 容量因子（capacity factor）：限制每个专家最多处理多少token，超出部分的token会被直接跳过或走残差，避免个别专家忙不过来。

需要强调一点：Router本身的计算量很小，线性层的参数规模大约为H×E，相比专家的参数量可以忽略。但Router的分配结果决定了后续所有计算和通信的走势，所以对Infra影响很大。

1.3 一个MoE层的前向骨架

用伪代码看，一个简化版MoE前向过程长这样：

def moe_ffn(x, router, experts, top_k=2): # x: [T, H]，T是token数，H是hidden size # 1. 路由打分 logits = router(x) # [T, E] probs = softmax(logits, dim=-1) # [T, E] # 2. 选择Top-k专家，并做权重归一化 topk_probs, topk_idx = topk(probs, k=top_k) topk_probs = topk_probs / topk_probs.sum(dim=-1, keepdim=True) # 3. dispatch：把每个token的hidden state发送到对应专家 expert_tokens = dispatch(x, topk_idx) # 4. 每个专家独立计算自己分到的token expert_outputs = [] for i, expert in enumerate(experts): expert_outputs[i] = expert(expert_tokens[i]) # 5. combine：把多个专家结果按权重加权求和 combined = combine(expert_outputs, topk_idx, topk_probs) # 6. 残差连接 return x + combined

这个骨架是后面所有细节的总纲。实际分布式场景里，dispatch和combine不是简单的数组索引，而是跨设备的All-to-All通信；expert计算也不是循环调E次，而是把不同专家的token拼好后做分组GEMM。但思路都是围绕这六步展开的。

2. 计算过程逐段拆解：从路由到回调

2.1 Router打分和Top-k选择

Router接收每个token的hidden state，经过一个线性层得到长度为E的logits。训练早期为了让路由分布更均匀，有些实现会在logits上加一点噪声，相当于让Router在探索阶段不要过早固化到某几个专家上。这个噪声只在训练时加，推理时去掉。

Softmax把logits转成概率后，选择Top-k个专家。有一个细节值得注意：如果先取Top-k再对k个logits做Softmax归一化，和先做全量Softmax再取Top-k再归一化，最终结果是一致的。因为归一化时公共分母会约掉。实际工程里为了少算E个概率值，很多实现会直接用Top-k的logits做sparse softmax，效果一样，但省掉一大部分元素级计算。

Top-1和Top-2的差别不只是精度。Top-1时，每个token只去一个专家，combine阶段不需要做加权求和，直接取那一个结果即可，计算和通信都是“单通道”。Top-2时，token需要把hidden state发给两个专家，再取回两个结果做加权平均，dispatch和combine的通信量翻倍，专家侧的计算量也翻倍。很多MoE模型会选择Top-2，因为多一个专家带来更丰富的表征，但代价是计算和带宽都跟着涨，Infra上必须提前把这个账算清楚。

2.2 dispatch、compute、combine三步流水

token被路由到专家后，真正的计算过程可以分成三步。

dispatch是把token按专家索引重新分组。如果所有专家在同一张卡上，这是一个本地的数据重排操作，把属于同一个专家的token行凑成连续块，方便后面做GEMM。如果专家分布在多张卡上，dispatch就是一个跨设备的All-to-All通信：源设备按照目标设备分组，把自己的token发过去。注意，每个token可能同时被分到多个专家，所以dispatch阶段会复制token，而不是简单地把token移动走。

compute阶段发生在专家所在的设备上。此时每张设备手上有“属于自己的若干专家”以及“路由到这些专家的token”。一个设备上可能有多个专家，每个专家分到的token数量不等。很多实现会把这些token拼接成大矩阵，做分组GEMM，也就是一次kernel调用里完成多个不同大小的矩阵乘法，避免一个专家一个kernel带来的启动开销。

combine阶段把多个专家的输出重新汇合。每个token需要把它在所有被选中专家上的输出按权重加起来。专家在不同设备上的话，这一步又是一次All-to-All，把结果送回token原本所在的设备。整个过程听起来像“发出去、算完、收回来”，但数据路径上是两条通信：一次dispatch，一次combine。这也是MoE层和稠密层在Infra上最大的差异——稠密层只需要片内GEMM，MoE层还必须支付跨设备的数据搬运开销。

2.3 容量因子与token丢弃：一个数值例子

路由结果不可能完全均匀。某个专家可能收到很多token，另一个专家可能很闲。为了让每张卡的负载可控，MoE引入专家容量（expert capacity）的概念。

假设一层有T=4096个token，E=8个专家，Top-k=2。理想情况下，每个专家应该处理大约：

T * k / E = 4096 * 2 / 8 = 1024 个token

但这个1024只是平均期望。如果某个专家被分配了1300个token，而容量因子设为1.0，那么这个专家只能处理前1024个token，剩下276个token会溢出。溢出的token如何处理取决于实现，常见做法是跳过这个专家的计算，直接走残差连接，相当于这个token“少学了一次该专家”。

容量因子就是用来控制这个“溢出”容忍度的：

容量因子越小，越可能丢token，影响模型质量；容量因子越大，容忍度越高，但计算量也会被撑大，甚至出现专家实际很闲、但按容量分配的资源却很浪费的情况。工程上常见的做法是在1.0到2.0之间调，具体取多少要看负载均衡损失压到什么程度。我见过一些实现会在训练初期用偏大的容量因子保证质量，后期再把容量因子压下来做收敛。

2.4 负载均衡辅助损失是怎么作用到计算上的

负载均衡损失的目标很朴素：让token尽可能均匀地分配给所有专家。如果一个专家收到太多token，它所在的设备就会成为热点，All-to-All通信也要等最慢的那一方，整体step时间会被拉长。

最常用的辅助损失长这样：

loss_aux = alpha * E * sum_i(f_i * P_i)

其中：

• f_i是第i个专家在本batch里实际分到的token比例；
• P_i是Router对所有token分配给第i个专家的平均概率。

当每个专家都均匀收到token时，f_i约等于1/E，P_i也约等于1/E，这个loss会接近一个归一化常数。如果某专家特别热门，f_i和P_i都偏高，loss就会变大，训练时通过梯度把Router的logits往均匀方向推。

这里有一个Infra上的微妙点：辅助损失不直接影响前向计算的吞吐，真正影响计算效率的是它训练出来的最终分布。如果辅助损失系数太小，专家不均衡，推理时照样热点；如果系数太大，Router被过度“平均化”，专家特征分化变差，模型精度反而下降。系数本身通常是个很小的值，但在MoE训练里它对计算效率的影响比很多超参数都大。

3. AI Infra视角：显存、通信、并行的落点

3.1 参数规模与计算量“解耦”的另一面

先说一个具体的计算例子。假设hidden size H=4096，专家内部的FFN隐藏层维度F=14336，专家数E=64，Top-k=2。

单个专家的参数量大约是：

W1: H * F = 4096 * 14336 ≈ 58.7M W2: F * H = 14336 * 4096 ≈ 58.7M 合计 ≈ 117.4M

64个专家加起来就是约7.5B参数。但每个token实际只访问2个专家，也就是约235M参数。这和64个专家全部激活的15B参数相比，计算量少了32倍。

不过“只访问2个专家”不等于“只放2个专家进显存”。无论token去了哪个专家，所有64个专家的权重都必须驻留在设备上，否则推理时遇到被调度的专家还要现场加载。这就带来AI Infra里很经典的问题：参数规模和计算规模解耦后，显存压力没有解耦。

更麻烦的是推理场景。GPU做GEMM时，是否算得快，取决于计算强度和硬件算力带宽比。BF16权重下，每个专家每处理n_i个token，它的算术强度大约等于n_i（单位是FLOP/Byte，因为每token计算量约4×H×F，专家权重加载字节数约4×H×F）。当每个专家每次只分到个位数token时，算术强度只有个位数，远低于主流训练卡的上千FLOP/Byte。也就是说，MoE推理经常是显存带宽瓶颈，而不是算力瓶颈。

3.2 All-to-All通信是MoE的“隐形税”

MoE层每次路由都可能把token发到另一张卡上，这个通信模式叫All-to-All。名字很直白：每张卡都要向其他卡发数据，也要从其他卡收数据，没有一个固定的“主从”关系。

通信量可以粗略估算。假设每层有T个token，hidden size为H，Top-k为K，数据用BF16（每个数2字节）。那么一次dispatch需要搬移的总数据量约为：

T * K * H * 2 字节

combine阶段同样要搬移一次。如果这层数据在D张卡的EP组内做All-to-All，平均每张卡收发各约：

T * K * H * 2 / D 字节

举个例子：T=65536，H=4096，K=2，D=16，每层dispatch总数据量约1.07GB，平摊到每张卡约67MB发送、67MB接收。combine再来一轮，一个MoE层就有约134MB/卡的通信量。如果整个模型有几十个MoE层，这个数字会迅速变成GB级。

但通信量不是全部，更隐蔽的是通信次数。All-to-All往往是在所有参与卡之间同步完成的，只要有一张卡慢一点，整轮通信都要等它。专家负载越不均衡，这种“木桶效应”越明显。这也是为什么前面说的负载均衡损失，对AI Infra的重要性远大于对模型精度的重要性。

3.3 专家并行与其他并行策略怎么咬合

训练一个大MoE模型，光靠Data Parallel不够，因为每张卡都要塞下全部专家参数，显存撑不住。业界常见的做法是专家并行（Expert Parallel）：把E个专家切到不同设备上，每张设备只保存一部分专家。

有了专家并行之后，token就必须跨设备路由。在数据流上，输入batch会被切到各张卡，每张卡先拥有部分token，然后通过All-to-All把token发到目标专家所在的卡。专家并行组内部，每张卡负责自己的专家计算，算完再All-to-All把结果还回去。

这里要注意和其他并行策略的组合。Attention部分的hidden维度可以走张量并行，但MoE专家通常不按hidden维度切分，而是按专家ID切分。两种并行方式混合时，MoE层和Attention层之间的hidden state可能需要额外的聚合或规约操作，形成通信接口。实际框架里会把专家并行单独划成一个通信组，避免和Attention的通信组混在一起。

训练时还有两个容易被忽略的点：

• 优化器状态同样会放大显存压力。Adam优化器对每个参数会保存额外的状态，7.5B专家参数带来的优化器状态远大于参数本身。专家并行的同时，往往还需要对优化器状态做分片，否则每张卡仍然可能OOM。
• 通信和计算要尽量重叠。dispatch发送完成后，可以先算本地已有的那部分专家，等剩下数据到齐再算另一个批次。如果实现得粗糙，All-to-All期间计算单元会空转，整个step时间会很难看。

3.4 推理阶段别把MoE当稠密模型用

推理时MoE最典型的问题是“参数都占显存，但每次只有少数参数被算”。这意味着同一张卡上的内存带宽成为稀缺资源，而不是算力。

一个常见做法是尽量提高batch size。batch越大，平均每个专家分到的token越多，算术强度越高，GPU算力利用率才会起来。如果batch很小，比如批量推理只有几十个token，每个专家可能只分到两三个token，此时All-to-All通信和权重加载的时间都会盖过实际计算时间。

针对特别大的专家数，还可以考虑把不常用的专家权重放到CPU内存或远端存储，用到时再加载。但这会引入额外延迟，需要设计预取策略。最简单的方案是先把经常被路由到的“热门专家”留在显存，冷门专家做成按需加载。问题在于“热门”可能随着输入分布变化，所以这个策略并不总是稳定。

推理阶段MoE另一个容易踩的坑是KV Cache和专家权重抢显存。MoE替换的是FFN部分，Attention仍然需要KV Cache。如果一个模型的KV Cache很大，再叠加全部专家权重，显存会非常局促。调显存时，不能只盯着专家参数，要把KV Cache、中间激活、通信缓冲一起算进去。

4. 实操中的热点排查与调优记录

4.1 专家不均衡：现象、定位、解法

专家不均衡是MoE最常见的问题。现象很直接：训练日志里某个专家被分到的token数明显偏高，对应设备上的GPU利用率也比其他卡高，整机训练时间被最忙的那张卡拖住。

定位时不要只看总loss，要把每个step的token分布统计打出来。可以每隔一定步数记录Router输出的专家索引，统计每个专家的token数量。早期的稀疏分配是随机的，如果训练中期某个专家占比超过20%，而平均是1/E，那基本可以确认需要干预。

解法通常有几个方向：

• 调高负载均衡损失的系数，但不要一次跳太多，一般是从0.01往0.05甚至0.1试，同时观察精度曲线。
• 检查Router输出logits是否有异常：如果某几个专家的logits始终特别高，可能是初始化问题，也可能是某个专家的权重更新过快，可以做梯度裁剪。
• 如果训练已经进入后期，分布基本固化，单纯调损失系数作用有限，需要回退到较早的checkpoint重新用不同的噪声种子训练。

从Infra角度，还可以在系统层做“克制”：当某个专家超过容量上限时，溢出的token直接跳过该专家，而不是让计算排队。这是一种保护机制，能避免单卡过载，但代价是部分token的信息流被切断，精度会受影响。

4.2 All-to-All成为瓶颈怎么定位

如果你发现GPU算力利用率不高，但step时间很长，很可能卡在通信上。在profiler上会看到一大段All-to-All等待时间，前后计算kernel都很快，就中间一段在等数据。

定位的时候要区分是带宽不够还是同步等待。如果是网络带宽不够，提高batch size或减少通信数据量会有效；如果是同步等待，也就是某张卡发得慢，那更多是负载不均或某张卡计算慢。一个很笨但有效的办法是：把All-to-All替成本地拷贝，如果step时间下降很多，说明通信开销占比确实高；再进一步缩小EP组范围，看通信时间是否等比例下降。

优化手段里，我实测下来最有效的是“合并通信”：把多个专家层的dispatch合并成一次大的All-to-All，而不是每层发一次。虽然数据总量没变，但通信次数大幅减少，同步次数少了，总耗时常常能下降30%以上。如果框架支持分组All-to-All，也可以把同一个目标设备的数据打包成连续buffer，减少小消息带来的额外开销。

4.3 显存OOM不是直接减少batch那么简单

MoE训练中遇到OOM，很多人第一反应是减batch。但如果在专家并行下，OOM往往不是因为batch太大，而是因为显存里还躺着额外的通信缓冲、优化器状态和中间激活。

通信缓冲尤其容易被忽略。All-to-All通常需要提前分配发送和接收buffer。如果MoE层的token数很大，这个buffer本身就是GB级。减少通信缓冲的一个方式是走“分批All-to-All”，也就是把token切成几段，一段一段发，避免一次性占用过大buffer。代价是通信次数变多，需要权衡。

另一个常见原因是分组GEMM的中间激活。某些实现为了把不同专家的token拼成一个大矩阵，会先padding到统一长度。padding越多，浪费的显存越多。如果观察到专家分布很不均匀，padding浪费会非常明显。调大负载均衡损失，除了改善计算热点，也能顺带降低padding浪费。

4.4 MoE模块调优检查单

我把平时排查MoE模块问题时会看的点整理成了一张检查单，按顺序走一遍，大部分问题都能定位到：

还有一个小细节很多人会忽略：MoE层里的Router其实特别敏感，它对训练后期的分布稳定性影响很大。建议在Router参数上做独立的梯度裁剪，或者把它的学习率调得比专家网络略小一点。我在某个垂直场景MoE Demo上试过，只调低Router学习率，专家分布方差就明显变小，通信尾部等待也减轻了很多。如果你也在调MoE，先别急着动大结构，把这几个小参数捋一遍，往往比换模型结构更见效。