1. 从一次显存爆炸说起为什么MoE值得单独开一篇第一次在生产环境里部署MoE模型的时候我盯着监控面板上那条显存曲线愣了很久。模型权重文件明明只有几十GB理论上两张卡就能塞下结果加载完直接OOM。后来把参数配置翻出来逐行核对才发现问题出在专家并行Expert Parallelism的切分策略上——专家层没有均匀分布到各张卡导致某张卡上堆了过多专家显存直接被打爆。这个坑让我意识到MoEMixture of Experts混合专家跟传统的Dense模型在计算逻辑上完全是两码事。你不能用部署Llama的那套直觉去套DeepSeek或者Qwen的MoE版本否则一定会在某个环节翻车。这篇笔记就围绕MoE的计算原理和过程展开把路由机制、专家计算、负载均衡、GEMM形状变化这些核心环节拆开讲清楚。适合已经在跑vLLM或者准备上手MoE模型推理的工程师也适合想搞清楚“MoE到底省不省显存”这个问题的朋友。我会尽量用实际部署中遇到的场景来解释原理而不是堆公式。2. MoE的整体计算架构稀疏激活到底省在哪里2.1 从Dense到MoE参数量与计算量的解耦传统Dense模型里每一层FFNFeed-Forward Network的参数量是固定的输入token经过这层时所有参数都会参与计算。一个7B的模型每个token过一遍FFN就要消耗完整的7B参数对应的计算量。MoE的思路完全不同。它把原来一个大FFN拆成N个小的FFN每个叫一个“专家”Expert然后加一个路由网络Router/Gate对每个token只选出Top-k个专家来计算。这样一来总参数量可以做得很大比如DeepSeek-V3总参数671B但每个token实际激活的参数量只有37B左右。单token计算量跟激活参数量挂钩而不是总参数量。这就是“稀疏激活”的核心含义。你可以理解为模型容量知识存储能力由总参数量决定而计算成本由激活参数量决定。两者解耦之后就能用相对可控的算力去跑一个知识容量极大的模型。但这里有个常见的误解需要澄清MoE省的是计算量不是显存。所有专家的参数都得加载到显存里待命因为路由是动态的你事先不知道哪个token会选哪个专家。所以显存占用看的是总参数量计算开销看的才是激活参数量。2.2 路由网络一个token该找谁算路由网络是MoE的调度中心。它的工作流程大致是这样拿到token的隐状态向量维度通常是hidden_size。通过一个线性层gate layer映射到N维N是专家数量。对N维输出做Softmax得到每个专家的概率分布。选出概率最高的Top-k个专家。把这k个专家的输出按概率加权求和。用伪代码表示就是# 假设 hidden_states: [num_tokens, hidden_size] # gate_weight: [num_experts, hidden_size] logits hidden_states gate_weight.T # [num_tokens, num_experts] probs softmax(logits, dim-1) topk_probs, topk_indices topk(probs, k2, dim-1) # topk_probs: [num_tokens, 2], topk_indices: [num_tokens, 2]看起来简单但实际部署时这里有几个关键决策点。Top-k的k取多少常见的是k2比如Mixtral 8x7B和DeepSeek系列都用2。k越大计算量越大但路由的容错性更好。k1的话计算最省但一旦路由选错这个token就没有补救机会。还有一个细节是Softmax的位置。有些实现是在Top-k之后再做Softmax归一化有些是先Softmax再取Top-k。这两种方式在数学上不等价对训练稳定性和推理结果都有影响。目前主流实现倾向于先算logits取Top-k之后再对选中的k个值做归一化这样能保证权重和为1同时避免未选中专家的概率干扰。2.3 专家层计算GEMM形状的变化这是MoE计算中最容易被忽视但影响最大的环节。在Dense模型里FFN的计算是两个连续的GEMM通用矩阵乘法# 第一层升维 intermediate activation(input W1) # [T, H] [H, I] - [T, I] # 第二层降维 output intermediate W2 # [T, I] [I, H] - [T, H]其中T是token数量H是hidden_sizeI是intermediate_size。到了MoE这里因为每个token选的专家不同你不能直接把所有token拼在一起做一个大GEMM。实际的计算过程变成了根据路由结果把token按专家分组。比如专家0收到了100个token专家1收到了80个token以此类推。对每个专家用它自己的W1和W2对分到的token做GEMM。把各专家的输出按原始token顺序散射scatter回去再按路由权重加权求和。这个过程在代码里通常叫“permute group GEMM unpermute”。permute是把token按专家重排group GEMM是对每个专家分别做矩阵乘unpermute是还原顺序。问题在于每个专家收到的token数量是不固定的。这导致GEMM的M维度token数动态变化GPU的利用率会受影响。如果某个专家只收到几个token做一次GEMM的启动开销可能比计算本身还大。这就是为什么MoE推理对kernel实现的要求比Dense模型高得多。3. 负载均衡MoE训练和推理中最容易翻车的地方3.1 为什么会出现负载不均路由网络是学习出来的它天然倾向于把token分配给“表现好”的专家。训练初期这种倾向还不明显但随着训练进行某些专家会被越来越多地选中另一些则几乎收不到token。这就是“赢者通吃”现象。负载不均的后果很直接训练时某些专家所在的GPU算力被打满其他GPU闲置整体吞吐下降。更严重的是闲置GPU上的专家参数收不到梯度逐渐变成“死专家”。推理时如果用了专家并行某个专家堆了太多token对应GPU的显存和算力都会成为瓶颈拖慢整个batch的延迟。3.2 辅助损失与容量因子为了解决这个问题MoE训练通常会加一个辅助损失auxiliary loss惩罚路由分布的方差。具体做法是计算每个专家被选中的频率然后让频率分布尽量均匀。这个loss的权重需要调太大影响主任务效果太小起不到均衡作用。另一个手段是设置容量因子capacity factor。每个专家能处理的token数量有一个上限超过上限的token会被丢弃或者走残差连接直接跳过。容量因子的计算公式通常是capacity capacity_factor * (num_tokens / num_experts)capacity_factor一般取1.0到1.5之间。取1.0意味着每个专家最多处理平均数量的token超了就丢。取1.5留了一些余量。这个参数在训练时很重要推理时一般可以放宽或者不设限。3.3 实际部署中的专家并行策略在生产环境部署MoE模型时专家并行的切分方式直接影响显存和通信开销。常见的策略有策略做法优点缺点专家均匀分布每个GPU放相同数量的专家显存均衡可能跨卡通信多按热度分布热门专家放同一张卡减少通信显存可能不均专家张量并行混合专家分卡每个专家内部再做TP适合超大专家实现复杂以DeepSeek-V3为例它有256个路由专家如果部署在8张卡上每张卡放32个专家。但实际部署时还要考虑共享专家Shared Expert的放置以及attention层的并行方式。这些因素叠加起来让MoE的部署配置比Dense模型复杂不少。我在实际配置时踩过一个坑vLLM的--tensor-parallel-size和专家并行的关系不是简单的倍数关系。如果你设了TP8但专家数量不能被8整除vLLM可能会报错或者自动padding导致显存浪费。所以部署前一定要确认专家数量和你打算用的并行度是否匹配。4. 完整推理流程拆解一个token在MoE层里经历了什么4.1 从输入到路由决策假设我们有一个batch的输入经过前面的attention层之后得到了隐状态hidden_states形状是[batch_size * seq_len, hidden_size]。这个张量进入MoE层后第一步就是路由计算。路由的gate layer是一个[num_experts, hidden_size]的权重矩阵。做一次矩阵乘之后得到[num_tokens, num_experts]的logits。然后取Top-k得到每个token选中的专家索引和对应的权重。这里有个实现细节Top-k操作在GPU上做的时候如果num_experts很大比如256排序开销不可忽略。有些实现会用近似方法比如先做一次阈值筛选再在候选集里做精确Top-k。vLLM在这方面做了不少优化具体kernel实现可以参考它的fused_moe相关代码。4.2 Token重排与分组GEMM拿到路由结果后下一步是把token按专家分组。假设有256个专家每个token选2个专家那么总共会产生num_tokens * 2个“token-专家”对。这些对需要按照专家索引排序让同一个专家的token在内存里连续排列。这个排序操作在CUDA里通常用argsort或者bincount cumsum来实现。排序之后每个专家分到的token就变成了一段连续的区间。然后对每个专家取出对应的token做GEMM。这里的关键是group GEMM的实现。如果对每个专家单独调用一次cuBLASkernel启动开销会很大。所以实际实现会用一些技巧比如把所有专家的权重拼成一个大矩阵然后用block-diagonal的稀疏GEMM来计算。或者用cutlass的grouped GEMM接口一次启动处理多个专家。4.3 SwiGLU激活与加权求和专家内部的FFN通常用SwiGLU激活函数。SwiGLU的公式是SwiGLU(x) Swish(xW1) * (xW3)其中Swish(x) x * sigmoid(x)。注意这里有两个权重矩阵W1和W3所以每个专家的参数量比标准FFN多了一个矩阵。这也是为什么MoE模型的单专家参数量看起来不大但总参数量很可观。计算完专家输出后需要把结果按原始token顺序散射回去。因为一个token可能选了2个专家所以它在两个专家那里各有一个输出需要按路由权重加权求和。这个加权求和的步骤通常在unpermute之后做可以用一个简单的element-wise乘加完成。4.4 共享专家的特殊处理DeepSeek系列的MoE架构里有一个特殊设计共享专家Shared Expert。这个专家对所有token都激活不经过路由选择。它的作用是捕捉一些通用的、所有token都需要的特征减轻路由专家的负担。在计算时共享专家的输出会直接加到路由专家的加权输出上。这意味着共享专家的计算是Dense的每个token都要过。所以共享专家的参数量不能太大否则就失去了稀疏激活的意义。DeepSeek-V3里共享专家的intermediate_size设置得比路由专家小一些具体数值可以参考它的技术报告。5. 性能优化与常见问题排查5.1 GEMM效率低下的排查思路如果你发现MoE层的推理速度远低于预期第一个要查的就是GEMM效率。可以用Nsight或者nvprof抓一下kernel耗时看看group GEMM那部分的利用率。常见原因有几个专家分到的token太少如果batch size小或者路由过于集中某些专家可能只分到几个token。这时候GEMM的M维度很小GPU利用率上不去。解决办法是增大batch size或者调整路由的temperature让分布更均匀。内存访问不连续permute之后token在内存里的排列可能不是最优的。如果group GEMM的实现没有处理好内存对齐会出现大量的非合并访问带宽利用率下降。kernel启动开销如果专家数量很多每个专家单独启动一个kernel启动开销会累积。这时候需要用grouped GEMM或者fused kernel来减少启动次数。5.2 显存不足的几种可能MoE模型显存不足的原因跟Dense模型不太一样排查时要有针对性现象可能原因排查方法加载权重时OOM专家并行切分不均检查每张卡的专家数量是否一致推理过程中OOMKV Cache增长 专家激活峰值监控不同batch下的显存曲线特定层OOM某层专家参数量异常大逐层打印参数量多卡OOM通信buffer占用调整NCCL buffer大小我遇到过一次比较隐蔽的情况模型加载时显存刚好够但跑起来之后随着batch内token数量增加某些专家的中间激活值intermediate activation占用的显存超出了预期。后来把max_num_batched_tokens调小才解决。这个参数在vLLM里控制单次前向的最大token数对MoE模型尤其重要。5.3 路由崩溃与死专家训练MoE时如果发现loss突然飙升或者某些专家的被选次数长期为0大概率是路由崩溃了。排查步骤打印每个专家在一个batch内的被选次数看分布是否严重倾斜。检查辅助损失的权重是否太小。检查学习率是否过大导致路由权重更新过猛。考虑加router z-loss惩罚logits的绝对值过大让路由分布更平滑。推理时一般不会出现死专家因为权重已经固定了。但如果模型在训练时就有死专家推理时这些专家的参数相当于浪费了显存。所以训练阶段的负载均衡做得好不好直接影响推理部署的成本。5.4 vLLM部署MoE模型的参数配置要点在vLLM里部署MoE模型有几个参数需要特别注意vllm serve deepseek-ai/DeepSeek-V2-Lite \ --tensor-parallel-size 2 \ --max-model-len 8192 \ --max-num-batched-tokens 4096 \ --gpu-memory-utilization 0.9 \ --enable-expert-parallel--enable-expert-parallel开启专家并行。注意这个选项和--tensor-parallel-size的配合关系不是所有版本都支持任意组合。--max-num-batched-tokens控制单次前向的token总数。MoE模型对这个参数比Dense模型更敏感设太大容易OOM设太小吞吐上不去。--gpu-memory-utilization显存利用率上限。MoE模型的显存占用波动比Dense模型大建议留一些余量不要设到0.95以上。另外不同版本的vLLM对MoE的支持程度不一样。有些版本对特定模型架构比如DeepSeek的MLA MoE组合有专门优化升级版本时最好看一下release notes里有没有相关改动。6. 几个容易被忽略的细节6.1 Top-k的k值对推理结果的影响k1和k2在推理结果上可能有肉眼可见的差异。k1时每个token只走一个专家路由一旦选错这个token就没有修正机会。k2相当于做了一个ensemble结果更鲁棒。但k2的计算量接近k1的两倍因为要过两个专家所以这是一个精度和速度的权衡。实际部署时如果模型训练时用的是k2推理时改成k1效果可能会下降。反过来训练时k1推理时k2效果不一定提升因为模型没有学过如何组合两个专家的输出。所以k值最好跟训练时保持一致。6.2 专家并行的通信开销专家并行意味着token需要在不同GPU之间传输。如果一个token选的专家在不同的卡上就需要跨卡通信。这个通信量跟batch size、Top-k、专家分布都有关系。在NVLink环境下通信开销相对可控。但如果是PCIe环境跨卡通信可能成为瓶颈。这时候可以考虑把热门专家放在同一张卡上减少跨卡传输。或者用一些路由约束让token尽量选择本卡的专家。不过这些优化手段需要根据具体硬件拓扑来调整没有通用方案。6.3 量化对MoE的影响MoE模型的量化比Dense模型更复杂因为不同专家的权重分布可能差异很大。如果用统一的量化参数某些专家可能精度损失严重。有些实现会对每个专家单独做量化校准但这样会增加部署复杂度。目前实践中MoE模型常用的是W8A8或者W4A16量化。W4A16对显存的节省更明显因为MoE的瓶颈主要在显存容量而不是计算。但4bit量化对路由精度的影响需要评估特别是gate layer的量化可能会改变Top-k的选择结果。7. 写在最后MoE的计算过程拆开来看每一步都不算特别复杂但组合在一起之后工程实现的复杂度就上来了。路由、permute、group GEMM、unpermute、加权求和每个环节都有优化空间也都有踩坑的可能。我在实际部署中最大的体会是不要用Dense模型的直觉去调MoE模型。显存够不够不能只看参数量还要看专家并行的切分方式速度快不快不能只看FLOPs还要看GEMM的M维度是否够大效果好不好不能只看loss还要看路由分布是否健康。如果让我给刚接触MoE的朋友一个建议那就是先把一个小规模的MoE模型比如DeepSeek-V2-Lite或者Qwen的MoE版本在单卡上跑通把路由和专家计算的流程用日志打出来看一遍。搞清楚每个token到底走了哪些专家每个专家的负载是多少然后再去调并行策略和性能参数。这样比直接上大模型盲调要高效得多。