简而言之
今天,我们宣布 Pathways on Google Cloud Platform正式可用(General Availability)。在本文中,你将了解 Pathways——Google 为训练自家大模型(包括 Gemini)而构建的分布式系统。我们将讨论 Pathways 如何通过引入一个客户端、把整个计算描述为一张图来应对协调难题,以及分布式运行时如何放置计算、搬运张量并让加速器保持忙碌。
大多数大模型训练都以单程序多数据(SPMD)程序的方式运行。每块加速器都以锁步方式运行同一程序,并通过 AllReduce 等集合通信操作进行协调。
当你的作业就是在每台设备上运行同一程序时,这种方式效果很好。但当作业不再统一时,它就会吃力,例如各阶段相互等待的流水线,或者采样器、奖励模型和训练器是必须重叠运行的不同程序的强化学习循环。
大规模训练还必须能承受硬件故障,这些故障曾导致 Gemini 2.5 训练每小时中断多次. 在 SPMD 中,一台机器丢失就会让所有进程停止,并使本次运行回退到上一个检查点。在每种情况下,挑战都在于协调。你必须决定什么在哪里运行、以什么顺序运行,张量如何在设备之间移动,以及当某台机器掉线时作业如何继续。
在“Pathways on Google Cloud”系列的本篇文章中,你将了解 Pathways——Google 为训练自家大模型(包括 Gemini)而构建的分布式系统。本文讨论 Pathways 如何通过引入一个客户端、把整个计算描述为一张图来应对协调难题,以及分布式运行时如何放置计算、搬运张量并让加速器保持忙碌。
读完后,你将了解 SPMD 在哪些情况下并非最优、单控制器执行带来了哪些改变,以及 Pathways 如何在整个集群范围内协调整个训练作业。
当扩展变成协调问题
如果你训练 PyTorch 模型,几乎肯定会用到 分布式数据并行(DDP) 或 完全分片数据并行(FSDP). DDP 在每块加速器上运行你脚本的一个副本,并通过 AllReduce 等集合通信操作对梯度求平均。FSDP 将参数、梯度和优化器状态分片到各个工作节点上,以便模型能够放得下。
图 1:使用 DDP 和 FSDP 的一个训练步骤。DDP 和 FSDP 都遵循单程序多数据(SPMD)执行模型。每个进程运行你脚本的同一副本,在整个运行期间独占其加速器,并与其他进程锁步推进,因为它们唯一的协调方式就是 AllReduce 等集合通信操作。在一个紧密连接的加速器组内,这些集合通信操作通过高速互连运行,开销很低。当作业扩展到多个组时,同样的集合通信操作就要跨越数据中心网络,而后者要慢得多。例如,在 TPU v5p上,每块芯片在其组内可获得 1,200 GB/s 的双向芯片间带宽,但到数据中心网络只有 50 Gb/s(吉比特,约 6 GB/s),大约低了两个数量级。
而且由于每一步都会阻塞在其集合通信操作上,这种通信会直接计入步骤耗时。跨越的组越多,每一步中用于传输数据的时间占比就越大,因此吞吐量的增长会慢于你所增加的加速器数量。
只要你的训练作业是一个能填满所给硬件的程序,这种执行模型就很合适。你可以像只有一台设备那样编写模型,其余交给集合通信操作。但把模型装进内存只是问题的一部分。让每块加速器都保持忙碌才是更难的部分,而这正是 SPMD 模型开始显露局限的地方。
以 流水线并行 为例。你将模型拆分成多个块分布到各设备上,每台设备现在对更小的微批次运行不同的代码。激活值向前流动,激活梯度向后流动,当设备 0 处理第一个微批次时,设备 1 和 2 在等待。除了你仍需为其付费的空闲加速器之外,这还会降低利用率和吞吐量,即使每台设备在稳态下都有工作可做。
图 2:三台设备,四个微批次。斜线填充的时隙就是流水线气泡。另一个例子是 LLM 的强化学习后训练。你有一个生成轨迹的采样器、一个为轨迹打分的奖励模型,以及一个更新策略并把新权重发回采样器的训练器。这是三个应当重叠运行的不同程序;否则采样器在训练器运行时空闲,训练器在采样器生成时空闲。
图 3:同一个强化学习循环,顺序运行与重叠运行(示意调度)。在这两种场景中,作业现在都是运行在不同设备上的多个不同程序,它们的依赖关系已无法再通过单一集合通信操作来管理。你现在必须协调哪些设备共享哪些张量、哪些进程在其他进程尚未完成时运行,以及提前多久将下一个缓冲区入队,才能让加速器永远不等待输入。由于每个工作节点只能看到本地设备和数据,没有任何单个进程能独自管理所有这些依赖关系。这会带来若干后果。
第一,由于每个进程在整个运行期间都独占自己的加速器,而且每次集合通信都会阻塞,直到所有设备都完成,因此只要有一台机器故障,集合通信就无法完成,整个作业随之停止,只能从上一个检查点重新启动。你无法在运行中途添加设备,也无法归还设备。
第二,由于每个进程从头到尾都占着自己的加速器,算力无法转移到需要它的地方。在同一个作业内,某个阶段需要的芯片比另一个阶段少,却仍然占着多出来的那些。它们闲置着,你还要继续为它们付费,而作业的其他部分在运行结束前都无法使用它们。在整个团队层面,集群规模上也会发生同样的情况。每个作业在运行期间都独占自己的加速器,因此共享一个集群就意味着要把它拆分成按作业划分的分配,而每个新作业都要在自己的主机上启动自己的进程。
最后,开发过程本身也变得更困难。在单台主机上,你可以快速迭代,因为你可以运行脚本、进入 Python shell、设置断点并调试。一旦作业跨越许多主机,情况就不同了。你需要一个启动器把脚本推送到每台主机,会遇到在单机上从未出现过的协调 bug,而且没有一个可以进入调试的单一进程。当某个环节卡住时,每个工作节点都阻塞在集合通信上等待其他节点,没有任何一个地方能展示整个作业的状态。
所以这里共同的挑战就是协调本身。一旦作业不再是单个程序步调一致地运行,你就必须决定什么在哪里运行、以什么顺序运行,以及张量如何从一个设备传到下一个设备。
描述这一点的一种方式是把它看作一张图:每个块是一个节点,在块之间移动的每个张量是一条边,而设备则是节点被放置的位置。在 SPMD 中,你可以手动逐个进程地把这张图连接起来。但如果由一个客户端来声明这张图,包括每个块做什么、需要什么、在哪里运行,而由一个分布式运行时来处理沿边的传输、依赖关系和顺序,会怎样呢?
图 3:同一个强化学习循环,顺序运行与重叠运行(示意调度)。这就是 Pathways背后的理念,它是 Google 为训练自家最大的模型(包括 Gemini)而构建的分布式系统。Pathways 在 TPU 上运行 JAX 程序,并且 已在 Google Cloud 上提供.
用 Pathways 协调整个计算
借助 Pathways,你只需从单个 Python 客户端定义端到端的计算,无需编写自定义协调代码,由分布式运行时来执行它。该客户端与硬件解耦。它运行在普通的 CPU 机器上,你针对虚拟设备编写程序;资源管理器把它们映射到真实芯片上,而每台加速器主机上的工作节点则运行它收到的已编译程序。
图 5:虚拟设备、资源管理器,以及主机故障后的重新映射。由于你的代码引用的是虚拟设备而非物理设备,运行时具有灵活性,允许用户更改用于给定计算的设备。这让你保持控制权,同时由运行时处理到硬件的映射。例如,在强化学习循环中,训练器和采样器很少需要相同数量的芯片,因此你可以通过请求更多虚拟设备来给采样器阶段分配更多芯片。
此外,当主机故障或硬件尚未就绪时,同样的能力也能派上用场。如果某台主机宕机,资源管理器会把工作重新映射到健康的芯片上,运行继续进行,而不是重新启动整个作业,这正是 弹性训练的基础。而且由于程序面向虚拟设备,编译可以在物理芯片配置完成之前就开始,因此环境搭建与资源调配可以重叠进行,而不必等待后者。
资源管理器拥有芯片,这也意味着一个 Pathways 集群可以服务多个客户端,因此团队中的研究人员可以共享其容量。借助 Shared Pathways Service,集群作为长期运行的服务保持在线,让你可以在 GKE 托管的 TPU 上进行交互式开发。每个客户端按需从笔记本电脑或虚拟机连接,从共享池中获得自己的切片(紧密连接的 TPU 芯片组),并在断开连接时归还。当所有芯片都忙碌时,下一个工作负载会等待容量释放。
以图 2 中的三块模型为例。块 A 在一组设备上运行,块 B 在第二组上运行,块 C 在第三组上运行。你声明每个块计算什么、目标设备是什么,以及它的输出去向何处。Pathways 把这些依赖关系表达为分布式数据流图,其中每个块是一个已编译的 SPMD 程序,每条边是运行时必须传递的一个张量。
图 6:分布式数据流图。实线箭头跨组,虚线边保持本地。在单个 SPMD 程序中,无论是否使用 Pathways,框架都会处理让组内设备保持同步的集合通信,你的代码无需更改。一旦作业在不同的组上运行多个程序(多程序多数据,MPMD),比如这三个块,或者强化学习循环中的采样器、奖励模型和训练器,差异就显现出来了。此时没有任何单个进程能同时看到两边(这就是之前的限制:每个工作节点只知道自己的设备),因此把张量从块 A 传到块 B 是一次你必须自己协调的交接。
Pathways 弥合了这一差距。你只需声明一次依赖关系,运行时就会为你移动张量并重新分片,在组内通过高速互连,在组间通过数据中心网络。
为了让加速器保持忙碌,Pathways 通过 future 和并行异步调度提前准备启动。当块 A 还在计算时,块 B 尚不能运行,但它的输入形状和 dtype 已经已知,因此运行时会在 A 的输出尚未产生之前就分配 B 的缓冲区并安排传输。每个块的主机侧准备工作都提前完成,每个块在输入到达时启动。
图 7:顺序准备与异步调度对比(示意时长)。这也是防止单一客户端成为瓶颈的原因。由于调度是异步的,协调工作与芯片上已在运行的计算重叠进行,因此客户端始终不在关键路径上,而不会在各步骤之间拖住加速器。
在 Pathways 中,当一个组有多个操作就绪时,由运行时决定执行顺序。在流水线训练中,B 的设备通常同时持有即将到来的微批次的前向传播和较早批次的向后传播。如果每台设备各自做决定,它们的集合通信就会冲突,整个组就会停滞。
Pathways 有一个中央调度器,把每个程序的操作排入一个一致的顺序,因此必须通信的操作会一起启动,而不会死锁。
图 8:没有共享顺序时组会停滞;有一个调度器时它就能持续运转。而且因为它仍然是一个客户端,我们之前提到的开发循环又回来了。无论你驱动一台主机还是数千台主机,都是同一个 Python 进程,所以你可以 进入 shell、设置断点,并在运行过程中检查。当出现停滞时,有一个地方知道每台设备上排队的内容以及哪个依赖尚未满足,而不是让挂起状态分散在各个独立进程中。
综合来看,Pathways 将模型构建与执行分离。你通过客户端选择模型如何在硬件上分区。而运行时负责传输、排序和调度,让加速器保持忙碌,覆盖整个集群而不只是你那一部分。
这样描述下来,Pathways 听起来可能像你已熟悉的编排器。如果调度工作、搬运张量、让整个集群保持忙碌让你想到 Ray,那也应当如此。但两者工作在不同的层次。
这不就是 Ray 做的事吗?
Ray 是一个开源框架,用于在大型集群上编排 Python 原生的 AI 工作负载。它负责配置集群,把你的 Python 代码调度为 任务和 actor,并在各主机上启动你的工作负载,将数据加载、训练、服务和调优置于同一个控制平面之下。对 Ray 来说,加速器只是一种可调度的资源。
Ray 不改变的是你的训练如何在芯片上运行。每个工作节点内的作业 仍然以 SPMD 方式运行。框架插入集合通信,每个进程只能看到自己的设备。Ray 负责放置并启动该作业,但不会把它变成一个驱动整个集群的单一程序。
这正是 Pathways 所处的层次。它改变底层的执行模型,使一个客户端协调每一块芯片。由于它们位于不同层次,两者是互补的。原则上,Ray 可以编排集群和流水线,而 Pathways 负责训练步骤如何执行。
所以具体来说,Ray 的工作是编排,而 Pathways 的工作是执行。使用 Ray,你把数据、训练、服务和调优串联起来,并把异构工作分散到加速器和云上。使用 Pathways,你在多组芯片上运行一个模型,表达超出 SPMD 的模式,并在主机故障时恢复,而不是重新启动整个作业。
结论与后续步骤
当训练作业不再是一个统一的程序时,挑战就变成了协调各阶段之间的依赖关系、跨网络传输的张量,以及让每块加速器都保持供给的排序。
SPMD 要求你自己处理这些依赖关系。使用 Pathways,你从单一客户端一次性把计算描述为块及其之间张量的图。并且你有一个运行时,把块放置到虚拟设备上,搬运并重新分片张量,在整个集群范围内排序工作,并在主机故障时重新映射。
使用 Pathways,模型构建与执行分离,你考虑的是整个集群而不是其中一部分。
如果你想进一步了解 Pathways,可以查看以下资源:
- Google Cloud 上的 Pathways - 文档
- Pathways:面向机器学习的异步分布式数据流 - 论文
- 使用 Pathways 进行弹性训练 - 文档
- Pathways 工具 - GitHub 仓库
- MaxText - GitHub 仓库
祝你构建愉快!
发布于: 上一篇 下一篇 相关文章
AI
云
操作指南
行业趋势
每任务成本:使用廉价分类器在 Gemini 模型和思考层级之间路由
2026 年 11 月 3 日
AI
问题解决
学习
智能体分心
2026 年 11 月 6 日