最近面试中被问到:为什么在 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 实现。
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 具有更大优势,它的显存占用更少且更稳定。
03
总结
EP 和 TP 各有优劣,其选择取决于具体的训练场景和需求:
计算效率:EP 在 expert 的计算效率上具有优势,减少了 kernel launch 次数,增加了每次 launch 的 workload。
通信开销:在 topk=2 且 token 分配均匀的情况下,EP 和 TP 的通信量相近。但在topk>2或分配不均匀的情况下,EP 的通信开销高于 TP。
显存占用:TP 的显存占用更低且更稳定,适合长序列训练或显存敏感的场景;而 EP 在不均衡分配时可能引发显存溢出问题。
数据并行性:EP 不影响数据并行的规模,可以在固定的资源下处理更多的数据。而 TP 则会减少数据并行的数量,可能导致迭代效率降低。
模型规模和架构:但 TP 在 attention 比重较高的模型中可能更有优势。
如何学习AI大模型?
大模型时代,火爆出圈的LLM大模型让程序员们开始重新评估自己的本领。 “AI会取代那些行业
?”“谁的饭碗又将不保了?
”等问题热议不断。
不如成为「掌握AI工具的技术人
」,毕竟AI时代,谁先尝试,谁就能占得先机!
想正式转到一些新兴的 AI 行业,不仅需要系统的学习AI大模型。同时也要跟已有的技能结合,辅助编程提效,或上手实操应用,增加自己的职场竞争力。
但是LLM相关的内容很多,现在网上的老课程老教材关于LLM又太少。所以现在小白入门就只能靠自学,学习成本和门槛很高
那么针对所有自学遇到困难的同学们,我帮大家系统梳理大模型学习脉络,将这份 LLM大模型资料
分享出来:包括LLM大模型书籍、640套大模型行业报告、LLM大模型学习视频、LLM大模型学习路线、开源大模型学习教程
等, 😝有需要的小伙伴,可以 扫描下方二维码领取🆓↓↓↓
👉[CSDN大礼包🎁:全网最全《LLM大模型入门+进阶学习资源包》免费分享(安全链接,放心点击)]()👈
学习路线
第一阶段: 从大模型系统设计入手,讲解大模型的主要方法;
第二阶段: 在通过大模型提示词工程从Prompts角度入手更好发挥模型的作用;
第三阶段: 大模型平台应用开发借助阿里云PAI平台构建电商领域虚拟试衣系统;
第四阶段: 大模型知识库应用开发以LangChain框架为例,构建物流行业咨询智能问答系统;
第五阶段: 大模型微调开发借助以大健康、新零售、新媒体领域构建适合当前领域大模型;
第六阶段: 以SD多模态大模型为主,搭建了文生图小程序案例;
第七阶段: 以大模型平台应用与开发为主,通过星火大模型,文心大模型等成熟大模型构建大模型行业应用。
👉学会后的收获:👈
• 基于大模型全栈工程实现(前端、后端、产品经理、设计、数据分析等),通过这门课可获得不同能力;
• 能够利用大模型解决相关实际项目需求: 大数据时代,越来越多的企业和机构需要处理海量数据,利用大模型技术可以更好地处理这些数据,提高数据分析和决策的准确性。因此,掌握大模型应用开发技能,可以让程序员更好地应对实际项目需求;
• 基于大模型和企业数据AI应用开发,实现大模型理论、掌握GPU算力、硬件、LangChain开发框架和项目实战技能, 学会Fine-tuning垂直训练大模型(数据准备、数据蒸馏、大模型部署)一站式掌握;
• 能够完成时下热门大模型垂直领域模型训练能力,提高程序员的编码能力: 大模型应用开发需要掌握机器学习算法、深度学习框架等技术,这些技术的掌握可以提高程序员的编码能力和分析能力,让程序员更加熟练地编写高质量的代码。
1.AI大模型学习路线图
2.100套AI大模型商业化落地方案
3.100集大模型视频教程
4.200本大模型PDF书籍
5.LLM面试题合集
6.AI产品经理资源合集
👉获取方式:
😝有需要的小伙伴,可以保存图片到wx扫描二v码免费领取【保证100%免费】🆓