【AI大模型大厂面试真题】腾讯大模型面试：MoE训练用TP还是EP？

前言

最近面试中被问到：为什么在 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 吗？

1、通信开销对比

（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。

2、计算开销对比

（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 具有更大优势，它的显存占用更少且更稳定。

3、总结

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 学习资料已经上传优快云，朋友们如果需要可以微信扫描下方优快云官方认证二维码免费领取【保证100%免费】

大模型知识脑图

为了成为更好的 AI大模型 开发者，这里为大家提供了总的路线图。它的用处就在于，你可以按照上面的知识点去找对应的学习资源，保证自己学得较为全面。

经典书籍阅读

阅读AI大模型经典书籍可以帮助读者提高技术水平，开拓视野，掌握核心技术，提高解决问题的能力，同时也可以借鉴他人的经验。对于想要深入学习AI大模型开发的读者来说，阅读经典书籍是非常有必要的。

实战案例

光学理论是没用的，要学会跟着一起敲，要动手实操，才能将自己的所学运用到实际当中去，这时候可以搞点实战案例来学习。

面试资料

我们学习AI大模型必然是想找到高薪的工作，下面这些面试题都是总结当前最新、最热、最高频的面试题，并且每道题都有详细的答案，面试前刷完这套面试题资料，小小offer，不在话下

640套AI大模型报告合集

这套包含640份报告的合集，涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师，还是对AI大模型感兴趣的爱好者，这套报告合集都将为您提供宝贵的信息和启示。

这份完整版的大模型 AI 学习资料已经上传优快云，朋友们如果需要可以微信扫描下方优快云官方认证二维码免费领取【保证100%免费】