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

前言最近面试中被问到为什么在 MoE 训练中使用 Expert ParallelismEP而不是 Tensor ParallelismTP我的回答是使用 EP 不会减少数据并行DP的数量因为每个 EP 处理不同的数据。而且EP 和 TP 对通信的要求都很高一般不会让 EP 和 TP 跨机。根据我们的实验结果EP 的吞吐量比 TP 更高。当 EP 开启到 8 时我们就不再使用 TP。面试结束后对这个问题进行了更深入的思考觉得还有一些未考虑的细节值得分析。翻了下DeepSeek的技术报告。在 v1 中他们使用了 PP、EP、TP 和 Zero1而在 v2236B 参数、21B 激活中配置为 8EP 16PPzero bubble Zero1没有使用 TP。对于这个参数和激活规模8EP 8PP Zero1 应该就足够了。不知道为什么用了 16PP是因为真的能实现 zero bubble 吗1、通信开销对比1EP通信分析Expert Parallelism 的逻辑如下图所示每个 EP rank 上只包含一部分 expert而每个 EP rank 上的 token即 token 对应的 hidden state会根据 gating 结果分发到其他 EP rank 上的 expert。这个过程通过all-to-all通信完成。2All-to-All Dispatch 逻辑以 4 个 expert、2 个 EP rank 和 topk2 为例 下图中每个 EP rank 上有 3 个 tokenEP rank 0 的 token 分配如下Token 1 → Expert 1 和 Expert 2Token 2 → Expert 1 和 Expert 2Token 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 一致所以总的通信量为4TP 通信分析在 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 为 6EP 的通信量要显著高于 TP。2、计算开销对比1Expert 计算对于 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-L435Megatron-LM fork 了这个库在此基础上支持了 multi stream带来了一定加速。这种场景很适合 multi stream因为每个 expert 的 gemm 都是相互独立的。2DP 数量开 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 固定时完成相同数量的数据需要更少的 GASgradient 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。通信开销在 topk2 且 token 分配均匀的情况下EP 和 TP 的通信量相近。但在topk2或分配不均匀的情况下EP 的通信开销高于 TP。显存占用TP 的显存占用更低且更稳定适合长序列训练或显存敏感的场景而 EP 在不均衡分配时可能引发显存溢出问题。数据并行性EP 不影响数据并行的规模可以在固定的资源下处理更多的数据。而 TP 则会减少数据并行的数量可能导致迭代效率降低。模型规模和架构但 TP 在 attention 比重较高的模型中可能更有优势。最后的最后感谢你们的阅读和喜欢作为一位在一线互联网行业奋斗多年的老兵我深知在这个瞬息万变的技术领域中持续学习和进步的重要性。为了帮助更多热爱技术、渴望成长的朋友我特别整理了一份涵盖大模型领域的宝贵资料集。这些资料不仅是我多年积累的心血结晶也是我在行业一线实战经验的总结。这些学习资料不仅深入浅出而且非常实用让大家系统而高效地掌握AI大模型的各个知识点。如果你愿意花时间沉下心来学习相信它们一定能为你提供实质性的帮助。这份完整版的大模型 AI 学习资料已经上传CSDN朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】大模型知识脑图为了成为更好的 AI大模型 开发者这里为大家提供了总的路线图。它的用处就在于你可以按照上面的知识点去找对应的学习资源保证自己学得较为全面。经典书籍阅读阅读AI大模型经典书籍可以帮助读者提高技术水平开拓视野掌握核心技术提高解决问题的能力同时也可以借鉴他人的经验。对于想要深入学习AI大模型开发的读者来说阅读经典书籍是非常有必要的。实战案例光学理论是没用的要学会跟着一起敲要动手实操才能将自己的所学运用到实际当中去这时候可以搞点实战案例来学习。面试资料我们学习AI大模型必然是想找到高薪的工作下面这些面试题都是总结当前最新、最热、最高频的面试题并且每道题都有详细的答案面试前刷完这套面试题资料小小offer不在话下640套AI大模型报告合集这套包含640份报告的合集涵盖了AI大模型的理论研究、技术实现、行业应用等多个方面。无论您是科研人员、工程师还是对AI大模型感兴趣的爱好者这套报告合集都将为您提供宝贵的信息和启示。这份完整版的大模型 AI 学习资料已经上传CSDN朋友们如果需要可以微信扫描下方CSDN官方认证二维码免费领取【保证100%免费】