最近面试中被问到: 为什么在 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 比重较高的模型中可能更有优势。
如何系统学习掌握AI大模型?
AI大模型作为人工智能领域的重要技术突破,正成为推动各行各业创新和转型的关键力量。抓住AI大模型的风口,掌握AI大模型的知识和技能将变得越来越重要。
学习AI大模型是一个系统的过程,需要从基础开始,逐步深入到更高级的技术。
这里给大家精心整理了一份
全面的AI大模型学习资源
,包括:AI大模型全套学习路线图(从入门到实战)、精品AI大模型学习书籍手册、视频教程、实战学习、面试题等,资料免费分享
!
1. 成长路线图&学习规划
要学习一门新的技术,作为新手一定要先学习成长路线图,方向不对,努力白费。
这里,我们为新手和想要进一步提升的专业人士准备了一份详细的学习成长路线图和规划。可以说是最科学最系统的学习成长路线。
2. 大模型经典PDF书籍
书籍和学习文档资料是学习大模型过程中必不可少的,我们精选了一系列深入探讨大模型技术的书籍和学习文档,它们由领域内的顶尖专家撰写,内容全面、深入、详尽,为你学习大模型提供坚实的理论基础。(书籍含电子版PDF)
3. 大模型视频教程
对于很多自学或者没有基础的同学来说,书籍这些纯文字类的学习教材会觉得比较晦涩难以理解,因此,我们提供了丰富的大模型视频教程,以动态、形象的方式展示技术概念,帮助你更快、更轻松地掌握核心知识。
4. 大模型行业报告
行业分析主要包括对不同行业的现状、趋势、问题、机会等进行系统地调研和评估,以了解哪些行业更适合引入大模型的技术和应用,以及在哪些方面可以发挥大模型的优势。
5. 大模型项目实战
学以致用 ,当你的理论知识积累到一定程度,就需要通过项目实战,在实际操作中检验和巩固你所学到的知识,同时为你找工作和职业发展打下坚实的基础。
6. 大模型面试题
面试不仅是技术的较量,更需要充分的准备。
在你已经掌握了大模型技术之后,就需要开始准备面试,我们将提供精心整理的大模型面试题库,涵盖当前面试中可能遇到的各种技术问题,让你在面试中游刃有余。
全套的AI大模型学习资源已经整理打包,有需要的小伙伴可以
微信扫描下方优快云官方认证二维码
,免费领取【保证100%免费
】