腾讯大模型面试：MoE训练用TP还是EP？

最新推荐文章于 2025-04-22 14:44:06 发布

程序员笑武

最新推荐文章于 2025-04-22 14:44:06 发布

阅读量1.6k

点赞数 14

文章标签：面试职场和发展人工智能 langchain 自动驾驶音视频语言模型

本文链接：https://blog.csdn.net/m0_59164304/article/details/144751601

版权

最近面试中被问到：为什么在 MoE 训练中使用 Expert Parallelism（EP）而不是 Tensor Parallelism（TP）？

我的回答是，使用 EP 不会减少数据并行（DP）的数量，因为每个 EP 处理不同的数据。

而且，EP 和 TP 对通信的要求都很高，一般不会让 EP 和 TP 跨机。根据我们的实验结果，EP 的吞吐量比 TP 更高。当 EP 开启到 8 时，我们就不再使用 TP。

面试结束后，对这个问题进行了更深入的思考，觉得还有一些未考虑的细节值得分析。

翻了下DeepSeek的技术报告。在 v1 中，他们使用了 PP、EP、TP 和 Zero1，而在 v2（236B 参数、21B 激活）中，配置为 8EP + 16PP（zero bubble）+ Zero1，没有使用 TP。

对于这个参数和激活规模，8EP + 8PP + Zero1 应该就足够了。不知道为什么用了 16PP，是因为真的能实现 zero bubble 吗？

01.通信开销对比

（1）EP通信分析

Expert Parallelism 的逻辑如下图所示，每个 EP rank 上只包含一部分 expert，而每个 EP rank 上的 token（即 token 对应的 hidden state）会根据 gating 结果分发到其他 EP rank 上的 expert。

这个过程通过 all-to-all 通信完成。

（2）All-to-All Dispatch 逻辑

以 4 个 expert、2 个 EP rank 和 topk=2 为例，下图中每个 EP rank 上有 3 个 token：

EP rank 0 的 token 分配如下：

Token 1 → Expert 1 和 Expert 2
Token 2 → Expert 1 和 Expert 2
Token 3 → Expert 0 和 Expert 3

在 all-to-all 通信前，需要对 local token 按照 gating 结果进行permute/group，将发往同一 expert 的 token 分组。随后，这些 token 通过 all-to-all 通信发送到对应的 expert rank。

（3）通信量计算

每个 EP rank 发送/接收的 token 数量为：

对于 half precision，为通信量近为：

在 local experts 上计算结束后需要发送原本的 ep rank，是一个 all-to-all 的逆过程，对应上图中的 all-to-all combine，通信量和 all-to-all dispatch 一致，所以总的通信量为：

（4）TP 通信分析

在 Tensor Parallelism 中，MLP（或 expert）前向计算需要一次 all-reduce 操作。

对于半 half precision，通信量为：

最前面的 2 是由于 ring all-reduce 包含reduce-scatter和 all-gather 两个步骤，它们的通信量相等。

这里通信量的计算也是有近似的，实际上 reduce-scatter 只需要发送和接收 tp-1 次，而不是 tp 次，细节可以参考 OneFlow：手把手推导 Ring All-reduce 的数学性质。

类似地，Transformer 中的 attention 中的 linear 也会被切分，进一步增加 TP 的通信开销。

对于一个 Transformer 层，TP 的前向通信量为：

对比 EP 和 TP 的通信量，当 topk 等于 2 时，通信量一致，也就是 Mixtral 8x7B 这种配置，但是这是在 token 分配完全均匀的假设下，真实训练场景中，不可能是均匀的，由于木桶效应，ep 的通信延迟会更高。

MoE 训练中会出现这样一个现象，随着训练的进行，吞吐会提升，尤其在训练早期，这是由于一开始 token 分配非常不均匀，随着训练的进行，分配更加均匀，吞吐趋于稳定。

当 topk 大于 2 时，EP 的通信量要高于 TP，像deepseek v2做了 expert segmentation 后，topk 为 6，EP 的通信量要显著高于 TP。

02.计算开销对比

（1）Expert 计算

对于 EP，完成 All-to-all dispatch 后，所有 token 都被分发到了对应目标 expert 所在的 EP rank，接着执行矩阵乘法运算。

对于 TP，每个 TP rank 都包含所有 expert，但每个 expert 的参数只有 1/TP 份。

由于包含所有 expert，无需将 token 发送到其他 rank，可以直接在本地完成计算。

EP 和 TP 在 expert 的 FLOPS 数相同，但 EP 的 expert 计算对硬件更友好。

以上面两图为例，EP 执行两个大的矩阵乘法（因为 local rank 的 expert 参数量更大，且从其他 rank 上收到分配给 local expert 的 token），而 TP 则执行 4 个小的矩阵乘法。GPU 在大矩阵乘法上的效率更高。

FLOPS 数并不一定重要，更应该考虑计算对硬件是否友好。

例如 Mamba1，尽管它的 FLOPS 数比 attention 少，且可以使用 parallel scan 并行训练，但由于 parallel scan 只能使用 CUDA core 而无法利用 tensor core，其速度反而比能够利用 tensor core 的 full attention 慢。不过，Mamba2 已经解决了这个问题。

除此之外，矩阵乘法的次数也不同。在一个 ep rank 上，矩阵乘法次数等于 local expert 的个数（total_experts / ep_world_size）。

而在一个 tp rank 上，矩阵乘法次数等于 total expert 的个数。这需要对 local expert 进行一次 for loop，执行 local expert 数量次 kernel launch。

比如 deepseek v2 160 个 expert，开启 EP 8，每个 ep rank 负责 20 个 expert 的计算，TP 8 则负责 160 个 expert 的计算，恐怖…

总的来说，ep 在 expert 计算上比 tp 具有显著优势：一次 kernel launch 有更大的 workload，且 kernel launch 次数更少。

这里都会使用 grouped gemm 来加速计算，本质也是减少 kernel launch，只需要一次 launch ，增加一次 kernel launch 的 workload。

这样缓解了 wave quantization 的问题，感兴趣的可以看看 How To Write A CUDA Program: The Ninja Edition。

对 grouped gemm 感兴趣的可以看看 Writing Grouped GEMMs in Triton Nvidia以及 triton 官方 tutorial。

但是实际生产中，megablocks使用了这个库，而这个库并非真正的 grouped gemm，仍是通过 for loop 实现。

https://github.com/tgale96/grouped_gemm/blob/main/csrc/grouped_gemm.cu#L418-L435

Megatron-LM fork 了这个库在此基础上支持了 multi stream，带来了一定加速。这种场景很适合 multi stream，因为每个 expert 的 gemm 都是相互独立的。

（2）DP 数量

开 EP 不会影响 DP 数量，这是因为每个 EP rank 处理不同的数据。

相比之下，同一个 TP group 中的所有 TP rank 处理相同的数据，在固定 world size 的情况下，开启 TP 会使 DP 变为原来的 1/TP。

举例来说，当 world size 为 64 时，启用 EP 8 后 DP 仍为 64，但启用 TP 8 后 DP 就只有 8。

这表明在总卡数相同的情况下，使用 EP 而非 TP 可以在每次 forward 中处理更多数据。

当 global batch size 固定时，完成相同数量的数据需要更少的 GAS（gradient accumulation step）。

另外的一个间接影响：在有 pipeline parallelism 的情况下，较大的 DP 会导致 micro batch 数减小，从而产生更大的 pipeline bubble。

在计算效率这块来说，EP 比 TP 有显著优势。

（3）显存占用

TP 相比 EP 多切分了 attention 中的 linear 层，但由于 attention 在 MoE 架构中占比较低，这个优势并不显著。

在负载不均衡的情况下，某个 rank 上分配的 token 可能过多，导致显存使用激增，甚至出现 OOM。

这种情况下，micro batch 中的 token 数量越多，不均衡分配带来的显存压力就越大。

当 micro batch size 或 sequence length 增加时，单个 micro batch 中的 token 数也会相应增加。因此在长文本训练中，如果 EP 出现显存溢出，可以考虑使用 TP。

因此从显存角度看，TP 具有更大优势，它的显存占用更少且更稳定。

总结

EP 和 TP 各有优劣，其选择取决于具体的训练场景和需求：

计算效率：EP 在 expert 的计算效率上具有优势，减少了 kernel launch 次数，增加了每次 launch 的 workload。

通信开销：在 topk=2 且 token 分配均匀的情况下，EP 和 TP 的通信量相近。但在topk>2或分配不均匀的情况下，EP 的通信开销高于 TP。

显存占用：TP 的显存占用更低且更稳定，适合长序列训练或显存敏感的场景；而 EP 在不均衡分配时可能引发显存溢出问题。

数据并行性：EP 不影响数据并行的规模，可以在固定的资源下处理更多的数据。而 TP 则会减少数据并行的数量，可能导致迭代效率降低。

模型规模和架构：但 TP 在 attention 比重较高的模型中可能更有优势。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述