企业官网建设流程全解析

1. 从 Dense GEMM 到 TP MoE 的演进逻辑

1.1 为什么单纯堆 Dense GEMM 已经走不通了

过去几年做大模型训练和推理的人都有一个共同感受：Dense GEMM 的优化空间正在被榨干。所谓 Dense GEMM，就是稠密矩阵乘法，输入矩阵里几乎没有零元素，每个权重都参与计算。Transformer 里的注意力投影、FFN 的两层线性变换，本质上都是 Dense GEMM。早期模型参数量小，一张卡或者几卡就能放下，大家比拼的是怎么把 cuBLAS 的利用率从 60% 拉到 85%，怎么用 Tensor Core 把 FP16 的吞吐打满。

但模型规模一上来，问题就变了。一个 70B 参数的模型，光权重就是 140GB（FP16），单卡根本放不下。这时候你有两个选择：一是把模型切开放到多卡上，二是想办法让每次计算只激活一部分参数。前者是并行策略的问题，后者就是 MoE（Mixture 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) : (r+1) * (intermediate_dim // tp_size)]。如果按行切，则是weight[e][r * (hidden_dim // tp_size) : (r+1) * (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 的选择有个细节：k=1 的时候，每个 token 只走一个专家，计算量最小，但负载均衡最难做；k=2 的时候，每个 token 走两个专家，计算量翻倍，但路由更平滑。Switch Transformer 用的是 k=1，但加了 capacity factor 和 auxiliary loss。我实测下来，k=2 在推理时效果更稳，尤其是当专家数量不多（比如 8 个）的时候，k=1 容易导致某些专家被完全忽略。

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 还是 Tree，NCCL_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 的结果必须等通信完成才能用于 GEMM，GEMM 的结果也必须等计算完成才能发回去。我一般用 event 来做细粒度的同步，而不是全局的torch.cuda.synchronize()，后者会把所有重叠都打断。

4.3 不同规模下的参数调优对照

TP MoE 的参数没有一套放之四海而皆准的配置，得根据模型规模、卡数、拓扑来调。我整理了一个对照表，是我在实际项目中总结出来的经验值：

这个表里的“预期吞吐提升”是相对于同等激活参数量的 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_size=2），这样专家计算可以并行，all-to-all 的粒度也更细。这一步把吞吐提升到 1.8x。

第三步是 NCCL 调参：把NCCL_ALGO从默认改成 Tree，NCCL_PROTO改成 Simple，并且把NCCL_BUFFSIZE从 4MB 调到 16MB。这一步把吞吐提升到 2.0x。

整个过程花了大概两周，核心就是不断抓 timeline、找瓶颈、做实验。TP MoE 的调优没有银弹，就是个体力活加经验活。但一旦调好，收益是非常可观的——同样的硬件，能跑更大的模型，或者同样的模型，能省一半的卡。