企业官网建设流程全解析

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

1.1 为什么单卡 GEMM 已经喂不饱现代模型

做过大模型训练或者推理优化的人都有一个直观感受：模型参数量的增长速度远远快于单卡算力和显存的增长速度。早几年做一个 Transformer 结构的模型，一层里最重的计算就是多头注意力和前馈网络里的两次矩阵乘法，也就是我们常说的 Dense GEMM。所谓 Dense，指的是参与计算的权重矩阵是稠密的，每一个输入 token 都要和整个权重矩阵做乘加运算，计算量随参数量线性增长。

Dense GEMM 在单卡上的优化其实已经非常成熟了。从最朴素的 tiling 分块，到 shared memory 的 bank conflict 规避，再到 tensor core 的 MMA 指令排布，社区里能挖的优化空间基本被挖干净了。但问题在于，当模型从几亿参数涨到几百亿、上千亿参数时，单卡的显存根本放不下完整的权重矩阵，即便放得下，单卡的计算吞吐也撑不住。这时候就出现了两条路：一条是把模型切开放到多卡上，也就是各种并行策略；另一条是改变模型本身的计算结构，让每个 token 只激活一部分参数，这就是 MoE（Mixture of Experts，混合专家）。

我个人的判断是，Dense GEMM 到 MoE 的转变，本质上不是单纯的算子优化问题，而是计算范式的迁移。Dense GEMM 追求的是把一次矩阵乘法做到极致，而 MoE 追求的是在总参数量巨大的前提下，让每次前向只触碰其中一小部分。这个思路的转变，直接决定了后面所有工程实现的重心都不一样了。

1.2 MoE 到底解决了什么问题，又引入了什么新问题

MoE 的核心思想可以用一句话概括：把原来一个大 FFN 拆成 N 个小 FFN（每个叫一个 expert），然后通过一个门控网络（gate / router）为每个 token 选择 top-k 个 expert 进行计算。这样一来，模型的总参数量可以做得非常大，但每个 token 的实际计算量只和 top-k 个 expert 的参数量相关。

举个具体数字感受一下。假设原来一个 FFN 的中间维度是 16384，现在拆成 64 个 expert，每个 expert 中间维度是 2048，top-2 激活。那么总参数量是 64 × 2048 对应的权重，而每个 token 只算 2 × 2048。总参数涨了，单 token 计算量反而可能下降。这就是 MoE 最吸引人的地方：用稀疏激活换更大的模型容量。

但天下没有免费的午餐。MoE 引入的新问题同样棘手：第一，门控网络的路由决策是动态的，每个 token 走的 expert 不一样，导致计算负载极不均衡，有的 expert 被挤爆，有的 expert 闲着；第二，expert 分布在不同设备上时，token 需要跨设备传输，通信开销可能吃掉计算收益；第三，训练时的负载均衡损失、容量因子、dropless 还是 dropping 这些策略，每一个都会影响最终效果。这些问题在纯 Dense GEMM 时代是不存在的，所以从 Dense 迁移到 MoE，工程复杂度是成倍上升的。

1.3 TP MoE 里 TP 和 MoE 是怎么结合的

标题里的 TP MoE，指的是 Tensor Parallelism（张量并行）和 MoE 的结合。这里需要把两个概念分清楚：TP 解决的是单个权重矩阵太大、单卡放不下的问题，做法是把矩阵按行或按列切分到多张卡上，每张卡算一部分，最后通过 all-reduce 或 all-gather 把结果拼起来；MoE 解决的是总参数量太大、但每次只激活一部分的问题，做法是把不同 expert 放到不同卡上，token 根据路由结果被发送到对应的卡。

当这两个东西叠在一起时，就出现了二维甚至三维的切分：一个 expert 内部的权重矩阵可能还要做 TP 切分，同时不同 expert 又分布在不同设备组上。这时候通信模式就变得非常复杂，既有 TP 组内的 all-reduce，又有 expert 之间的 all-to-all。理解 TP MoE 的关键，就是理解这两层通信是怎么嵌套和调度的。我见过不少团队在这一步踩坑，要么是通信和计算没有 overlap 好，要么是 expert 放置策略和 TP 切分策略冲突，导致 GPU 利用率上不去。

2. Dense GEMM 的核心细节与实操要点

2.1 Dense GEMM 的性能瓶颈到底在哪

要理解从 Dense 到 MoE 的迁移，得先把 Dense GEMM 的瓶颈吃透。一个标准的 GEMM 计算 C = A × B，其中 A 是 M×K，B 是 K×N，C 是 M×N。在 GPU 上，这个计算的性能上限由三个因素决定：算力、显存带宽、以及数据搬运的效率。

对于大模型里的 GEMM，M 通常是 batch × sequence length，K 和 N 是隐藏维度或中间维度。当 M 很大时，GEMM 是 compute-bound 的，这时候 tensor core 的利用率是核心指标；当 M 很小时（比如 decode 阶段每次只有一个 token），GEMM 退化成 GEMV，变成 memory-bound，这时候瓶颈在权重矩阵的读取带宽上。这个区别非常关键，因为 MoE 在 decode 阶段的收益和 prefill 阶段完全不同，后面会展开讲。

实操中判断一个 GEMM 是 compute-bound 还是 memory-bound，有个简单的算法：计算 arithmetic intensity，也就是总浮点运算数除以总字节搬运量。以 FP16 为例，如果 intensity 大于 GPU 的算力带宽比（比如 A100 大约是 312 TFLOPS / 1555 GB/s ≈ 200），那就是 compute-bound，否则是 memory-bound。这个估算在做优化方案选型时非常有用，能帮你快速判断该往哪个方向使劲。

2.2 手写一个高效的 Dense GEMM kernel 要注意什么

虽然现在大部分场景直接用 cuBLAS 或 CUTLASS 就够了，但理解手写 kernel 的要点，对后面理解 MoE 的 grouped GEMM 很有帮助。一个高效的 GEMM kernel，核心是分块策略：把大矩阵切成能塞进 shared memory 和寄存器的 tile，让每个 thread block 负责一个输出 tile，通过循环沿 K 维度累加。

具体来说，几个关键参数需要仔细调：BM、BN、BK 分别是输出 tile 的行、列和 K 维分块大小。BM 和 BN 决定了每个 block 的输出规模，BK 决定了每次从 global memory 加载多少数据到 shared memory。这几个值不是随便定的，要结合 shared memory 容量、寄存器数量、以及 tensor core 的 MMA 形状来算。比如用 16×8×16 的 MMA 指令，那 BM 和 BN 最好是 16 和 8 的倍数。

还有一个容易被忽略的点是 double buffering。在沿 K 维循环时，如果加载下一块数据和计算当前块是串行的，那 memory latency 会直接暴露出来。用 double buffering，也就是开两块 shared memory，一块算的时候另一块在加载，能把 latency 藏起来。这个技巧在 MoE 的 grouped GEMM 里同样适用，而且因为 expert 数量多、每个 expert 的矩阵小，double buffering 的收益更明显。

注意：调 tile 参数时不要只盯着单个 kernel 的峰值性能，要看端到端的表现。我见过为了把单个 GEMM 调到理论峰值的 95%，结果 shared memory 占用太高导致 occupancy 下降，整体吞吐反而掉了的情况。

2.3 Dense GEMM 在多卡上的切分方式

当单个权重矩阵放不下一张卡时，就要做切分。最常见的两种是列并行和行并行。列并行是把 B 矩阵按列切成 N 份，每张卡算 C 的一部分列，这种切法不需要在 forward 里做通信，但 backward 需要 all-reduce。行并行是把 B 按行切，每张卡算部分和，forward 需要 all-reduce。

在 Transformer 的 FFN 里，通常第一个线性层用列并行，第二个线性层用行并行，这样中间只需要一次 all-reduce，通信量最小。这个模式在 MoE 里依然成立，只不过每个 expert 的 FFN 都要独立做这个切分。理解这一点，对后面设计 TP MoE 的通信模式至关重要。

实测下来，TP 的通信开销和切分粒度强相关。切得越细，单卡计算量越小，但 all-reduce 的次数和数据量不变，通信占比就越高。一般建议 TP degree 不要超过 8，再大通信收益就不划算了。这个经验值在 MoE 场景下同样适用，甚至要更保守，因为 MoE 本身还有 all-to-all 的通信。

3. MoE 的核心机制与工程实现

3.1 门控网络与路由策略的设计取舍

门控网络是 MoE 的大脑，它决定了每个 token 去哪些 expert。最简单的实现是一个线性层加 softmax，输出每个 expert 的权重，然后取 top-k。但这里有几个设计选择会显著影响效果和工程复杂度。

第一个选择是 top-k 里的 k 取多少。k=1 最省计算，但负载均衡最难做，容易塌缩到少数几个 expert；k=2 是目前主流，兼顾了效果和均衡性；k 再大计算量就上去了，收益递减。第二个选择是门控权重的归一化方式，是在 top-k 内部重新归一化，还是保留原始 softmax 值。前者更稳定，后者保留了更多信息但可能导致数值问题。第三个选择是是否加噪声，训练时给门控 logits 加高斯噪声可以促进探索，避免早期就塌缩。

我个人的经验是，门控网络的学习率要单独调，通常比主干网络小一个量级。因为门控的梯度信号相对弱，学太快容易震荡，学太慢又跟不上主干的收敛节奏。这个细节在很多论文里不会写，但实际训练时非常关键。

3.2 负载均衡：MoE 训练里最头疼的问题

负载均衡做不好，MoE 基本就废了。所谓负载不均衡，就是大部分 token 都涌向少数几个 expert，其他 expert 得不到足够训练，相当于浪费了参数容量。更糟的是，被挤爆的 expert 所在的设备成为瓶颈，整个训练被拖慢。

解决负载均衡有两个层面。算法层面，最常用的是加一个 auxiliary loss，惩罚 expert 使用率的不均衡。这个 loss 的系数需要仔细调，太大影响主任务效果，太小起不到均衡作用。工程层面，可以用 capacity factor 限制每个 expert 最多处理多少 token，超出的 token 直接丢弃（dropping）或者绕过（dropless）。dropping 实现简单但会损失信息，dropless 效果好但需要动态调度，实现复杂。

提示：auxiliary loss 的系数建议从 0.01 开始试，观察 expert 使用率的分布再调整。如果发现某些 expert 使用率长期接近 0，说明系数太小；如果主任务 loss 明显高于同规模 Dense 模型，说明系数太大。

3.3 Expert 的放置策略与 all-to-all 通信

当 expert 分布在多张卡上时，token 需要从它所在的卡被发送到目标 expert 所在的卡，算完再发回来。这个通信模式就是 all-to-all，是 MoE 里最重的通信开销。

Expert 放置策略直接影响 all-to-all 的效率。最简单的是一张卡放一个 expert，但这样 expert 数量受限于卡数。更常见的是一个 expert 跨多张卡做 TP 切分，或者一张卡放多个小 expert。放置策略要和并行策略一起考虑：如果 expert 和 attention 在同一组卡上，可以复用通信组；如果分开，就要额外的通信域。

实测下来，all-to-all 的性能对网络拓扑非常敏感。在 NVLink 全连接的机器上，all-to-all 效率很高；跨节点走 InfiniBand 时，延迟和带宽都会成为瓶颈。所以设计 TP MoE 时，尽量让 all-to-all 发生在节点内，跨节点的通信能少则少。这一点在后面的实操部分会结合具体配置再展开。

4. TP MoE 的实操过程与核心环节

4.1 从 Dense FFN 改造为 TP MoE 的完整步骤

假设我们有一个已经跑通的 Dense Transformer，现在要把它改造成 TP MoE。我把这个过程拆成几个可落地的步骤。

第一步，确定 MoE 的粒度。是每层都换 MoE，还是隔层换？是只换 FFN，还是 attention 也换？主流做法是只换 FFN，且每隔一层换一次，这样能在效果和开销之间取得平衡。确定粒度后，把对应的 FFN 替换成 MoE 层。

第二步，设计 expert 的并行布局。假设有 8 张卡，TP degree 设为 2，那么可以分成 4 个 expert 组，每组 2 张卡做一个 expert 的 TP。这样每个 expert 的权重矩阵被切成 2 份，同时有 4 个 expert 并行。这个布局要在初始化时就确定好，因为涉及进程组的划分。

第三步，实现门控和路由。门控网络本身很小，通常不做切分，每张卡都有一份完整的门控权重。路由时，每张卡先算出本地 token 的目标 expert，然后通过 all-to-all 把 token 发送到对应 expert 所在的卡。

第四步，实现 grouped GEMM。每个 expert 收到的 token 数量不一样，不能直接用普通 GEMM，要用 grouped GEMM，也就是把多个不同大小的矩阵乘法打包成一个 kernel 执行。这一步是性能关键，后面单独讲。

第五步，处理反向传播。MoE 的反向比 forward 更复杂，因为梯度要沿着 all-to-all 的路径回传，同时门控网络也要更新。这里要注意梯度累积和通信的 overlap，否则反向会成为瓶颈。

4.2 Grouped GEMM 的实现要点与参数选择

Grouped GEMM 是 TP MoE 里最核心的算子。它的挑战在于：每个 expert 处理的 token 数不同，矩阵的 M 维度是动态的，而且各个 expert 的矩阵大小差异可能很大。

实现上有两种思路。一种是串行执行，循环每个 expert 调一次 GEMM，简单但 kernel launch 开销大，且小矩阵的利用率低。另一种是融合执行，把所有 expert 的 GEMM 打包成一个 kernel，用 group 的概念让不同 block 处理不同 expert。后者效率高但实现复杂，CUTLASS 里有 grouped GEMM 的支持，可以直接用。

参数选择上，关键是处理好 M 维度不均衡的问题。如果某个 expert 收到的 token 特别多，它的 GEMM 会成为长尾。解决办法一是用 capacity factor 限制每个 expert 的最大 token 数，二是把大 expert 再切分到多张卡上。实测下来，capacity factor 取 1.25 到 1.5 之间比较合适，既能容纳一定的波动，又不会浪费太多计算。

注意：grouped GEMM 的 tile 大小不能照搬普通 GEMM。因为每个 group 的 M 可能很小，如果 tile 的 BM 设得太大，会导致大量 block 空转。建议 BM 设小一点，比如 64 或 32，让更多 block 有活干。

4.3 通信与计算的 overlap 调度

TP MoE 的性能上限，很大程度上取决于通信和计算能不能 overlap。理想情况下，当一张卡在等 all-to-all 的数据时，它应该在做别的计算，比如 attention 或者门控。

实现 overlap 的常用手段是把通信拆成多个 chunk，算完一个 chunk 就发一个，同时开始算下一个。这需要把计算和通信放在不同的 CUDA stream 上，用 event 做同步。PyTorch 里可以用torch.cuda.Stream手动管理，或者用一些框架提供的通信算子。

我踩过的一个坑是：overlap 做过头反而会拖慢。因为通信和计算抢 SM 资源，如果计算本身已经把 SM 占满了，通信再插进来只会互相干扰。所以 overlap 要适度，建议先 profile 一下计算和通信各自占用的资源，再决定 overlap 的比例。一般来说，all-to-all 的通信量在总时间占比 20% 到 40% 之间时，overlap 收益最明显。

4.4 一个具体的配置示例与实测数据

为了让大家有个直观感受，我分享一个实测过的配置。硬件是 8 卡 A100 80G，NVLink 全连接。模型是 24 层 Transformer，隐藏维度 4096，FFN 中间维度 16384。改造方案是每隔一层换 MoE，每层 8 个 expert，top-2 激活，TP degree 设为 2。

改造前，Dense 模型单步训练时间约 320ms，显存占用每卡 62G。改造后，MoE 模型单步训练时间约 410ms，显存占用每卡 71G。看起来慢了，但要注意 MoE 的总参数量是 Dense 的约 3.5 倍。如果按参数量归一化，MoE 的每参数训练效率反而更高。推理阶段，prefill 时 MoE 比同参数量的 Dense 快约 2.3 倍，decode 时因为 batch 小、all-to-all 占比高，优势缩小到 1.4 倍左右。

这个数据说明一个关键点：MoE 的收益和 batch size 强相关。batch 越大，all-to-all 的通信能被更多计算摊薄，收益越明显。所以部署 MoE 时，尽量把 batch 做大，或者用 continuous batching 把多个请求拼起来。

5. 常见问题与排查技巧实录

5.1 Expert 负载不均衡的排查与调优

负载不均衡是最高频的问题。排查的第一步是打印每个 expert 的 token 数分布。如果发现最大和最小的 expert 之间差了几倍，那就是不均衡。

调优的顺序是：先调 auxiliary loss 系数，从 0.01 开始，每次翻倍或减半，观察分布变化；如果调 loss 效果有限，再考虑加 capacity factor 做硬限制；最后才考虑改门控结构，比如加噪声或者换归一化方式。我见过有人一上来就改门控结构，结果引入新问题，反而更难排查。

还有一个隐蔽的坑是：负载不均衡可能不是门控的问题，而是数据的问题。如果训练数据本身分布很偏，某些类型的 token 天然倾向于某些 expert，那再怎么调门控也没用。这时候要考虑在数据层面做均衡，或者给不同数据源加不同的采样权重。

5.2 all-to-all 通信超时的定位方法

all-to-all 超时通常表现为训练卡住或者报 NCCL timeout。定位时先看是偶发还是必现。偶发的话，多半是某张卡的计算慢了，导致其他卡等它，触发超时。这时候要检查是不是有 expert 负载特别重，或者某张卡的温度、功耗有问题。

必现的话，先检查通信组配置是否正确。TP MoE 里通常有多个通信组：TP 组、expert 组、以及全局组。如果某个 token 被路由到了一个不在当前通信组里的 expert，就会死锁。这种问题在改并行配置时特别容易出现，建议每次改配置后先用小规模数据跑通再上大规模。

提示：NCCL 的环境变量对排查很有帮助。设置NCCL_DEBUG=INFO可以看到通信的详细信息，NCCL_ASYNC_ERROR_HANDLING=1能让超时更快暴露而不是一直卡着。

5.3 显存溢出的常见原因与解决

MoE 的显存占用比 Dense 复杂，因为除了权重和激活，还有 all-to-all 的缓冲区、grouped GEMM 的临时空间、以及门控的中间结果。

显存溢出最常见的原因是 all-to-all 的缓冲区开太大。这个缓冲区的大小和 batch size、序列长度、top-k、隐藏维度都相关，估算公式是 batch × seq × top-k × hidden × dtype_size × 2（收发各一份）。如果显存紧张，可以减小 batch 或者用更小的 dtype 存缓冲区。

另一个原因是 grouped GEMM 的 workspace。不同实现的 workspace 大小差异很大，有的会预分配一大块，有的按需分配。如果用的是预分配的实现，要确保 workspace 大小和实际 expert 数量匹配，否则要么浪费要么溢出。

5.4 训练不收敛的排查清单

MoE 训练不收敛的原因很多，我整理了一个排查清单，按优先级排序。

这个清单是我从多次踩坑中总结的，按这个顺序排查，基本能覆盖 90% 的不收敛问题。剩下的 10% 往往是多个因素叠加，需要逐个隔离变量来定位。

6. 一些实操心得与后续扩展方向

6.1 从 Dense 迁移到 MoE 的时机判断

不是所有场景都值得上 MoE。我的判断标准是：如果模型参数量已经大到单卡放不下，且推理时 batch 能做得比较大，那 MoE 的收益就很明显。反之，如果模型本身不大，或者推理场景是低延迟的单请求，那 MoE 的通信开销可能吃掉所有收益，不如老老实实做 Dense 的量化或者蒸馏。

还有一个容易被忽略的点是训练成本。MoE 的训练虽然单步计算量可能没涨多少，但因为通信和调度的开销，实际训练时间往往比同参数量的 Dense 长。所以如果训练资源有限，要慎重评估。

6.2 TP MoE 后续可以怎么优化

从工程角度看，TP MoE 还有不少优化空间。一是通信压缩，all-to-all 传的是激活值，可以用 FP8 或者更激进的量化来减少通信量；二是 expert 的弹性调度，根据负载动态调整每个 expert 的容量，而不是固定 capacity factor；三是和流水线并行结合，把 MoE 层放到不同的 pipeline stage，进一步降低单卡的显存压力。

从算法角度看，门控网络的设计还有很多可探索的。比如用强化学习来学路由策略，或者用可微分的 top-k 来让梯度更好地回传。这些方向目前还在研究中，但已经有了一些有意思的结果。

我个人在实际操作中的体会是，TP MoE 的难点从来不是某一个单点技术，而是各个模块之间的配合。门控、路由、通信、计算，任何一个环节没调好，整体性能就上不去。所以做这类项目，一定要有全局视角，先想清楚数据流和通信模式，再动手写代码。上来就抠 kernel 细节，往往事倍功半。