ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

【AI Infra小白入门⑩】集合通信基础:AllReduce、Ring算法与NCCL全解析

【AI Infra小白入门⑩】集合通信基础:AllReduce、Ring算法与NCCL全解析

【AI Infra小白入门⑩】集合通信基础:AllReduce、Ring算法与NCCL全解析

本文是「AI Infra 小白入门」系列的第十篇。分布式训练/推理的本质是"计算 + 通信"。集合通信就是多卡之间"传话"的方式,是理解DDP、ZeRO、TP通信开销的钥匙。本文用最通俗的语言,讲清楚通信原语、Ring/Tree算法、NCCL,以及NVLink和IB的带宽差异。


一、写在前面:为什么需要集合通信?

想象一下:你有8张GPU,要训练一个大模型。每张卡上有一部分数据,各自算各自的梯度。

但问题来了:每张卡算的梯度不一样,怎么让所有卡都拿到"平均后的梯度",来同步更新参数呢?

这就需要"集合通信"——多张GPU之间,按照某种模式,交换数据、汇总数据。

集合通信是分布式训练的基础。不懂集合通信,你就理解不了:

  • DDP是怎么同步梯度的
  • ZeRO为什么能省显存、为什么通信量变大了
  • 张量并行的通信开销到底有多大
  • 为什么3D并行要"TP节点内、DP跨节点"

先上一张全景图:


二、六大通信原语:都是什么意思?

通信原语(Collective Primitives)就是"通信的基本模式"。最常用的有六种:

原语语义大白话典型用途
Broadcast1 → N 广播一个人说话,所有人都听到参数初始化同步
ReduceN → 1 规约所有人把数据给一个人,这个人汇总老版梯度同步(只在主卡上有结果)
AllReduceN → N 全部得到规约结果所有人都把数据汇总,每个人都有最终结果DDP梯度平均(最常用!)
AllGather各自片段 → 全部得到完整数据每个人把自己的那片数据发给所有人,最后每个人都有完整的ZeRO-3/FSDP参数收集、TP输出拼接
ReduceScatter规约 + 结果分片先汇总,再把结果切成片,每个人拿一片ZeRO-2/3梯度分片
AllToAll两两交换每个人给每个人发不同的数据MoE专家路由、序列并行

最重要的一个关系

记住这个等式,后面经常用到:

AllReduce = ReduceScatter + AllGather

什么意思?

  • 你可以先做ReduceScatter(规约+分片),再做AllGather(收集)
  • 合起来的效果,和直接做AllReduce是一样的
  • Ring AllReduce就是这么实现的!

三、Ring AllReduce:带宽最优的算法

AllReduce是最常用的集合通信原语,而Ring算法是实现AllReduce的经典方式。

3.1 基本思想

把N个rank排成一个环(Ring),数据切成N块,然后分两个阶段:

阶段一:Scatter-Reduce(N-1步)

  • 每个rank把自己的一块数据传给下一个rank
  • 同时接收上一个rank传来的数据,和自己对应块做规约(相加)
  • 这样传N-1轮之后,每个rank手上都有一块"完整规约后"的数据

阶段二:AllGather(N-1步)

  • 每个rank把自己手上那块"完整结果"传给下一个rank
  • 传N-1轮之后,每个rank都收集齐了所有N块,也就是完整的结果

整个过程,每个rank要传多少次?2(N-1)次。

3.2 关键结论:通信量与rank数无关

每个rank的总通信量是多少?

每个rank通信量 = 2 × (N-1)/N × S ≈ 2S
(S是总数据量,N很大时,约等于2S)

等等,这意味着什么?不管有多少个rank,每个rank的通信量大约都是数据量的2倍!

8卡和128卡,每个卡传的数据量差不多?听起来反直觉,但这就是Ring AllReduce的厉害之处——带宽利用率最优。

3.3 优缺点

优点缺点
带宽利用率高,大数据量下效率高延迟随rank数线性增长(要传2(N-1)步)
实现简单,逻辑清晰rank数多的时候,延迟会比较高

适用场景:大数据量的通信——比如梯度同步,数据量大,带宽是瓶颈,延迟相对不重要。


四、Tree AllReduce:延迟最优的算法

Ring算法带宽最优,但延迟高。那有没有延迟低的算法?有,就是Tree算法。

4.1 基本思想

Tree算法(也叫Recursive Doubling / Halving)用"树状"的方式做规约:

  1. 第1步:rank0和rank1规约,rank2和rank3规约……
  2. 第2步:rank0/1的结果和rank2/3的结果规约……
  3. ……
  4. log₂N步之后,就完成了

因为是"两两合并",所以只需要log₂N步,比Ring的2(N-1)步少多了。

4.2 优缺点

优点缺点
延迟低(logN步)小rank数时,带宽利用率不如Ring
适合小消息/延迟敏感场景大数据量下效率不如Ring

适用场景:小数据量、延迟敏感——比如barrier同步、小tensor的通信。

4.3 怎么选?

实际应用中,NCCL会自动选择:

  • 消息大 → 用Ring
  • 消息小 → 用Tree

不用你操心,NCCL都帮你优化好了。


五、NCCL:GPU集合通信的事实标准

5.1 NCCL是什么?

NCCL(NVIDIA Collective Communications Library)是NVIDIA官方的GPU集合通信库,也是现在的事实标准。

特点:

  • 拓扑感知:自动发现NVLink、PCIe、IB的拓扑结构,选择最优路径
  • 多流并发:多个通信可以同时进行
  • 与CUDA Stream集成:通信和计算可以在不同stream上重叠执行

PyTorch的DDP/FSDP,backend="nccl"指的就是它。

5.2 调试工具

  • 查看拓扑和通道:设置环境变量NCCL_DEBUG=INFO,NCCL会打印详细的拓扑发现和通信通道信息
  • 测试带宽:用nccl-tests里的all_reduce_perf,可以测不同消息大小下的AllReduce带宽

💡调试分布式性能问题时,第一步就是看NCCL的日志——确认拓扑是不是对的、带宽是不是正常。


六、NVLink vs IB:带宽差一个数量级

集合通信的性能,很大程度上取决于硬件互联的带宽。我们来对比一下:

维度节点内 NVLink(A100)跨节点 InfiniBand HDR
带宽600 GB/s(每GPU)~25 GB/s(200Gbps)
延迟~微秒级~几微秒
拓扑NVSwitch全互联交换机组网
适合的通信TP(每层AllReduce/AllGather)DP(每步一次梯度AllReduce)、PP(stage间点对点)

带宽差了多少?600 ÷ 25 = 24倍!差了一个数量级还多。

对并行策略的影响

这个带宽差异,直接决定了分布式训练的"部署铁律":

TP做在节点内(NVLink域),DP/PP跨节点(IB域)

为什么?

  • TP(张量并行):每层都要做AllReduce/AllGather,通信非常频繁,数据量也不小——必须用NVLink,走IB根本扛不住
  • DP(数据并行):每步只通信一次梯度,虽然数据量大,但频率低——IB能承受,但也是瓶颈之一(所以才有梯度bucket化、计算通信重叠这些优化)
  • PP(流水线并行):阶段之间只传激活,通信量相对小——IB可以扛

量化感受一下:

  • 8卡TP,每层通信几MB,频率极高 → 必须NVLink
  • DP梯度同步,每步一次几百MB → IB能承受,但已经是瓶颈之一
  • 这就是为什么要做"梯度bucket化 + 计算通信重叠"——把通信藏在计算后面,让两者并行

七、一张图总结:通信与并行的关系

最后,我们把通信原语和并行策略对应起来:

并行策略用到的主要通信原语通信频率通信量适合的互联
数据并行(DP/DDP)AllReduce每步1次大(整个梯度)IB/NVLink均可
ZeRO-1Reduce + Broadcast每步1次大IB/NVLink均可
ZeRO-2ReduceScatter + AllGather每步1次大IB/NVLink均可
ZeRO-3ReduceScatter + AllGather(参数)每层多次很大最好NVLink
张量并行(TP)AllReduce / AllGather每层多次中到大必须NVLink
流水线并行(PP)Send/Recv(点对点)每步多次中IB可承受
序列并行(SP)AllGather + ReduceScatter每层多次小NVLink最好

八、总结一下

今天这篇文章,我们讲了集合通信的基础知识:

  1. 六大通信原语:Broadcast、Reduce、AllReduce、AllGather、ReduceScatter、AllToAll。记住:AllReduce = ReduceScatter + AllGather
  2. Ring算法:带宽最优(每个rank约传2S数据,与N无关),但延迟随N线性增长。适合大数据量
  3. Tree算法:延迟最优(logN步),但带宽利用率略低。适合小数据量
  4. NCCL:GPU集合通信的事实标准,拓扑感知、自动选算法
  5. NVLink vs IB:带宽差一个数量级,决定了"TP节点内、DP/PP跨节点"的部署铁律

下一篇预告:我们来聊聊AI Infra工程师的"工具箱"——性能分析工具链。gdb、perf、nsys、ncu……这些工具分别是干什么的?什么时候用?


📚本系列导航:

  • 第1篇:一文搞懂什么是AI基础设施:五大方向与不可能三角
  • 第2篇:大厂AI Infra组织架构大揭秘
  • 第3篇:AI Infra岗位路线图:级别、薪资与成长路径
  • 第4篇:AI Infra学习路线图:从入门到精通的完整指南
  • 第5篇:AI Infra必备编程基础:Python、C++与Linux速通指南
  • 第6篇:AI Infra数学基础速成:线代、概率与微积分(工程向)
  • 第7篇:Transformer架构工程拆解:从Attention到MLA的优化密码
  • 第8篇:PyTorch框架Infra实战指南:Tensor、autograd与性能调试
  • 第9篇:GPU硬件概论:从SM到Memory Wall,搞懂优化的物理基础
  • 第10篇:集合通信基础:AllReduce、Ring算法与NCCL全解析(本文)
  • 第11篇:AI Infra工具链指南:性能分析与调试工具一网打尽
  • 更多内容持续更新中……
RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进