论文标题和作者

文章标题：PATHWAYS: 面向机器学习的异步分布式数据流
作者：Paul Barham, Aakanksha Chowdhery, Jeff Dean, Sanjay Ghemawat, Steven Hand, Dan Hurt, Michael Isard, Hyeontaek Lim, Ruoming Pang, Sudip Roy, Brennan Saeta, Parker Schuh, Ryan Sepassi, Laurent El Shafey, Chandramohan A. Thekkath, Yonghui Wu
机构：Google Research

A1 主要贡献

本文介绍了一种名为PATHWAYS的新型大规模分布式机器学习系统。当前系统在模型、硬件和软件的共同演进中，存在过度专精于现有工作负载的风险。PATHWAYS旨在解决这一问题，支持未来机器学习工作负载所需的功能，同时保持现有模型的顶尖性能。

核心问题与研究目标:
1. SPMD模型的局限性：当前最先进的机器学习工作负载大多采用“单程序多数据”（SPMD）模型，如MPI【16，The MPI message passing interface standard，1994，Programming Environments for Massively Parallel Distributed Systems】。所有加速器同步运行相同计算，通信依赖于AllReduce等集合操作。然而，对于超大规模语言模型中使用的流水线并行【61，PipeDream: Generalized pipeline parallelism for DNN training，2019，ACM SOSP；70，DeepSpeed: System optimizations enable training deep learning models with over 100 billion parameters，2020，ACM SIGKDD；63，Efficient large-scale language model training on GPU clusters using Megatron-LM，2021，arXiv】和专家混合（MoE）模型【79，Outrageously large neural networks: The sparsely-gated mixture-of-experts layer，2017，ICLR】中的计算稀疏性，SPMD模型限制性太强。
2. 硬件异构性与资源利用率：随着新一代加速器的出现，机器学习集群日益异构。为大量同构加速器提供独占访问权限成本高昂且浪费。研究人员开始转向“多程序多数据”（MPMD）计算，以更灵活地将计算子部分映射到更易获得的较小加速器集群上，从而提高利用率。
3. 基础模型的需求：基础模型（Foundation Models）【11，On the opportunities and risks of foundation models，2021，arXiv】的兴起带来了新的机遇，例如通过在多任务间复用资源和高效共享状态来提高集群利用率。这需要系统支持多用户并发微调、共享子模型上的训练和推理等。

创新点与主要贡献:
PATHWAYS系统旨在匹配现有系统的功能和性能，同时提供支持未来ML工作负载所需的能力。
1. 灵活的单控制器架构：PATHWAYS采用客户端-服务器架构，其运行时可以在系统管理的计算集群上为多个客户端执行程序。这使其易于表达非SPMD计算，并支持集中式资源管理和虚拟化，以提高加速器利用率。
2. 可扩展的异步分布式数据流模型：PATHWAYS是首个设计用于透明高效地执行跨多个TPU "pod"【24，Cloud TPU，2021，https://cloud.google.com/tpu】的程序系统。它采用一种新颖的异步分布式数据流设计，通过分片数据流图和异步操作符（消费和产生future）来协调数千个加速器上的异构并行计算和数据传输。这种设计允许控制平面在数据平面存在依赖的情况下并行执行，从而解决了传统单控制器模型的性能瓶颈。
3. 性能与功能验证：
- 在2048个TPU上运行SPMD计算时，PATHWAYS的性能与最先进的多控制器系统相当（加速器利用率约100%）。
- 对于跨16个阶段进行流水线处理的Transformer模型，或分片到两个通过数据中心网络连接的加速器集群上的模型，其吞吐量与SPMD情况相当。

A3 背景知识/关键Observation/设计原则

2 设计动机

现有分布式ML系统的设计驱动因素：分布式机器学习系统的设计选择通常由底层目标硬件加速器的特性驱动。附录A详细讨论了这些特性及其对分布式ML系统的影响。本节重点关注现有分布式ML系统的设计和实现选择如何使其难以支持大型、稀疏或不规则模型。

多控制器架构的优势与局限性：用于训练最先进SPMD模型的分布式ML系统，如MPI【16，The MPI message passing interface standard，1994，Programming Environments for Massively Parallel Distributed Systems】、PyTorch【68，PyTorch: An imperative style, high-performance deep learning library，2019，NeurIPS】、JAX【12，JAX: Composable transformations of Python+NumPy programs，2018，http://github.com/google/jax】以及新版TensorFlow【78，Mesh-TensorFlow: Deep learning for supercomputers，2018，NeurIPS；2，TensorFlow Eager: A multi-stage, Pythonembedded DSL for machine learning，2019，arXiv】，通常采用多控制器架构。该架构在所有主机上直接运行相同的客户端可执行文件，并在程序执行期间独占这些主机的资源。其关键优势是分发加速器计算的延迟低（见图1a），因为用户代码的相同副本在每个加速器主机上运行，分发仅涉及通过相对较快的PCIe链路进行通信。然而，这种架构与使用流水线或计算稀疏性的现代ML工作负载不匹配。任何超出标准集合操作的通信都需要用户实现自己的协调原语。多控制器方法还假定对硬件资源的独占所有权，这不仅将确保昂贵加速器高利用率的责任转移给了用户，也使资源虚拟化和复用等功能的设计复杂化，而这些功能是构建高效集群级ML基础设施所必需的。

单控制器架构的灵活性与挑战：单控制器系统，如TensorFlow v1【1，TensorFlow: A system for large-scale machine learning，2016，OSDI】，提供了非常通用的分布式数据流模型，包括优化的图内控制流【91，Dynamic control flow in large-scale machine learning，2018，EuroSys】。一个TensorFlow (TF) Python客户端构建一个计算图，并将其交给一个协调器运行时，该协调器将图划分为每个工作节点的子图，并委托工作节点上的本地运行时执行。工作节点之间的协调通过数据和控制边在数据中心网络（DCN）上传递消息来完成。尽管单控制器设计提供了灵活的编程模型和资源虚拟化，但它也带来了实现上的挑战。

单控制器系统的实现挑战：首先，多控制器系统仅需通过PCIe通信来分发加速器计算（图1a），而单控制器系统中的客户端“距离更远”，分发延迟涉及DCN通信，通常比PCIe慢一个数量级（图1b）。其次，为了支持MPMD程序与SPMD子计算的并发执行，运行时必须有某种机制来支持对加速器计算进行组调度（gang-scheduling）。组调度对于TPU至关重要，因为它们是单线程且只运行不可抢占的内核，如果通信计算没有以一致的顺序入队，系统将死锁。即使对于可以执行并发计算的GPU或其他加速器，组调度也能更有效地执行集合操作【22，Gang scheduling performance benefits for fine-grain synchronization，1992，Journal of Parallel and Distributed Computing】。因此，用于ML的单控制器系统需要一个分布式调度机制来对代表不同程序入队的计算进行排序。最后，一个用于现代ML工作负载的系统必须被设计为能在数千个加速器上运行分布式计算，并对分片表示和数据结构提供一流的支持。例如，一个表示M路分片计算和N路分片计算之间边的朴素数据流图将需要M+N个节点和M×N条边，很快就会变得难以管理。

图1. 多控制器和单控制器系统之间分发开销和通信模式的比较。(a) Jax或PyTorch SPMD通过快速PCIe独立异步地将加速器计算入队；(b) TensorFlow v1 SPMD需要通过较慢的DCN发送控制消息；(c) TensorFlow v1非SPMD程序需要通过显式的发送(S)和接收(R)操作进行跨主机协调或数据传输。

TensorFlow v1的局限性：TF v1的实现选择过于专精，假定只有一个小型的、独占的加速器集群。这种过度专精使得TF在实践中无法用于当代或未来的ML工作负载。虽然TF可以通过发送和接收操作运行需要跨主机协调或数据传输的计算（图1c），但目标端的主机侧工作（如分发加速器计算）只有在传输完成后才被触发。在涉及许多跨主机传输的程序中，例如具有大量阶段的流水线模型，这些分发延迟会累积，导致加速器利用率低下。虽然TF v1用户可以通过使用控制边在单个程序内（低效地）强制执行一致的组调度顺序，但像TF v1这样的单控制器系统中缺乏集中式调度器，使得无法确保跨程序的计算之间的一致顺序。TF还物化了完整的分片计算图，当分片数量达到数千时，这会在图序列化和执行中引入巨大的开销，导致子计算之间产生数百万条图边。

PATHWAYS的设计目标：PATHWAYS旨在结合单控制器框架的灵活性与多控制器的性能。我们采用单控制器模型，因为它为新颖高效的ML计算提供了更好的机会，既可以通过利用计算稀疏性和异构性，也可以通过启用促进资源共享和虚拟化的集群管理系统。我们的设计与旧的单控制器ML系统不同，它使用异步分发来匹配多控制器系统的性能，支持集中式资源管理和调度，并对SPMD加速器计算的组（gangs）提供一流支持，同时使用分片数据流系统进行高效协调。

3 PATHWAYS 编程模型

对TensorFlow和JAX的支持：我们已经实现了从TensorFlow和JAX编写的源程序来定位PATHWAYS的支持，但在本文的评估中，我们主要关注JAX。JAX用户可以使用装饰器显式地包装标准Python代码，以指示应编译为（可能是SPMD）XLA计算的代码片段。这些XLA计算通常具有已知的输入输出类型和形状、有界循环以及少量（若有）条件判断（详见附录B），这使得提前估算计算的资源需求变得可行。我们将这些具有已知资源需求的计算称为“已编译函数”（compiled functions）。每个这样的函数都映射到PATHWAYS程序中的一个（分片的）计算节点。

扩展JAX至多TPU pod：目前的JAX无法扩展到单个TPU pod之外，因为在多控制器配置中运行的JAX程序使用XLA集合操作传输所有数据，而这些操作目前只能在TPU的ICI（Inter-Chip Interconnect）上使用。PATHWAYS可以作为JAX后端的插件替代品，允许JAX代码无需修改即可运行，区别在于SPMD计算现在不仅可以访问本地连接的TPU核心，还可以访问系统中配置的任意数量的核心。由于PATHWAYS可以通过ICI和DCN进行通信，它首次允许JAX程序扩展到包含数千个TPU核心的多个TPU pod。

图2. 在PATHWAYS上运行跨多个TPU集群的分片计算的Python用户代码示例。

虚拟设备与程序追踪：虽然能够运行未经修改的JAX代码很方便，但这并未完全释放PATHWAYS的性能。PATHWAYS用户可以请求一组“虚拟设备”，并可选择性地对设备类型、位置或互连拓扑施加约束，然后能够将特定的已编译函数放置在这些设备上（图2）。系统将自动处理相关计算之间的所有数据移动和重新分片。

程序追踪器：默认情况下，我们将每个已编译函数转换为一个独立的PATHWAYS程序，该程序只包含一个（分片的）计算。这意味着如果用户想连续运行许多函数，每个函数都需要一个单独的Python调用和从客户端到协调器的RPC。因此，我们还实现了一个新的程序追踪器（图2），用户可以用它来包装一个调用多个已编译函数的Python代码块。该追踪器会生成一个单一的PATHWAYS程序，其中每个已编译函数都由数据流图中的一个计算节点表示。

与JAX生态系统的契合：JAX支持对追踪代码进行转换的理念与我们希望探索的研究方向非常契合。例如，JAX有一个名为FLAX【30，Flax: A neural network library and ecosystem for JAX，2020，http://github.com/google/flax】的配套库，用于表达分层的DNN模型，我们编写了一个库，可以自动将FLAX模型转换为流水线化的PATHWAYS程序。此外，JAX支持将“逐样本”的Python函数向量化的转换，从而生成高效的批处理代码，这种转换为探索新形式的数据依赖向量化控制流提供了良好的基础，我们将在后面简要描述（§6.3）。

A2 方法细节

4 PATHWAYS 系统架构

PATHWAYS广泛地构建在先前的系统之上，包括使用XLA【82，XLA: Optimizing compiler for TensorFlow，2019，https://www.tensorflow.org/xla】来表示和执行TPU计算，使用TensorFlow图和执行器【1，TensorFlow: A system for large-scale machine learning，2016，OSDI】来表示和执行分布式CPU计算，以及包括JAX【12，JAX: Composable transformations of Python+NumPy programs，2018，http://github.com/google/jax】和TensorFlow API在内的Python编程框架。通过利用这些构建模块，我们能够专注于PATHWAYS新颖的协调方面，同时能够以最少的代码更改运行现有的ML模型。

4.1 资源管理器

资源管理器的结构与功能：PATHWAYS后端由一组加速器组成，这些加速器被分组为紧密耦合的集群（islands），集群之间通过DCN连接（图3）。PATHWAYS有一个“资源管理器”，负责对所有集群中的设备进行集中管理。客户端可以请求集群的“虚拟切片”（virtual slices），并指定适合其通信模式的特定2D或3D网格形状。每个虚拟切片包含“虚拟设备”，允许客户端表达计算在网格上的布局方式。资源管理器会动态地为虚拟设备分配满足所需互连拓扑、内存容量等条件的物理设备。

图3. PATHWAYS系统概览。(左) 分布式计算表示为一个有向无环图(DAG)，其中每个节点代表一个独立的已编译函数，节点之间的边代表函数间的数据流。(中) 资源管理器为每个已编译函数分配集群加速器的子集(“虚拟切片”)。(右) 每个集群的集中式调度器对计算进行组调度，然后由每个分片的执行器进行分发。红色箭头表示控制消息，蓝色箭头表示数据路径传输。

资源分配策略：我们最初的资源管理器实现使用了一种简单的启发式方法，试图通过将计算分散到所有可用设备上来静态地平衡负载，并保持虚拟设备和物理设备之间的一对一映射。如果未来的工作负载需要，我们可以采用更复杂的分配算法，例如，考虑所有客户端计算的资源需求和系统的当前状态，来近似地实现物理设备到计算的最优分配。

动态资源管理与虚拟化：PATHWAYS允许后端计算资源动态地添加和移除，由资源管理器跟踪可用设备。通过我们的单控制器设计实现的虚拟设备和物理设备之间的间接层，未来将使我们能够支持透明的挂起/恢复和迁移等功能，即客户端的虚拟设备可以被临时回收或重新分配，而无需用户程序的协作。

4.2 客户端

客户端工作流程：当用户想要运行一个被追踪的程序时，它会调用PATHWAYS客户端库。该库首先为之前未运行过的任何计算分配虚拟设备，并向资源管理器注册这些计算，触发服务器在后台编译这些计算。然后，客户端为程序构建一个设备位置无关的PATHWAYS中间表示（IR），该表示以自定义的MLIR【50，MLIR: Scaling compiler infrastructure for domain specific computation，2021，IEEE/ACM CGO】方言形式表达。该IR通过一系列标准的编译器遍（passes）逐步“降低”（lowered），最终输出一个包含物理设备位置的低级表示。这个低级程序考虑了物理设备之间的网络连接性，并包含将输出从源计算分片传输到其目标分片位置的操作，包括在需要数据交换时的散播（scatter）和收集（gather）操作。在虚拟设备位置不变的常见情况下，重复运行低级程序是高效的；如果资源管理器更改了虚拟设备和物理设备之间的映射，则可以重新降低程序。

分片缓冲区抽象以实现可扩展性：在旧的单控制器系统中，客户端可能会迅速成为性能瓶颈，因为它需要协调分布在数千个加速器上的数千个单独计算和数据缓冲区。PATHWAYS客户端使用一种分片缓冲区抽象来表示一个可能分布在多个设备上的逻辑缓冲区。这种抽象通过在逻辑缓冲区的粒度上摊销簿记任务（包括引用计数）的成本，而不是在单个分片上进行，从而帮助客户端实现扩展。

4.3 协调实现

协调运行时的需求：协调运行时还必须支持沿着分片边的稀疏数据交换，其中消息可以在动态选择的分片子集之间发送，使用标准的进度跟踪机制【3，MillWheel: Fault-tolerant stream processing at internet scale，2013，VLDB；60，Naiad: A timely dataflow system，2013，ACM SOSP】来检测何时一个分片的所有消息都已接收。高效的稀疏通信是避免DCN成为加速器上数据依赖控制流瓶颈的一个要求，而这是我们希望PATHWAYS启用的关键能力之一。协调基板用于发送处于关键路径上的DCN消息，用于传输调度消息和数据句柄（图4），因此它必须以低延迟发送关键消息，并在需要高吞吐量时批量处理发往同一主机的消息。

使用PLAQUE作为协调基板：PATHWAYS依赖PLAQUE来完成所有使用DCN的跨主机协调。PLAQUE是Google内部一个现有的（闭源）生产级分片数据流系统，用于许多需要高扇出或高扇入通信且对可扩展性和延迟都有重要要求的面向客户的服务。低级的PATHWAYS IR被直接转换为一个PLAQUE程序，表示为一个数据流图。PATHWAYS对其协调基板有严格的要求，而PLAQUE满足了所有这些要求。首先，用于描述PATHWAYS IR的表示必须为每个分片计算包含一个单一节点，以确保跨越许多分片的计算有一个紧凑的表示。其次，PLAQUE运行时实现中，每个节点生成的输出数据元组都带有目标分片标签，因此在执行数据并行时，N个数据元组会流动，每个相邻IR节点对之间一个。

通用数据流引擎的便利性：使用一个可扩展、通用的数据流引擎来处理DCN通信也很方便，因为这意味着PATHWAYS也可以用它来执行后台的内务管理任务，如分发配置信息、监控程序、清理程序、在失败时传递错误等。

替代方案的可行性：我们认为，使用其他分布式框架如Ray【58，Ray: A distributed framework for emerging AI applications，2018，OSDI】来替代PLAQUE，实现低级协调框架，从而重新实现完整的PATHWAYS设计是可行的。在这样的实现中，PATHWAYS的执行器和调度器将被实现PATHWAYS调度的长期运行的Ray actor所取代，而执行器可以使用PyTorch进行GPU计算和集合操作。可能需要一些补充来实现可比的性能（见§5），因为Ray缺乏例如HBM对象存储，或者在GPU互连上传输远程对象的高效原语。

4.4 Gang调度动态分发

Gang调度的必要性与实现：如前所述（§2），在共享的加速器集上支持SPMD计算的一个要求是支持高效的组调度（gang-scheduling）。PATHWAYS运行时为每个集群包含一个集中式调度器，该调度器对集群中的所有计算进行一致的排序。当PATHWAYS将一个程序入队执行时，PLAQUE数据流程序负责（i）在每个加速器上将本地已编译函数的执行入队，以缓冲区future作为输入；（ii）将函数执行输出的缓冲区future的网络发送操作入队到远程加速器；以及（iii）与调度器通信，以确定在集群上运行的所有程序之间函数执行的一致顺序。调度器必须实现以毫秒级时间尺度分配加速器的策略。我们当前的实现只是简单地按FIFO顺序将工作入队，但更复杂的调度器可能会根据估计的执行时间重新排序计算。

4.5 并行异步分发

传统异步分发的局限性：在加速器上运行计算时，系统可以利用异步API来重叠计算与协调【48，Nimble: Lightweight and parallel GPU task scheduling for deep learning，2020，NeurIPS】。考虑图4a中的三节点图，其中方块对应于在主机A、B和C上的加速器上运行的三个节点A、B和C。所有节点计算都是常规的已编译函数。主机A将节点A入队，接收A输出的future，并将该future传输给主机B。主机B分配B的输入，将输入缓冲区地址传输给主机A，并执行启动节点B函数的大部分准备工作。当节点A完成时，其输出通过加速器互连直接发送到节点B的输入缓冲区，然后主机B启动节点B。一个节点完成到下一个节点启动之间的延迟可以做到仅比数据传输时间稍长。上述设计在前辈节点的计算时间长于调度、资源分配和主机间协调所花费的时间时效果很好。但是，如果计算时间太短（如图中所示），异步流水线会停滞，主机侧的工作成为执行整个计算序列的关键瓶颈。

图4. 顺序分发与并行分发的对比（针对一个三节点程序）。当设备上的计算执行时间短于调度、资源分配和协调所花费的时间时，由于主机侧工作的存在，顺序分发中的异步流水线会发生停滞。并行异步分发通过并行运行主机侧工作来克服这一瓶颈，利用了常规已编译函数静态已知的资源使用情况。为简洁起见，图中省略了调度器。

并行异步分发的创新设计：鉴于已编译函数都是常规的，后继节点的输入形状实际上可以在前辈计算入队之前就算出。因此，我们引入了一种新颖的并行异步分发设计（图4b），它利用常规已编译函数的静态已知资源使用情况，并行运行一个计算的多个节点的大部分主机侧工作，而不是将一个节点的工作序列化到其前辈节点入队之后。由于只有当函数是常规的时才能并行调度工作，PATHWAYS将并行调度视为一种优化，并在节点的资源需求直到前辈计算完成才可知时（例如，由于数据依赖的控制流），回退到传统模型。当一个计算的子图可以被静态调度时，程序向调度器发送一条描述整个子图的单一消息，调度器能够将子图中所有活动分片背靠背地序列执行。使用单一消息旨在最小化网络流量，但这并不要求调度器实际将所有子图的分片作为一个批次入队：计算仍然可以与并发执行的其他程序提交的计算交错进行。我们在§5中评估了不同分发机制的成本。

4.6 数据管理

分片对象存储：每个主机管理一个分片对象存储，类似于Ray的对象存储【58，Ray: A distributed framework for emerging AI applications，2018，OSDI】，但扩展到还可以跟踪每个分片中保存在加速器HBM中的缓冲区。客户端程序可以持有对远程主机或加速器内存中对象的引用，客户端和服务器使用不透明的句柄来引用它们，这使得系统在需要时可以迁移它们。中间程序值也保存在对象存储中，例如，当系统等待在加速器之间传输它们或将它们传递给后续计算时。这些对象被标记了所有权标签，以便在程序或客户端失败时可以进行垃圾回收。

内存管理机制：我们可以使用简单的反压机制来暂停一个计算，如果它因为其他计算的缓冲区暂时占用了HBM而无法分配内存。

A4 实验

实验环境

硬件配置:
- TPU配置(A): 每台主机4个TPU，最大实例包含512台主机，总计2048个通过ICI连接的TPU。
- TPU配置(B): 每台主机8个TPU，最大实例包含64台主机，总计512个TPU。
- TPU配置(C): 四个TPU集群（islands），每个集群有4台主机和32个TPU。
- GPU配置: 用于Ray评估，使用p3.2xlarge虚拟机，主机通过DCN连接，并使用Amazon放置组进行调度。
软件配置:
- 框架: JAX, PATHWAYS, TensorFlow, Ray v1.3, PyTorch 1.8.1。
- 对比基线: 主要与多控制器JAX进行比较，因为它在行业标准基准测试中表现出顶尖性能【57，MLPerf: An industry standard benchmark suite for machine learning performance，2020，IEEE Micro】。部分微观基准测试也与TensorFlow (TF)和Ray进行了比较。

实验结果

5.1 单控制器分发开销

实验内容：通过微观基准测试比较JAX多控制器与单控制器框架（PATHWAYS, TF, Ray）的开销。程序重复运行一个包含单个标量AllReduce和标量加法的简单组调度计算。测量吞吐量（每秒计算次数）。比较了三种入队方式：
- OpByOp (-O): 每次计算都单独调用。
- Chained (-C): 一次调用执行一个包含128个计算节点的链。
- Fused (-F): 一次调用执行一个包含128个计算的融合节点。

实验结果与分析（图5）：
- OpByOp: JAX多控制器吞吐量远优于单控制器系统，尤其是在加速器数量增加时。PATHWAYS的主要开销来自客户端等待协调器完成一次计算入队才进行下一次。
- Fused (-F): 当足够多的工作被融合进一个节点时，PATHWAYS在多达1000个TPU核心上性能与JAX相当。
- Chained (-C): PATHWAYS Chained在多达256个核心上性能优于JAX OpByOp，因为PATHWAYS可直接从C++执行背靠背的加速器计算，而JAX OpByOp每次计算都需要切换回Python。
- TF和Ray: 由于缺乏设备对象存储，性能较差。TF在多核心上性能慢是因为其使用集中式屏障进行序列化。
- 结论: 实验表明，如果将PATHWAYS的设计理念（如设备上对象存储）应用于Ray等系统，有望实现相当的性能。

图5. PATHWAYS与TF、JAX和Ray的分发开销对比。PATHWAYS在所有配置上均优于TF和Ray等单控制器系统，并在Fused (-F)和Chained (-C)配置下，分别在多达1000和256个TPU核心上与多控制器JAX的性能持平。每个计算包含一个标量AllReduce后跟一个标量加法。

计算规模对性能影响（图6）：
- 实验内容: 改变每个计算的耗时，找出PATHWAYS与JAX吞吐量持平的最小计算规模。
- 实验结果: 在配置(B)的16台主机（128个TPU）上，计算耗时仅需2.3ms即可持平。在配置(A)的512台主机（2048个TPU）上，计算耗时需35ms才能掩盖PATHWAYS的单控制器开销。

图6. PATHWAYS与JAX达到相同吞吐量的最小计算规模，以掩盖单控制器开销。在配置(B)的16台主机（128个TPU）上，计算规模至少为2.3ms时PATHWAYS与JAX吞吐量持平；在配置(A)的512台主机（2048个TPU）上，计算规模至少为35ms。

并行异步分发性能（图7）：
- 实验内容: 构建一个流水线基准测试，每个计算在不同的4个TPU核心（不同主机）上运行，数据通过ICI传输。评估并行异步分发机制（§4.5）的优势。
- 实验结果: 随着主机数量增加，并行异步分发能摊销固定的客户端和调度开销。与强制使用顺序异步分发相比，并行异步分发显著提升了性能。

图7. PATHWAYS中的并行与顺序异步分发对比。每个流水线阶段在不同的4个TPU核心（不同主机）上运行，通过ICI将数据传输到下一阶段。通过并行异步分发，PATHWAYS在大量流水线阶段的情况下摊销了固定的客户端开销和调度开销。

5.2 多租户

实验内容：验证PATHWAYS在并发程序之间进行时间复用加速器的能力（在配置(B)上进行）。
实验结果与分析（图8, 图9）：
- 吞吐量: PATHWAYS在多个客户端并发提交不同程序时，能达到与JAX至少相同的聚合吞吐量，表明在程序间进行上下文切换没有开销（当资源能同时放入HBM时）。
- 利用率: 对于非常小的计算，PATHWAYS的最大吞吐量超过了JAX，因为它能从远程客户端接收比JAX本地Python分发更多的计算，从而实现更高的TPU利用率。
- 公平性: 图9的追踪图显示，PATHWAYS能对4个独立客户端提交的程序进行组调度，并控制加速器时间分配以实现公平性（例如，强制按比例共享）。

图8. 并发程序的聚合吞吞吐量（计算时间单位为ms）。PATHWAYS能高效地在程序间时间复用加速器，上下文切换无开销。

图9. PATHWAYS上一组核心样本的追踪图，显示了4个客户端之间按1:1:1:1（上）和1:2:4:8（下）比例共享的组调度并发程序的交错执行。

5.3 大规模模型性能

实验内容：在训练真实机器学习模型（SPMD程序）时评估PATHWAYS的性能，并与JAX和TF进行比较。数值结果验证为相同。

Text-to-Text Transformer模型 (与JAX对比) (表1):
- 模型: 使用(Raffel et al., 2019)【69，Exploring the limits of transfer learning with a unified text-to-text Transformer，2019，arXiv】中的配置，模型参数最高达110亿。
- 结果: 在所有测试的模型规模下，JAX和PATHWAYS的训练吞吐量（tokens/秒）相同，因为真实的计算规模足以掩盖单控制器的开销。

表1. Text-to-text Transformer模型在JAX多控制器和PATHWAYS上的训练吞吐量(tokens/s)。

Decoder-only Transformer语言模型 (SPMD vs. 流水线) (表2):
- 模型: 30亿参数的TF模型，包含62个Transformer层。比较SPMD配置与类似GPipe【32，GPipe: Efficient training of giant neural networks using pipeline parallelism，2019，NeurIPS】的流水线调度。
- 结果:
- PATHWAYS的训练吞吐量随每个流水线阶段的TPU核心数成比例增加。
- 增加流水线阶段数（从4到16）带来的开销极小。
- 在此实例中，流水线性能与SPMD相当，因为SPMD的集合通信开销高于流水线气泡开销。

表2. 3B Transformer语言模型在PATHWAYS上使用C个TPU核心的训练吞吐量(tokens/s)，分为SPMD或多个流水线阶段。对于流水线并行模型，有S个阶段，每个批次分为M个微批次。

跨集群训练（图10）：
- 实验内容: 验证PATHWAYS能高效地在通过DCN连接的TPU集群上训练模型。
- 结果: 使用配置(C)的4个32核集群与使用配置(B)的单个128核集群，获得了相同的吞吐量（131.4k tokens/sec）。图10的追踪图显示DCN传输时间被计算有效重叠。

图10. 3B Transformer模型在128个TPU上进行流水线训练：PATHWAYS可以在通过DCN连接的TPU集群上高效训练模型，在配置(C)的4个32核集群上实现了与配置(B)的单个128核集群相同的吞吐量(131.4k tokens/sec)。

更大规模模型训练:
- 模型: 64B和136B参数的Decoder-only Transformer模型。
- 实验: 在两个通过DCN连接的加速器集群上训练。
- 结果: PATHWAYS实现了与使用两倍设备数量的单个集群相比约97%的吞吐量。对于136B（64B）模型，在两个1024（512）核心的集群上训练，全局规约通过岛内ICI和岛间DCN传输1030GB（457GB）数据。

A7 补充细节

6 讨论

PATHWAYS的设计与实现：PATHWAYS的设计目标是大型TPU加速器集群。使用TPU而非GPU影响了许多底层设计决策。TPU与GPU最大的区别在于，TPU支持将更长运行时间和更复杂的计算融合成一个单一的TPU内核，因为它支持丰富的控制流和通信原语，而这些在GPU系统上必须由驱动代码执行。相比之下，GPU与主机内存系统和DCN的集成更紧密【65，NVIDIA GPUDirect technology，2021，http://developer.download.nvidia.com/devzone/devcenter/cuda/docs/GPUDirect Technology Overview.pdf】。TPU非常适合PATHWAYS，因为XLA可以编译包含融合集合操作的高性能函数，而大规模的高性能TPU互连允许灵活调度各种不同规模的计算。尽管如此，我们相信在PATHWAYS中做出的大部分高级架构选择，对于大规模GPU系统同样有效。

资源管理：PATHWAYS旨在允许各种细粒度的动态资源管理策略。我们初步的研究集中在TPU计算的高效动态时间复用上。对于未来更复杂的多租户用例，PATHWAYS需要处理更多样化的资源类型，包括但不限于设备和主机内存，以及ICI、DCN和PCIe带宽。PATHWAYS的单控制器模型赋予系统大规模跟踪可用资源和分配资源的能力。我们计划探索常见的多租户需求，如优先级、性能隔离、访问控制和资源核算，但时间尺度比以往工作小得多，且资源池规模大几个数量级。

数据依赖的向量化控制流：目前几乎所有ML模型在每一步训练中都根据每个训练样本更新所有模型权重。我们希望通过细粒度的控制流来支持研究，使得可以根据每个样本，甚至子样本（如图像的补丁或句子的单词）更新不同的模型权重。像专家混合（MoE）【79，Outrageously large neural networks: The sparsely-gated mixture-of-experts layer，2017，ICLR】和路由胶囊网络【31，Matrix capsules with EM routing，2018，ICLR；8，Machine learning systems are stuck in a rut，2019，HotOS】这样的模型，通过根据学习函数将不同的（子）样本“路由”到托管不同模型权重子集的加速器来利用计算稀疏性。这种路由需要节点之间细粒度的数据依赖数据交换。我们的ML研究同事表示，他们希望在训练更大模型、处理更多任务时更有效地利用稀疏性，但现有框架限制了他们对新颖模型架构的实验。支持具有简洁编程模型和良好性能的数据依赖向量化控制流是未来的工作。

7 相关工作

超越SPMD多控制器的需求：我们在§2中详细探讨了紧密相关的工作。本节扩展讨论了那些需要超出SPMD多控制器所能提供能力范围的ML工作负载的相关研究，并验证了我们的PATHWAYS设计选择。

资源共享与虚拟化：
- 粗粒度共享: 传统的资源共享是粗粒度的。云提供商通常将加速器专用于单个用户，而集群调度器虽然优化了异构性和公平性【62，Heterogeneity-aware cluster scheduling policies for deep learning workloads，2020，OSDI；88，Gandiva: Introspective cluster scheduling for deep learning，2018，OSDI】，但资源仍以秒级或更长的时间尺度专用于单个作业。
- 细粒度共享: 近期工作表明，更细粒度的共享可以进一步提高资源效率。这包括虚拟化加速器【90，AvA: Accelerated virtualization of accelerators，2020，ASPLOS】、GPU内存虚拟化【73，vDNN: Virtualized deep neural networks for scalable, memory-efficient neural network design，2016，MICRO】或DRAM卸载【70，DeepSpeed: System optimizations enable training deep learning models with over 100 billion parameters，2020，ACM SIGKDD】，以及并发ML任务执行【53，Zico: Efficient GPU memory sharing for concurrent DNN training，2021，USENIX ATC】。这些细粒度共享技术展示了共享加速器的机会，但若没有像PATHWAYS这样的单控制器系统，这些机会很难大规模利用。

非SPMD计算模型：
- 许多工作表明，偏离SPMD计算可以在大型工作负载上提高效率。
- 流水线并行【32，GPipe: Efficient training of giant neural networks using pipeline parallelism，2019，NeurIPS；61，PipeDream: Generalized pipeline parallelism for DNN training，2019，ACM SOSP；89，Pipemare: Asynchronous pipeline parallel DNN training，2021，MLSys】将ML模型划分为跨加速器的静态异构计算。
- 固有异构任务: 图神经网络训练【38，Improving the accuracy, scalability, and performance of graph neural networks with Roc，2020，MLSys】、神经架构搜索【69，Exploring the limits of transfer learning with a unified text-to-text Transformer，2019，arXiv】以及多模态多任务学习系统【54，Modeling task relationships in multi-task learning with multi-gate mixture-of-experts，2018，ACM SIGKDD】等，都是本质上异构和动态的任务，不自然地适应SPMD模型。
- 未来模型: 我们预计，即将到来的大规模高效ML模型可能会形成一个由共享层和专属层组成的集合【11，On the opportunities and risks of foundation models，2021，arXiv】，这很自然地可以表示为MPMD。

A5 结论

PATHWAYS在当前的单租户SPMD机器学习模型上，实现了与最先进的多控制器系统相当的性能。我们确保了与多控制器JAX的严格兼容性，并且如§5所示，PATHWAYS在极大规模的系统上，除了最小的计算外，其性能均与JAX相匹配。

同时，PATHWAYS颠覆了JAX程序的执行模型，将用户代码拉回到单控制器模型中，并在客户端和加速器之间引入了一个集中的资源管理和调度框架。单控制器编程模型使用户能够简单地访问更丰富的计算模式。资源管理和调度层允许重新引入集群管理策略，包括多租户共享、虚拟化和弹性，所有这些都针对ML工作负载和加速器的需求量身定制。我们的微观基准测试显示了并发客户端工作负载的交错执行和高效的流水线执行，有力地证明了我们构建的系统机制是快速和灵活的，为研究利用这些机制的新颖策略奠定了坚实的基础。

我们已经证明，通过精心的系统设计和工程实现，我们可以“两全其美”，在当今的ML模型上匹配性能，同时提供编写未来模型所需的功能。

A6 附录

A 加速器设计考量

硬件加速的重要性：硬件加速对现代深度学习至关重要，但实现高性能是一个复杂的系统工程。以下是深度学习系统中为实现良好性能而普遍采用的既定技术。

A.1 批处理（Batching）：由于Dennard缩放定律的终结，加速器采用硬件并行（如SIMT或收缩阵列）。这使得内存带宽成为关键资源，需要使用昂贵且容量有限的高带宽内存（HBM）。批处理通过解锁并行性和实现内存重用（一个浮点数从内存中读取一次用于多次计算）来缓解带宽压力。然而，批处理会增加HBM内存压力，且过大的批量可能会减慢模型收敛速度。
A.2 异步编程：加速器抽象依赖异步编程模型来实现高性能。同步抽象会因PCIe延迟、内核调度开销和中断延迟而浪费大量计算资源。计算被放入流（streams）中以备将来在加速器上执行，只要维持足够大的工作流水线，这种异步抽象就能有效掩盖小型操作的分发延迟。
A.3 高性能互连：现代深度神经网络的大小远超单个加速器HBM内存的容量，因此需要将模型分片到多个加速器上，这使得加速器之间的高速互连至关重要。GPU使用NVLink连接小范围内的加速器，并使用RDMA在集群间通信。TPU则内置了自定义网状网络，芯片间可直接通信。这些专用互连通常通过MPI原语（如AllReduce）暴露给应用程序，且必须进行组调度。
A.4 单租户（Single-tenancy）：加速器通常不像计算机中的其他资源那样被多个程序同时共享。深度学习模型易于通过增加参数或批量大小来扩展内存使用，因此程序实际上会消耗大部分可用的HBM内存。由于PCIe带宽远小于HBM或加速器互连带宽，细粒度的上下文切换（将大量HBM数据换出到主机DRAM）会导致大量加速器周期浪费。因此，当一个主机程序未充分利用加速器时，计算资源就被闲置了。
A.5 GPU与TPU的对比：
- 互连与编程: GPU系统通常有小规模的NVLink连接集群（例如，一台主机内8个GPU），通过Infiniband或DCN连接更大的聚合体。GPU编程通过向加速器分发许多小的预编译“内核”来进行。TPU系统则有数千个设备全互连，通过XLA执行长时间运行的函数，这些函数可包含跨ICI网络的集合通信。
- 编译与优化: 由于TPU上的XLA计算运行时间远长于GPU内核，编译器可以进行更多优化，如静态缓冲区分配和自动重物质化。这导致TPU对动态形状的支持有限，非常适合PATHWAYS的常规已编译函数概念。
- 抢占与调度: TPU一次只能运行一个程序，没有本地抢占，这使得集中式组调度至关重要，以避免死锁。虽然GPU可以并发执行程序而无需集中调度，但对于以吞吐量为重的ML训练工作负载，将整个GPU或其静态部分专用于单个精心大小的计算，比让GPU驱动和硬件运行时动态复用其资源更有效。因此，类似PATHWAYS的设计对GPU效率也有好处。

B 典型ML程序的结构

“已编译函数”（Compiled Functions）：现代ML工作负载中加速器执行的计算主要由我们称之为“已编译函数”的子计算构成。这些子计算具有以下特点：
- 静态已知的元数据: 输入输出类型以及任何输入/输出张量的形状，在输入数据计算出来之前就已知。
- 有界循环: 任何循环的边界在节点计算被调度时已知，或指定为具有潜在提前终止的最大迭代次数。
- 函数式条件: 条件是“函数式”的，即两个分支具有相同的输出类型，并且会提前为任一分支分配足够的资源。

对框架的影响：由于已编译函数的资源需求是预先知道的，现代ML框架利用这一特性，在前辈节点运行之前异步地将已编译函数入队，允许主机端工作与加速器计算并行进行。框架尽可能将已编译函数图提交给“即时”（JIT）编译器，以利用布局分配和融合等优化，从而显著提高生成的加速器代码的效率。这使得尽管客户端代码可能用高级语言编写，但性能敏感的节点计算通常被降低为可序列化并易于发送到远程主机执行的内部表示（IR）。

C 输入数据处理

利用TensorFlow进行数据处理：JAX有意避免重新实现数据加载流水线，而tensorflow/datasets【83，TensorFlow Datasets: A collection of ready-to-use datasets，2021，https://www.tensorflow.org/datasets】通常用于JAX的输入处理，因此JAX程序很容易适应将输入处理卸载到在PATHWAYS工作节点上运行的基于CPU的TensorFlow执行器。PATHWAYS在每个主机上实例化一个基于CPU的TensorFlow执行器，以便用户程序可以将输入处理序列化为一个TensorFlow图，并将其分布在工作节点上。我们计划支持流式数据协议，以便基于CPU的计算可以在一组独立管理的服务器上执行，从而将昂贵的连接TPU的主机与可用于输入处理的CPU资源解耦。

D 评估工作负载追踪

多租户工作负载追踪（图11）：图11展示了图8中工作负载的追踪情况，其中不同数量的客户端并发提交程序（§5.2）。单个客户端的每个程序计算时间非常短（0.33ms），不足以使加速器饱和。借助PATHWAYS的多租户支持，使用多个客户端可将设备利用率提高到约100%。所有客户端程序都在所有核心上进行组调度，并以毫秒级或更小的时间尺度交错执行，显示出很小的上下文切换开销。

图11. 图8中TPU核心样本的追踪图。PATHWAYS展示了组调度并发程序的交错执行。

跨集群训练追踪（图12）：图12显示了当64B Decoder-only Transformer模型在两个各含512个芯片的加速器集群上进行数据并行训练时，多个训练步骤的追踪剖面图（§5.3）。前八行（蓝色）对应第一个集群中一台主机的TPU计算，后八行（绿色）对应第二个集群中一台主机的TPU计算。在这种情况下，每个集群计算梯度，然后将这些梯度的传输操作入队到另一个集群。当通过DCN的梯度传输完成时，每个集群应用接收到的梯度并开始下一个训练步骤。即使在每对128台主机的规模上，DCN传输也只产生最小的开销，与使用ICI通信的总计相同数量芯片的SPMD配置相比，训练吞吐量达到了97.2%。

图12. 64B Transformer模型在两个各含512个TPU的集群上进行数据并行训练。该追踪图突显了使用DCN进行跨集群传输的开销相对较小。

PATHWAYS: ASYNCHRONOUS DISTRIBUTED DATAFLOW FOR ML