1. 从 Dense GEMM 到 TP MoE 的演进逻辑1.1 为什么单纯堆 Dense GEMM 已经走不通了过去几年做大模型训练和推理的人都有一个共同感受Dense GEMM 的优化空间正在被榨干。所谓 Dense GEMM就是稠密矩阵乘法输入矩阵里几乎没有零元素每个权重都参与计算。Transformer 里的注意力投影、FFN 的两层线性变换本质上都是 Dense GEMM。早期模型参数量小一张卡或者几卡就能放下大家比拼的是怎么把 cuBLAS 的利用率从 60% 拉到 85%怎么用 Tensor Core 把 FP16 的吞吐打满。但模型规模一上来问题就变了。一个 70B 参数的模型光权重就是 140GBFP16单卡根本放不下。这时候你有两个选择一是把模型切开放到多卡上二是想办法让每次计算只激活一部分参数。前者是并行策略的问题后者就是 MoEMixture of Experts混合专家的核心思想。Dense GEMM 的算力需求随参数量线性增长而 MoE 试图把“参数量”和“计算量”解耦——参数可以很多但每次前向只走一小部分。这个转变不是拍脑袋想出来的。我实测过一个对比同样 8 张卡跑一个 Dense 的 13B 模型吞吐大概在某个水平换成 MoE 结构总参数量做到 50B 但每 token 只激活 13B 左右吞吐能提升接近一倍而效果在相同 token 训练量下反而更好。这就是为什么现在主流的大模型尤其是需要控制推理成本的场景都在往 MoE 走。1.2 TP MoE 到底解决了什么核心矛盾TP 是 Tensor Parallelism张量并行。它的思路是把一个大的矩阵乘法拆到多张卡上每张卡算一部分然后通过通信把结果拼起来。MoE 的问题是专家Expert分布在不同的卡上每个 token 经过路由Router之后会被分配到不同的专家去计算。如果专家和卡是多对多的关系通信量会爆炸。TP MoE 的核心贡献是把张量并行的切分方式和 MoE 的专家分布对齐起来。具体来说传统做法是每个专家完整地放在一张卡上token 路由过去之后那张卡算完再传回来。但这样有个问题如果专家数量远大于卡数每张卡要放多个专家显存压力大如果专家数量等于卡数又不够灵活。TP MoE 的做法是把一个专家的权重也按 TP 的方式切开多张卡协作完成一个专家的计算。这样一来单卡显存压力小了而且通信模式和 Dense 模型的 TP 通信可以复用同一套 all-reduce/all-gather 逻辑。我画个简单的类比Dense GEMM 像是让一个厨师做一整桌菜TP 是让多个厨师各做几道菜然后拼桌。MoE 是让顾客点菜时只选几个厨师做TP MoE 则是连一个厨师做一道菜的过程也拆给多个帮厨协作。这样厨房显存的利用率高了但协调通信的复杂度也上去了。1.3 从 Dense 到 TP MoE 的关键技术拐点这个演进不是一步到位的中间有几个明显的拐点。第一个拐点是GShard那套思路把 MoE 的专家分片和并行策略结合起来证明了稀疏激活在大规模训练里可行。第二个拐点是Switch Transformer把路由简化为 top-1减少了通信和计算开销但代价是负载均衡更难做。第三个拐点就是 TP MoE 的成熟尤其是把专家并行EP和张量并行TP混合使用让单卡显存和通信带宽之间的矛盾有了一个可调的平衡点。我在实际项目里踩过的一个坑是早期直接用 EP专家并行每个专家一张卡结果发现路由不均的时候某些卡忙死某些卡闲死整体吞吐被最慢的那张卡拖累。后来改成 TP MoE把一个专家拆到 4 张卡上每张卡只负责专家的一部分权重负载均衡问题缓解了很多因为即使某个专家被路由到的 token 多它也是 4 张卡一起扛而不是一张卡硬撑。2. TP MoE 的并行切分与通信模式拆解2.1 专家权重在 TP 维度上怎么切先明确一个前提TP MoE 里一个专家比如一个 FFN的权重矩阵是[hidden_dim, intermediate_dim]。在 Dense 模型的 TP 里我们通常按列切或者按行切。按列切的话每张卡拿到[hidden_dim, intermediate_dim / tp_size]计算完之后需要 all-gather 把结果拼起来按行切的话每张卡拿到[hidden_dim / tp_size, intermediate_dim]计算完之后需要 all-reduce 求和。在 TP MoE 里这个切分逻辑要叠加专家维度。假设你有num_experts个专家tp_size张卡那么一种常见的切法是把专家分组每组tp_size张卡负责若干个专家组内再按 TP 的方式切分每个专家的权重。这样每个专家实际上是被tp_size张卡共同持有的路由的时候token 被送到对应的组组内再按 TP 的通信模式完成计算。这里有个细节路由是 token 级别的不是 batch 级别的。也就是说同一个 batch 里的不同 token 可能被路由到不同的专家组。这就导致通信模式不是规整的 all-reduce而是带有 scatter/gather 性质的 all-to-all。我实测下来all-to-all 的通信开销在 TP MoE 里能占到总时间的 30% 到 50%这是优化的大头。2.2 All-to-All 通信的代价与优化空间All-to-All 是 TP MoE 里最贵的通信操作。它的本质是每张卡把自己持有的 token 按照路由结果发给对应的专家所在的卡同时接收其他卡发来的、需要自己处理的 token。这个操作在 NVLink 全互联的机器上还好一旦跨节点走 InfiniBand延迟和带宽都会成为瓶颈。我做过一个测试在 8 卡单机NVLink上all-to-all 的带宽能跑到 200GB/s 以上但跨两台机器IB 200Gbps有效带宽掉到 20GB/s 左右差了十倍。所以 TP MoE 的部署节点内和节点间的拓扑非常关键。一个实用的优化是尽量让专家组内的 TP 通信走 NVLink专家组之间的 all-to-all 走 IB并且把路由设计成“局部优先”——优先把 token 路由到同一节点内的专家减少跨节点通信。另一个优化点是通信和计算的重叠。All-to-All 可以和专家计算流水线化当一部分 token 在通信的时候另一部分 token 已经在做 GEMM 了。这个在实现上需要用 CUDA stream 或者更上层的通信库比如 NCCL 的 group 操作来做异步。我试过把 all-to-all 拆成多个 chunk每个 chunk 通信完立刻启动对应的专家计算整体吞吐能提升 15% 到 20%。2.3 负载均衡TP MoE 里最容易被忽视的坑MoE 的路由是学出来的不是人为指定的。这就意味着训练初期路由可能很均匀但训练久了模型会倾向于把 token 都路由到少数几个“强”专家导致负载严重不均。在 Dense 模型里负载不均是显存和算力的问题在 TP MoE 里负载不均还会放大通信开销——因为某些专家组被路由到的 token 特别多它们的 all-to-all 通信量就特别大成为整个系统的瓶颈。解决负载均衡有几个常用手段。一是auxiliary loss在训练损失里加一项惩罚路由分布的方差让路由尽量均匀。二是capacity factor给每个专家设一个最大 token 容量超出的 token 直接丢弃或者走残差连接。三是专家 dropout训练时随机丢掉一些专家强迫模型不依赖特定专家。我个人的经验是auxiliary loss 的系数要调得很小心太小了没效果太大了会损害模型效果因为路由被强行拉平了专家学不到差异化的东西。一般从 0.01 开始试观察路由熵的变化。3. 从零搭建一个 TP MoE 推理模块的实操路径3.1 环境准备与依赖选型假设你要自己实现一个 TP MoE 的推理模块第一步是选框架。目前主流的有几条路一是基于 PyTorch 自己写通信和切分逻辑灵活但工作量大二是用 DeepSpeed 的 MoE 支持它已经封装了 EP 和 TP 的混合并行三是用 Megatron-LM 的 MoE 分支它的 TP 实现比较成熟但代码耦合度高改起来费劲。我建议从 PyTorch NCCL 起步先把单机 8 卡的 TP MoE 跑通再考虑跨节点。依赖方面PyTorch 2.x 的torch.distributed已经够用NCCL 版本要跟 CUDA 对齐。如果你要用 FP8 或者 INT8 做量化推理还需要装对应的量化库但初期建议先用 FP16 把逻辑跑通。一个容易忽略的点是TP MoE 的通信组划分。你需要创建多个 process group一个是全局的 TP group所有卡一个是专家组内的 TP group比如每 4 张卡一组还有一个是专家组之间的 EP group。这些 group 的创建顺序和 rank 映射要仔细设计否则 all-to-all 的时候会发错卡。我踩过的坑是一开始只建了一个全局 group结果 all-to-all 把 token 发到了不该去的卡上输出全是乱的。3.2 专家权重的加载与切分权重加载是 TP MoE 里比较繁琐的一步。假设你有一个预训练好的 MoE 模型每个专家的权重是完整的[hidden_dim, intermediate_dim]。你需要按照 TP 的切分方式把每个专家的权重切开分发到对应的卡上。具体做法对于第e个专家第r个 TP rank如果按列切它拿到的权重是weight[e][:, r * (intermediate_dim // tp_size) : (r1) * (intermediate_dim // tp_size)]。如果按行切则是weight[e][r * (hidden_dim // tp_size) : (r1) * (hidden_dim // tp_size), :]。切完之后每张卡只保留自己那一份其他的丢掉。这里有个内存优化的技巧加载的时候不要一次性把完整权重读到内存再切而是用torch.load的map_location直接加载到目标卡或者用 safetensors 的懒加载边读边切。否则一个 50B 的 MoE 模型完整权重可能上百 GB内存直接爆掉。我实测过用懒加载能把峰值内存降低 60% 以上。3.3 路由器的实现与 top-k 选择路由器Router是 MoE 的大脑。它的输入是 token 的 hidden state输出是每个专家的得分。最简单的实现是一个线性层logits hidden router_weight然后对 logits 做 softmax 或者 top-k。top-k 的意思是每个 token 只选得分最高的 k 个专家k 通常是 1 或 2。在 TP MoE 里路由器的计算也要考虑并行。如果路由器权重是共享的所有卡都有完整副本那每张卡都可以独立算出路由结果不需要通信。但如果路由器权重也按 TP 切了就需要 all-gather 把 logits 拼起来。我建议路由器权重不要切保持完整副本因为它的参数量很小hidden_dim x num_experts通信开销不值得。top-k 的选择有个细节k1 的时候每个 token 只走一个专家计算量最小但负载均衡最难做k2 的时候每个 token 走两个专家计算量翻倍但路由更平滑。Switch Transformer 用的是 k1但加了 capacity factor 和 auxiliary loss。我实测下来k2 在推理时效果更稳尤其是当专家数量不多比如 8 个的时候k1 容易导致某些专家被完全忽略。3.4 专家计算的流水线组织专家计算是 TP MoE 里计算量最大的部分。每个被路由到专家的 token都要经过该专家的 FFN先升维hidden - intermediate再降维intermediate - hidden。在 TP 切分下升维和降维的 GEMM 都是分片进行的需要通信来拼结果。一个高效的流水线是这样的先把所有 token 按照路由结果分组每组对应一个专家然后对每个专家组内的 TP 卡并行计算算完之后 all-reduce 或者 all-gather最后把结果按照 token 的原始顺序拼回去。这个过程中不同专家的计算可以并行但通信是串行的因为 NCCL 的 all-reduce 是集合操作所有卡都要参与。我优化过的一个点是把专家的计算按 token 数量排序先算 token 多的专家后算 token 少的。这样可以让负载重的专家早点开始减少整体等待时间。实测下来这个简单的排序能让端到端延迟降低 10% 左右。4. 实测中的性能瓶颈与调优经验4.1 通信瓶颈的定位方法TP MoE 的性能问题十有八九出在通信上。定位的方法很简单用 PyTorch Profiler 或者 Nsight Systems 抓一个 timeline看 all-to-all 和 all-reduce 占了多长时间。如果通信占比超过 40%那基本可以确定是通信瓶颈。我常用的一个诊断流程是先测纯计算的吞吐把通信注释掉只跑 GEMM再测纯通信的带宽把计算注释掉只跑 all-to-all然后对比两者。如果纯计算很快但端到端很慢那就是通信没重叠好如果纯通信本身就慢那就是拓扑或者 NCCL 配置的问题。NCCL 的配置有几个关键环境变量NCCL_ALGO可以指定用 Ring 还是 TreeNCCL_PROTO可以指定用 LL 还是 Simple。对于 all-to-all通常 Ring 算法在小消息下延迟低Tree 在大消息下带宽高。我实测下来跨节点 all-to-all 用 Tree Simple 的组合带宽能比默认配置高 20% 左右。4.2 显存与计算的重叠策略TP MoE 的显存压力主要来自三块专家权重、激活值、通信缓冲区。专家权重是固定的激活值随 batch size 增长通信缓冲区则是 all-to-all 的临时空间。这三块要平衡好否则很容易 OOM。一个实用的策略是把专家权重放在显存里常驻激活值用 gradient checkpointing 或者 activation offloading 来省通信缓冲区则尽量复用。我试过把 all-to-all 的输入输出缓冲区做成一个环形队列通信完的 buffer 立刻回收给下一轮用这样能把缓冲区占用降低一半。计算和通信的重叠核心是让 all-to-all 和 GEMM 在不同的 CUDA stream 上跑。PyTorch 的torch.cuda.Stream可以做到这一点但要注意同步点all-to-all 的结果必须等通信完成才能用于 GEMMGEMM 的结果也必须等计算完成才能发回去。我一般用 event 来做细粒度的同步而不是全局的torch.cuda.synchronize()后者会把所有重叠都打断。4.3 不同规模下的参数调优对照TP MoE 的参数没有一套放之四海而皆准的配置得根据模型规模、卡数、拓扑来调。我整理了一个对照表是我在实际项目中总结出来的经验值模型规模卡数tp_size专家数top-kcapacity factor预期吞吐提升10B84821.251.5x30B1641621.51.8x70B3283212.02.2x130B6486412.52.5x这个表里的“预期吞吐提升”是相对于同等激活参数量的 Dense 模型而言的。可以看到规模越大TP MoE 的优势越明显因为 Dense 模型在规模大了之后通信和显存的开销会急剧上升而 MoE 的稀疏激活能把计算量控制住。但要注意capacity factor 不能设得太大否则显存会爆也不能太小否则 token 被丢弃太多效果会掉。我一般从 1.25 开始试观察 token 丢弃率和效果的变化找到一个平衡点。4.4 一个真实的调优案例从 1.2x 到 2.0x最后分享一个我实际调过的案例。一开始我们的 TP MoE 推理只比 Dense 快了 1.2x远低于预期。抓了 timeline 之后发现all-to-all 占了 55% 的时间而且通信和计算完全没有重叠。第一步优化是拆分 all-to-all把一次大的 all-to-all 拆成 4 个 chunk每个 chunk 通信完立刻启动对应的专家计算。这一步把通信占比降到了 40%吞吐提升到 1.5x。第二步是调整专家组的划分原来是把 8 个专家分到 8 张卡上每个专家一张卡改成每 2 张卡负责一个专家tp_size2这样专家计算可以并行all-to-all 的粒度也更细。这一步把吞吐提升到 1.8x。第三步是 NCCL 调参把NCCL_ALGO从默认改成 TreeNCCL_PROTO改成 Simple并且把NCCL_BUFFSIZE从 4MB 调到 16MB。这一步把吞吐提升到 2.0x。整个过程花了大概两周核心就是不断抓 timeline、找瓶颈、做实验。TP MoE 的调优没有银弹就是个体力活加经验活。但一旦调好收益是非常可观的——同样的硬件能跑更大的模型或者同样的模型能省一半的卡。