谷歌下一代AI架构、JeffDean宣传大半年的Pathways终于有论文了

机器之心编辑部

「当前的 AI 模型只做一件事。Pathways 使我们能够训练一个模型，做成千上万件事情。」

在谈及当前的 AI 系统所面临的问题时，低效是经常被提及的一个。

谷歌人工智能主管 Jeff Dean 曾在一篇博文中写道，「今天的人工智能系统总是从头开始学习新问题 —— 数学模型的参数从随机数开始。就像每次学习一项新技能（例如跳绳），你总会忘记之前所学的一切，包括如何平衡、如何跳跃、如何协调手的运动等，然后从无到有重新学习。这或多或少是我们今天训练大多数机器学习模型的方式： 我们不是扩展现有模型来学习新任务，而是从无到有训练新模型来做一件事（或者我们有时将通用模型专门用于特定任务） 。结果是 我们最终为数千个单独的任务开发了数千个模型。以这种方式学习每项新任务不仅需要更长的时间，而且还需要更多的数据。 」

为了改变这种局面，Jeff Dean 等人去年提出了一种名叫「Pathways」的通用 AI 架构。他介绍说，Pathways 旨在用一个架构同时处理多项任务，并且拥有快速学习新任务、更好地理解世界的能力。

该架构的特点可以概括为：

能够训练一个模型来做成千上万件事情；

当前模型只注重一种感官，Pathways 可做到多种；

当前模型密集且效率低下，Pathways 会把模型变得稀疏而高效。

在发布想法大半年之后，Jeff Dean 终于公布了 Pathways 的论文，其中包含很多技术细节。

论文链接

论文写道，PATHWAYS 使用了异步算子的一个分片数据流图（sharded dataflow graph），这些算子消耗并产生 futures，并在数千个加速器上高效地对异构并行计算进行 gang-schedule，同时在它们专用的 interconnect 上协调数据传输。PATHWAYS 使用了一种新的异步分布式数据流设计，它允许控制平面并行执行，尽管数据平面中存在依赖关系。这种设计允许 PATHWAYS 采用单控制器模型，从而更容易表达复杂的新并行模式。

实验结果表明，当在 2048 个 TPU 上运行 SPMD（single program multiple data）计算时，PATHWAYS 的性能（加速器利用率接近 100%）可以媲美 SOTA 系统，同时吞吐量可媲美跨越 16 个 stage 或者被分割成两个通过数据中心网络连接的加速器岛的 Transformer 模型的 SPMD 案例。

以下是论文细节：

研究背景

在过去的十年里，深度学习在图像理解、自然语言处理等多个领域取得了显著的进展，这是 ML 模型、加速器硬件以及将两者联系在一起的软件系统协同进化的结果。这种协同进化带来的隐患是：深度学习系统可能过度专注于当前的工作负载，无法预测未来的需求。

PATHWAYS 是一个为分布式 ML 构建的新系统，剑指未来 ML 工作负载将需要的特定能力。当前，这些工作负载缺乏 SOTA 系统的支持。

例如，当今 SOTA ML 工作负载大多使用单程序多数据（SPMD）模型，该模型受到了 MPI 的启发，其中所有加速器都在同步运行相同的计算，加速器之间的通信由 AllReduce 等集体来描述。

但近年来，研究人员开始在 ML 计算中被 SPMD 掣肘。大型语言模型已经使用流水线并行而不是纯粹的数据并行来扩展；混合专家（MoE）等模型已经开始探索计算稀疏性，其最自然的表达方式是使用细粒度控制流和跨加速器的异构计算；系统设计者们已经开始采用巧妙的技术来在 MPI 风格的系统上执行流水线（pipelined）、同构 MoE 模型，但是，MPI 编程模型对于用户和底层系统来说都太受限制了。

另一方面，随着一代又一代新加速器的出现，ML 中的异构环境变得越来越普遍。提供对（通过高带宽 interconnect 连接的）同构加速器的大「岛」的独占访问是昂贵的，并且通常会造成浪费，因为单个用户程序必须努力保持所有加速器持续忙碌。这种限制进一步推动研究人员走向「多程序多数据」（MPMD）计算。MPMD 通过将整个计算的子部分映射到一组更容易获得的小加速器岛上来实现更大的灵活性。为了提高利用率，一些 ML 硬件资源管理研究人员以细粒度的方式在工作负载之间复用硬件，实现工作负载弹性，并提高容错能力。

最后，研究人员开始标准化一套基础模型（foundation model），这些模型是在大数据上大规模训练的，可以适应多种下游任务。通过在许多任务之间复用资源，并在它们之间有效地共享状态，这种模型的训练和推理提供了提高集群利用率的机会。例如，几个研究人员可能同时微调用于不同任务的一个基础模型，使用相同的加速器来保持固定的基础模型层。在共享的子模型上进行的训练或推理可以受益于一些技术，这些技术允许来自不同任务的示例被组合在一个 vectorized batch 中，以获得更高的加速器利用率。

本文提出的 PATHWAYS 在功能和性能上可以媲美 SOTA ML 系统，同时提供了支持未来 ML 工作负载所需的能力。它使用了一个 client-server 架构，该架构使得 PATHWAYS 的运行时能够代表许多 client 在系统管理计算岛上执行程序。

PATHWAYS 是第一个旨在透明、高效地执行跨多个 TPU pods 的程序的系统。通过采用新的数据流执行模型，它可以扩展到数千个加速器。PATHWAYS 的编程模型使得表达非 SPMD 计算变得很容易，并支持集中的资源管理和虚拟化，以提高加速器的利用率。

PATHWAYS 系统架构

PATHWAYS 构建在先前的系统的基础上，包括用于表征和执行 TPU 计算的 XLA (TensorFlow, 2019)、用于表征和执行分布式 CPU 计算的 TensorFlow 图和执行器 (Abadi et al., 2016)，以及包括 JAX (Bradbury et al., 2016) 在内的 Python 编程框架 (Bradbury et al., 2018) 和 TensorFlow API。利用这些构建块，PATHWAYS 在兼顾协调性的同时，仅用最少的代码更改就能运行现有的 ML 模型。

资源管理器

PATHWAYS 的后端由一组加速器组成，这些加速器组合成紧密耦合的 island，这些 island 又通过 DCN 相互连接，如上图 3 所示。PATHWAYS 有一个「资源管理器」，负责集中管理所有 island 上的设备。client 可能会要求 island 的「虚拟 slice」具有适合其通信模式的特定 2D 或 3D 网格形状。每个虚拟 slice 都包含「虚拟设备」，允许 client 表达计算在网格上的布局方式。资源管理器为满足所需互连拓扑、内存容量等的虚拟设备动态分配物理设备。

最初的资源管理器使用一个简单的启发式方法来实现，尝试通过在所有可用设备上传播计算来静态平衡负载，并在虚拟设备和物理设备之间保持一对一的映射。如果未来的工作负载需要，则可以采用更加复杂的分配算法，例如考虑所有 client 计算的资源需求和系统的当前状态，以近似计算物理设备的最佳分配。

PATHWAYS 允许动态添加和移除后端计算资源，由资源管理器跟踪可用设备。由单控制器设计启用的虚拟设备和物理设备之间的间接层将允许未来支持透明的挂起 / 恢复和迁移等功能，其中 client 的虚拟设备可以临时回收资源或重新分配而无需用户程序的协助。

client

当用户想要运行一个被跟踪的程序时，可以调用 PATHWAYS client 库，它首先将虚拟设备分配给之前没有运行过的任何计算，并用资源管理器注册计算，触发 server 在后台编译计算。

然后，client 为程序构建与设备位置无关的 PATHWAYS 中间表征 (IR)，表示为自定义 MLIR (Lattner et al., 2021) dialect。IR 通过一系列标准编译器 pass 逐渐降低级别，最终输出包含物理设备位置的低级表征。这种低级程序考虑了物理设备之间的网络连接，并包含将输出从源计算分片传输到其目标分片（shard）位置的操作，包括需要数据交换时的分散和收集操作。在虚拟设备位置不变的通常情况下重复运行低级程序是有效的，如果资源管理器改变了虚拟设备和物理设备之间的映射关系，可以 re-low 程序。

较旧的单控制器系统中的 client 可能很快成为性能瓶颈，因为它负责协调数千个单独的计算，还要协调分布在数千个加速器中的计算分片相应的数据缓冲区。PATHWAYS client 使用分片缓冲区抽象来表征可能分布在多个设备上的逻辑缓冲区。这种抽象通过以逻辑缓冲区而不是单个分片的粒度分摊 bookkeeping 任务（包括参考计数（reference counting））的成本来帮助 client 扩展。

协调实现

PATHWAYS 依赖 PLAQUE 完成所有使用 DCN 的跨主机协调。PLAQUE 是一种现有的（闭源）生产分片数据流系统，谷歌将它用于许多面向客户的服务，这些服务需要高扇出或高扇入通信，并且可扩展性和延迟都很重要。低级 PATHWAYS IR 直接被转换为 PLAQUE 程序，并表征为数据流图。PATHWAYS 对其协调 substrate 有严格的要求，而 PLAQUE 满足所有要求。

首先，用于描述 PATHWAYS IR 的表征必须包含每个分片计算的单个节点，以确保能够紧凑表征跨多个分片的计算，即带有 N 个计算分片的 2 个计算 A 和 B 的链式执行，无论 N 是多少，每个计算分片在数据流表征中都有 4 个节点：Arg → Compute (A) → Compute (B) → Result。在 PLAQUE 运行时实现中，每个节点都会生成带有目标分片标记的输出数据元组，因此在执行数据并行执行时，N 个数据元组将在每对相邻的 IR 节点之间流动。

协调运行时还必须支持沿分片边缘的稀疏数据交换，其中消息可以在动态选择的分片子集之间发送，使用标准的进度跟踪机制（Akidau et al., 2013; Murray et al., 2013）来检测何时已收到分片的所有消息。高效的稀疏通信能够避免 DCN 成为加速器上依赖于数据的控制流瓶颈，这是 PATHWAYS 启用的关键功能之一。

如下图 4 所示，协调 substrate 用于发送传输调度消息和数据 handle 的关键路径中的 DCN 消息，因此它必须以低延迟发送关键消息，并在需要高吞吐量时将消息批量发送到同一个 host。

使用可扩展的通用数据流引擎来处理 DCN 通信也很方便，因为这意味着 PATHWAYS 还可以将其用于后台管理任务，例如分发配置信息、监控程序、清理程序、在出现故障时提示错误等。

谷歌认为，使用 Ray (Moritz et al., 2018) 等其他分布式框架而不是 PLAQUE 来重新实现完整的 PATHWAYS 设计以实现低级协调框架是可行的。在这种实现中，PATHWAYS 执行器和调度器将被长期运行的 Ray Actor 所取代，这些 Ray Actor 将在底层 Ray 集群调度之上实现 PATHWAYS 调度，并且执行器可以使用 PyTorch 进行 GPU 计算和集合。

Gang-scheduled 动态调度

如前所述，在一组共享加速器上支持 SPMD 计算的一个要求是支持高效的 gang-scheduling。

PATHWAYS 运行时包括每个 island 的集中式调度器，它对 island 上所有计算进行一致性排序。当 PATHWAYS 将一个程序加入队列以执行时，PLAQUE 数据流程序负责以下操作：

在每个加速器上将本地编译函数执行加入队列，并将缓冲 future 作为输入；

将网络发送（network sends）加入到远程加速器的队列，以获得函数执行输出的缓冲 future；

与调度器通信，以确定在 island 上运行的所有程序中函数执行的一致顺序。

调度器必须实施以毫秒为单位分配加速器的策略。不过，当前的实现只是按照 FIFO 顺序将工作加入队列，但更复杂的调度器可能会根据估计的执行时间重新排序计算。

并行异步调度

当在加速器上运行计算时，系统可以利用异步 API 将计算与协调重叠。如下图 4a 中的三节点图所示，正方形分别对应三个节点 A、B 和 C，它们在连接到主机 A、B 和 C 的加速器上运行。所有节点计算都是常规编译函数。主机 A 将节点 A 加入队列，接收 A 输出的 future 并将它传输给主机 B。主机 B 分类节点 B 的输入，将该输入缓冲地址传输给节点 A，并执行大部分准备工作以启动节点 B 的功能。当节点 A 完成时，它的输出直接通过加速器互联发送至节点 B 的输入缓冲，然后主机 B 启动节点 B。一个节点完成和另一个节点启动之间的延迟时间要比数据传输时间更长。

当 predecessor 节点的计算时间超过主机之间调度、资源分类和协同所用时间时，上述设计运行良好。但如果计算时间太短，异步 pipeline 就会停止，主机端的工作成为执行整个计算序列过程中的关键瓶颈。考虑到编译的函数都是常规的，后续节点的输入形状实际上可以在 predecessor 计算加入队列之前进行计算。

因此，谷歌引入了一种全新的并行异步调度设计方案，具体如下图 4 b 所示。该方案利用常规编译函数的静态已知资源来并行运行计算节点的主机端工作，而不是在 predecessor 已经加入队列之后对节点工作进行序列化处理。考虑到常规函数下只能并行地调度工作，PATHWAYS 将并行调度作为一种优化手段，并在节点资源需求在 predecessor 计算完成时才知道的情况下回退到传统模型。

当计算的子图可以进行静态调度时，该程序会向调度器发送描述整个子图的单条消息，该调度器能够对子图中所有活动分片的执行进行背靠背排序。设计单条消息旨在最小化网络流量，但不需要调度器将所有子图的分片作为一个批次来加入队列：计算仍可能与其他并发执行程序提交的计算交错。

三节点程序的顺序调度（a）与并行调度（b）比较。

实验结果

谷歌展示了 PATHWAYS 在训练真实机器学习模型（它们可以被表示为 SPMD 程序）中的性能。首先与使用编码器 - 解码器架构运行 Transformer 模型的 JAX 多控制器进行比较。

下表 1 展示了在不同数量的加速器上训练时，不同大小的文本到文本 Transformer 模型的训练吞吐量（tokens / 秒）。正如所预期的一样，由于模型代码相同，在 JAX 和 PATHWAYS 上训练的模型在步数相同的情况下实现了相同的困惑度。

接着，谷歌比较了当仅用解码器架构训练 Transformer 语言模型时，PATHWAYS 在不同配置上的性能。如表 2 所示，PATHWAYS 的训练吞吐量与每个 pipeline 阶段的 TPU 核心数量成比例增加，这与其他系统保持一致。

上述结果与下图 5 一致， 表明 PATHWAYS 的吞吐量与主机数量呈线性缩放关系。增加 pipeline 阶段的数量会提高最小开销，当阶段数量从 4 增加到 16 时，吞吐量从 133.7k tokens / 秒减少到 131.4k tokens / 秒。谷歌将 pipelined 模型的性能与使用 SPMD 的等效模型进行了比较，并观察到至少在这种情况下，pipeline 的性能与 SPMD 相当，这是因为 SPMD 计算内部聚合通信产生的开销比 pipeline 泡沫（bubble）开销更高。