PyTorch 2.15 先把 NCCL 那笔债还了

1 阅读11分钟

一台 8 卡机器,训练跑到第 11 个小时,loss 曲线停了。进程还在,GPU 利用率掉到 0,日志没有任何报错。

这种现场最难查。崩溃至少留下栈,挂住什么也不留。而你已经在里面烧了 11 个小时的卡时。

PyTorch 2.15 里,这件事的处理方式变了。10 月 28 日发布的这个版本,把重写过的 NCCL 后端设为所有 nccl process group 的默认实现。它在 2.14 里还是需要手动打开的选项。同一版还带了一个「飞行记录仪」,collective 挂住或超时的时候自动把现场写下来。

核心开发者、Meta 分布式计算方向的工程师 Tristan Rice 在周四的直播里说得很直白:旧的那套 NCCL 接口「欠了太多技术债」,必须重写。

旧接口的问题:一个组里能长出几十个 communicator

要理解这次改动,得先知道 NCCL 在 PyTorch 里怎么用。

NCCL 是英伟达的集合通信库(读作「nickel」),2015 年在 SC15 上发布,2016 年 Facebook 做 PyTorch 时就用上了。它做的事是让多张 GPU 之间高效交换张量、梯度、激活值、权重。用标准 TCP 栈干这个,在大集群上会撞上带宽和 CPU 开销的墙。

这些交换操作叫 collective(集合操作),语义是「所有参与者都必须到场」。一个 all-reduce 里少了任何一张卡,剩下的人就一直等。

这些卡在框架里的组织单元叫 process group。旧的 ProcessGroupNCCL 允许在一个 process group 内部生成几十个 communicator。communicator 是 NCCL 层面的一条实际连接拓扑,ring 或者 tree。

多 communicator 本身不是错,问题是它们互相不知道对方的存在。两条 collective 序列在不同的 communicator 上交错执行,顺序一旦对不上,就死锁。表现形式是间歇性的:同样的代码,跑一百次挂一次,挂在哪一次看不出来。

新后端立了一条硬规则:初始化时,一个 process group 只允许一个 communicator。

这条规则把 crosstalk 从设计里删掉了。代价是并行度要重新组织:需要更多并行路径的程序,得显式地开更多 process group,而不是在一个组里偷偷多开几条连接。

Rice 的说法是,清干净的接口给后续功能留了「一张白纸」。

飞行记录仪:给挂住这件事留证

第二个改动更实用。

以前一个 GPU 集群卡住,用户拿到的是一句通用错误,最多告诉你「某个 collective 超时」。到底是哪张卡、卡在哪个 op、当时各张卡的进度差多少,全靠猜。

飞行记录仪在 collective 挂住或超时的时候,自动把诊断状态写成一系列快照,事后可以逐帧看。它的设计与后端无关:不绑死 NCCL,意味着跑 ROCm 或者混合硬件的团队也能用。

这一点针对的正是分布式训练里最难的一类故障:hang 不产生错误,只产生沉默。 你看到的是 GPU 利用率 0、进程活着、日志没动。传统的日志排查在这种场景下几乎无效,因为你没有「出错的那一行」。

有了快照,问题从「猜」变成「读」:最后完成的 op 是哪个、哪张卡的序列落后、哪条连接没建立。这三样东西基本能定位到具体的那张卡或那条链路。

挂住的时候,硬件在干什么

把「等一个永远不来的参与者」这句话展开,会看到更具体的画面。

all-reduce 的常见实现是 ring:N 张卡排成一个环,每张卡把自己的数据切成 N 份,沿着环传 N−1 步,同时累加。到第 N−1 步结束,每张卡手上都是完整的结果。这个设计的好处是带宽利用均衡,坏处是环上任何一环断了,整条链就停在那里等。

一张卡挂了,它的上游一直在往它发数据,发不动就阻塞在 send;它的下游一直在等它,等不到就阻塞在 recv。整个环上所有卡的状态都变成「等」:GPU 利用率掉到 0,进程活着,日志不动。

框架层还有一层超时:collective 卡太久会触发 watchdog,默认时间在十分钟这个量级。但超时只告诉你「某个 op 超时了」,不告诉你是谁、卡在哪一步。以前你能拿到的信息就这么多。

这也解释了「飞行记录仪」为什么是这个版本最实用的部分。它记下来的东西里没有错误,只有时间线上每一张卡的最后动作。而 hang 的现场恰恰只有时间线。

新旧对比

维度旧 ProcessGroupNCCL2.15 的新后端
每组 communicator 数可以几十个初始化时强制 1 个
死锁来源多 communicator 交错设计上消除 crosstalk
挂住诊断一句通用超时自动快照,可逐帧看
straggler只能整体重启可被编排器剔除并替换
重配置速度慢快路径需 NCCL 2.29+ / RCCL 2.30.7+
退路—2.15 可回滚,2.16 移除

straggler:把慢卡踢出去

第三个改动是弹性。一张卡崩了,或者跑得比别人慢,现在可以让第三方编排器把这张卡从组里排除,再补一个新的 worker 进来。

这条依赖外部配合。PyTorch 自己不负责调度,得由 torchrun 那类编排层去执行剔除和补充。所以它落到接口层面:后端提供「可重配置」的能力,编排器负责决策。

重配置的速度也有要求。官方给的版本门槛是 NCCL 2.29 以上,ROCm 那边是 RCCL 2.30.7。低于这个版本的,快速重配置路径用不了。

为什么一张坏卡能毁掉整个训练

这里的账值得算一下。

大模型训练是同步的:每一步,所有卡算完自己的那部分梯度,然后 all-reduce 求平均,再各自更新。这是一个 barrier。任何一张卡慢 20%,整个集群的每一步都要等它 20%;任何一张卡挂了,整个 job 停在那里等。原因不在计算,而在 collective 在等一个永远不来的参与者。

规模越大,这个概率越不可避免地往上走。一千张卡跑一个月,单卡故障率哪怕只有千分之一每天,也意味着平均每天都有一张要坏。

过去的处理方式是 checkpoint + 重启:把最近的权重存下来,杀掉 job,重新拉起,从 checkpoint 恢复。代价是 checkpoint 之间的所有进度作废,加上重新初始化的时间(加载权重、重建通信、预热)。在千卡规模上,这一套通常以十分钟计。

新方案的方向是把「一张卡坏了」从 job 级事件降级为 worker 级事件:剔掉它、补一个、继续跑。省下的是那十分钟乘以故障频率。

这也是为什么这次发布的定位落在可用性这一侧,不在性能那一侧。PyTorch 从研究工具往生产基础设施走,这类改动比多加几个算子重要得多。

这笔债是怎么欠出来的

NCCL 是一个 C 库,有自己的异步模型:collective 调用发出去立刻返回,真正的完成要等 stream 上的事件。PyTorch 要在 Python 和 C++ 两层之间包住这套语义,早期实现走的是「每次调用按需建一个 communicator」的路子。

按需建,在功能上没问题,跑起来也对。问题在规模上来之后才显形:一个进程组里同时活着几十个 communicator,每个都有自己的连接集合和完成的时序,任何一个时序上的错位都会变成死锁。而这些错位是概率事件,本地八卡跑不出来,集群上偶尔出现一次。

Rice 用的词是「技术债」。这类债的典型特征就是:能跑、能修、每次修都得在别人的代码里找原因。重写的价值不在新功能,在把「偶发」从系统里删掉。

那十分钟值多少

把 checkpoint 重启这套流程摊开,账是这么算的。

一次故障,先等 watchdog 超时(十分钟量级),确认没救,杀掉全部进程,重新拉起,从最近的 checkpoint 加载权重,重建通信,回到出故障前那一步。整个流程在千卡规模上通常是十几分钟起,慢的话半小时。

关键在「最近的 checkpoint」这四个字。checkpoint 不是每步都存,存一次要写几个 TB 到几十 TB(取决于模型大小和优化器状态),写盘本身要占时间。所以实际间隔往往在几十分钟到几小时之间。

折中下来,一次单卡故障的期望损失,是「重启时间 + 半个 checkpoint 间隔」。在千卡集群上,这套成本乘以「大致每天一次」的故障频率,一年就是几百小时的整集群卡时。

这也是为什么 straggler 剔除这件事听起来小、实际收益大。它把损失从「半个间隔」压到「补一张卡的时间」。

顺带说一句,checkpoint 的间隔设计本身就是个权衡:存得勤,训练被写盘打断得多;存得稀,故障损失大。这套权衡在弹性训练成熟之后会变——如果补一张卡只要几分钟,就没必要为了「少丢进度」而高频写盘。

迁移要注意的三件事

第一,NCCL 版本。 NCCL 2.29+ / RCCL 2.30.7+ 是硬门槛。很多团队的系统里,NCCL 是跟着 CUDA 镜像或者驱动一起锁死的,升 PyTorch 可能要连驱动一起动。这一点在升级前就该查,别等到跑起来才发现快速重配置不可用。

第二,旧后端的退路只有一个版本。 2.15 里还留着旧的 ProcessGroupNCCL 作为回滚选项,官方计划在 2.16(明年初)删掉它。也就是说你有一个版本周期来测新后端、把兼容问题解决掉。留到 2.16 再动,就没有退路了。

第三,一个 process group 一个 communicator 会改变性能特征。 原来靠「一个组里多开几条连接」拿到并行度的地方,现在会退化。这类代码要改成「多开 process group」。改法本身不复杂,但涉及通信拓扑的重新设计,不能只改一行。

这一版还有什么

除了 NCCL 后端,2.15 还带了几样:

  • 一个新的 Python transport API,用于异步的单向数据搬运,典型场景是 trainer 和 generator 之间同步权重;
  • Apple 硬件上的并行处理改进;
  • 对英伟达下一代 Blackwell GPU 的支持。

整个版本 3,117 个 commit,583 位贡献者。这个量级说明它不是一个小版本。

边界:还不清楚的部分

官方没有公布新后端的性能对比数据。一条 process group 只允许一个 communicator,理论上会减少连接数、降低建立开销,但在需要高并行度的拓扑上也可能带来额外延迟。这笔账在不同模型、不同互联(NVLink 还是 InfiniBand)上结论不同,得自己测。

飞行记录仪的快照格式、单次写入的开销、会不会在超大集群上产生可观的 IO,官方文档还没给出具体数字。一个千卡 job 挂住时,每张卡都写一份快照,这个量不小。

straggler 剔除的实际效果也取决于编排器实现。PyTorch 提供能力,谁来做决策、剔除之后怎么保证数据一致性(被剔除的 worker 已经算了一半的梯度怎么办),这些不在这次发布的范围内。

接下来看什么

两个时间点。

一是 2.16 什么时候发。旧后端被删掉的那一版,会把所有「先不动、等等看」的团队逼到台前。历史上 PyTorch 的 minor 版本间隔大约一个季度,如果按这个节奏,明年初就会到。

二是各大训练框架(Megatron-LM、DeepSpeed、FSDP 上层封装)什么时候跟进这套接口。PyTorch 改了底层,上层的并行策略要重新对齐,这一步通常比 PyTorch 自己发版慢几个月。

在那之前,能做的判断是:如果你的训练集群规模在百卡以上,且被「一张坏卡毁掉整个 job」折磨过,这一版值得排进测试计划。如果你的规模在十几张卡、故障率不高,等 2.16 更划算,反正那时候旧路也没了。


来源: