MOE训练用EP还是TP?
最近面试中被问到: 为什么在 MOE训练中使用_ExpertParallelism(EP)而不是TensorParallelism(TP)?
我的回答是，使用EP 不会减少数据并行(DP)的数量，因为每个EP 处理不同的数据。
而且，EP和 TP对通信的要求都很高，一般不会让EP和 TP跨机。根据我们的实验结果，EP的吞吐量比 TP 更高。当 EP 开启到8时，我们就不再使用 TP。
面试结束后，对这个问题进行了更深入的思考，觉得还有一些未考虑的细节值得分析。
翻了下DeepSeek的技术报告。在v1中，他们使用了PP、EP、TP和Zerol，而在v2(236B 参数、21B激活)中，配置为8EP+16PP(zerobubble)+Zero1，没有使用 TP。
对于这个参数和激活规模，8EP+8PP+Zero1应该就足够了。不知道为什么用了16PP，是因为真的能实现 zero bubble 吗?
1.通讯开销对比
①EP通信分析
Expert Parallelism 的逻辑如下图所示，每个 EPrank上只包含一部分expert，而每个EPrank上的token(即token 对应的 hidden state)会根据 gating 结果分发到其他 EP rank 上的 expert。
这个过程通过 all-to-all 通信完成。
② All-to-All Dispatch 逻辑
以4个expert、2个EPrank和topk=2为例，下图中每个EPrank上有3个token:
EP rank0 的 token 分配如下:
Token 1→Expert 1和 Expert 2
Token 2→Expert1和 Expert 2
Token3->Expert0和 Expert 3
在 all-to-all 通信前，需要对 local token 按照 gating 结果进行permute/group，将发往同-expert的 token 分组。
随后，这些token 通过 all-to-all 通信发送到对应的 expert rank。
③ 通信量计算
假设 local token的形状为[b x s,h]。在分组后,形状变为[b x s x topk,h]。若token 分配完全均匀:
每个 EP rank 发送/接收的 token 数量为:
b x s x topk *(ep_world_size-1)/ep_world_size
近似为b x s x topk
对于 half precision，为通信量近为:
2 x b x s x topk x h bytes
在 local experts 上计算结束后需要发送原本的 ep rank，是一个 all-to-all的逆过程，对应上图中的 all-to-all combine，通信量和 all-to-all dispatch一致，所以总的通信量为:
4 x b x s x topk x h bytes
④TP 通信分析
在Tensor Parallelism中，MLP(或 expert)前向计算需要一次 all-reduce操作。
Z1和 Z2 的形状为 [b x s,h]
对于半 half precision，通信量为:
2 x b x s x h x 2 bytes
最前面的2是由于 ring all-reduce 包含reduce-scatter和 all-gather 两个步骤，它们的通信量相等。
这里通信量的计算也是有近似的，实际上reduce-scatter只需要发送和接收 tp-1次，而不是 tp 次，细节可以参考 OneFlow: 手把手推导 Ring Allreduce 的数学性质。
类似地，Transformer 中的 attention 中的 linear 也会被切分，进一步增加 TP 的通信开销。
对于一个 Transformer 层，TP 的前向通信量为:
2 x 2 x b x s x h x 2 bytes = 8 x b x s x h bytes
对比 EP 和 TP 的通信量，当 topk 等于2 时，通信量一致，也就是 Mixtral8x7B 这种配置，但是这是在 token 分配完全均匀的假设下，真实训练场景中，不可能是均匀的，由于木桶效应，ep的通信延迟会更高。
MOE训练中会出现这样一个现象，随着训练的进行，吞吐会提升，尤其在训练早期，这是由于一开始 token分配非常不均匀，随着训练的进行，分配更加均匀，吞吐趋于稳定。
当 topk大于2 时，EP 的通信量要高于 TP，像deepseekv2做了 expertsegmentation后，topk为6，EP的通信量要显著高于TP。
2.计算开销对比
①Expert计算
对于 EP，完成 All-to-all dispatch后，所有 token 都被分发到了对应目标expert所在的 EPrank，接着执行矩阵乘法运算。
对于 TP，每个 TP rank都包含所有 expert，但每个 expert 的参数只有 1/TP 份。
由于包含所有 expert，无需将 token 发送到其他 rank，可以直接在本地完成计算。
EP 和 TP 在 expert 的 FLOPS 数相同，但 EP 的 expert 计算对硬件更友好。
以上面两图为例，EP执行两个大的矩阵乘法(因为localrank的 expert参数量更大，且从其他rank上收到分配给localexpert的token)，而TP则执行4 个小的矩阵乘法。GPU 在大矩阵乘法上的效率更高。
FLOPS 数并不一定重要，更应该考虑计算对硬件是否友好。
例如 Mambal，尽管它的 FLOPS 数比 attention 少，且可以使用 parallelscan 并行训练。
但由于 parallelscan 只能使用 CUDAcore 而无法利用 tensorcore，其速度反而比能够利用 tensor core 的 full attention 慢。不过，Mamba2 已经解决了这个问题。
除此之外，矩阵乘法的次数也不同。在一个eprank上，矩阵乘法次数等于localexpert的个数(total_experts/ep_world_size)
而在一个 tp rank 上，矩阵乘法次数等于 total expert 的个数。这需要对local expert 进行一次 for loop，执行 local expert 数量次 kernel launch。
比如deepseekv2160个expert，开启EP8，每个eprank负责20个expert的计算，TP8则负责160个expert的计算，恐怖…
总的来说，ep 在 expert计算上比tp 具有显著优势:-次 kernellaunch 有更大的 workload，且 kernellaunch 次数更少。
这里都会使用 grouped gemm 来加速计算，本质也是减少 kernellaunch，只需要一次 launch ，增加一次 kernellaunch 的 workload。
这样缓解了 wave quantization 的问题，感兴趣的可以看看 How To Write ACUDA Program: The Ninja Edition.
对 grouped gemm 感兴趣的可以看看 Writing Grouped GEMMs in TritonNvidia以及 triton 官方 tutorial。
但是实际生产中，megablocks使用了这个库，而这个库并非真正的groupedgemm，仍是通过 for loop 实现。
② DP数量
开 EP 不会影响 DP 数量，这是因为每个 EP rank 处理不同的数据。
相比之下，同一个 TP group 中的所有 TP rank 处理相同的数据，在固定 worldsize 的情况下，开启 TP 会使 DP 变为原来的 1/TP。
举例来说，当worldsize为64时，启用EP8后DP仍为 64，但启用 TP8后DP 就只有 8。
这表明在总卡数相同的情况下，使用 EP 而非 TP 可以在每次 forward 中处理更多数据。
当 global batch size 固定时，完成相同数量的数据需要更少的 GAS(gradientaccumulation step)。
另外的一个间接影响: 在有 pipeline parallelism 的情况下，较大的 DP 会导致micro batch 数减小，从而产生更大的 pipeline bubble。
③ 显存占用
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 的计算效率上具有优势，减少了 kernellaunch 次数，增加了每次launch 的 workload。
通信开销: 在 topk=2 且 token 分配均匀的情况下，EP 和 TP 的通信量相近。
但在topk>2或分配不均匀的情况下，EP的通信开销高于TP。
显存占用: TP 的显存占用更低且更稳定，适合长序列训练或显存敏感的场景;而 EP 在不均衡分配时可能引发显存溢出问题。
数据并行性: EP 不影响数据并行的规模，可以在固定的资源下处理更多的数据。而 TP 则会减少数据并行的数量，可能导致迭代效率降低。
模型规模和架构: 但 TP 在 attention 比重较高的模型中可能更有优势。
