文章：PATHWAYS: ASYNCHRONOUS DISTRIBUTED DATAFLOW FOR ML
机器学习的异步分布式数据流

XLA(Accelerated Linear Algebra)-加速线性代数，是Google推出的高性能机器学习领域编译器，它可以在不更改源代码的条件下加速Tensorflow模型。

PLAQUE is an existing (closed-source) production sharded dataflow system used at Google for many customer-facing services where high-fanout or high-fanin communication is necessary, and both scalability and latency are important.

摘要：我们提出了一种新的用于加速器的大规模编配层的设计。我们的系统，PATHWAYS，明确地设计用于探索新的系统和ML研究思想，同时保持当前模型的最先进的性能状态。PATHWAYS使用异步操作符的分片数据流图，这些操作符消费和产生futures，并在数千个加速器上高效地对异构并行计算进行分组调度，同时在它们专用的互连上协调数据传输。PATHWAYS使用了一种新颖的异步分布式数据流设计，它允许控制平面并行执行，尽管数据平面中存在依赖关系。这种经过精心设计的设计，允许PATHWAYS采用单控制器模型，从而更容易表达复杂的新并行模式。我们证明，当运行超过2048个tpu的SPMD计算时，在最先进的系统中，PATHWAYS可以实现性能平价(100%的加速器利用率)，同时还提供了可与跨越16个阶段的Transformer模型的SPMD案例相比的吞吐量。或者被分割成两个通过数据中心网络连接的加速器岛

1. 介绍

在过去的十年里，深度学习在图像理解等领域取得了显著的成就(Krizhevsky et al.， 2012;他等人，2016)到自然语言处理(Devlin等人，2019;Brown等人，2020年)。机器学习(ML)的这种快速发展的特点是ML模型、加速器硬件和将两者联系在一起的软件系统的共同进化。这种协同进化带来了系统对当前工作负载过于专一的危险，并且无法预测未来的需求。在这篇文章中,我们描述通路,新系统分布式毫升。通路设计目标需要特定的能力,我们相信通过未来毫升工作负载(院长,2021),因此需要今天支持研究那些糟糕的工作负载,但是支持的最先进的系统

例如，目前大多数最先进的ML工作负载都使用受MPI启发的单程序多数据(SPMD)模型(Clarke et al.， 1994)，其中所有加速器都以同步方式运行相同的计算，加速器之间的通信由AllReduce等集体描述。最近遇到了SPMD在ML计算的限制。非常大型的语言模型已经使用流水线而不是纯粹的数据并行进行了扩展(Narayanan等人，2019;Rasley等人，2020年;Narayanan等人，2021年)和混合专家(MoE)等模型(Shazeer等人，2017年)已经开始探索计算稀疏性，最自然的表达方式是使用细粒度控制流和跨加速器的异构计算。系统设计者采用了巧妙的技术来执行流水线操作(Narayanan等人，2021年;Rasley等人，2020年;Narayanan等人，2019年;Huang et al.， 2019)和齐次MoE (Lepikhin et al.， 2020;Fedus等人，2021)的MPI风格系统模型，但正如我们稍后将详细讨论的那样，MPI编程模型对用户和底层系统都太有限制了。

另一方面，随着每一代新加速器的出现，ML集群正变得越来越异质性(Jeon等人，2019;Chaudhary等人，2020年;翁等，2022)。提供对通过高带宽互连连接的同质加速器的大岛的独家访问是昂贵的，而且通常是浪费的，因为单个用户程序必须试图保持所有加速器持续忙碌。这些限制进一步推动研究人员转向多程序多数据(MPMD)计算，通过将整体计算的子部分映射到更容易获得的更小的加速器岛的集合，从而实现更大的灵活性。为了提高资源利用率，一些ML硬件资源管理研究者在工作负载之间以细粒度的方式复用硬件，使工作负载具有弹性，并提高容错能力

最后，研究人员开始将一组基础模型标准化(Bommasani等人，2021年;Dean, 2021)，在大规模的大数据上训练，并适应多个下游任务。这种模型的训练和推理为通过跨多个任务复用资源和在任务之间有效地共享状态来提高集群利用率提供了机会。例如，几个研究人员可能同时进行微调(Houlsby等人，2019;Zhang等人，2021)的不同任务的基础模型，使用相同的加速器来保持固定的基础模型层。对于共享子模型的训练或推断可以受益于允许将来自不同任务的示例组合到单个向量化批处理中以获得更好的加速器利用率的技术(Crankshaw等人，2017)。

本文描述了我们的系统，PATHWAYS，它与最先进的ML系统的功能和性能相匹配，同时提供支持未来ML工作负载所需的功能。PATHWAYS使用了一个客户端-服务器架构，使得PATHWAYS的运行时能够代表许多客户端在系统管理的计算孤岛上执行程序。PATHWAYS是第一个旨在透明和高效地执行跨越多个tpu单元的程序的系统(谷歌，2021)，它通过采用新的数据流执行模型，可以扩展到数千个加速器。PATHWAYS的编程模型使得表达非spmd计算变得很容易，并支持集中的资源管理和虚拟化，以提高加速器的利用率。

在本文的其余部分中，我们首先讨论当前分布式ML系统的局限性，并激励我们对路径(2)的设计选择，然后描述路径支持的灵活编程模型(3)。我们描述了路径的架构(4)，强调了我们如何使用分片数据流模型和异步组调度来解决旧的客户机-服务器ML系统的关键限制。我们展示了微基准和使用真实的ML模型的端到端评估，证明我们已经达到了与最先进的多控制器在现实工作负载下的性能匹配的目标(5)，并验证PATHWAYS机制非常适合支持研究和部署新型高效ML方法所需的功能

2. 设计动机

分布式ML系统的设计选择通常是由底层目标硬件加速器的属性驱动的。附录A描述这些特性如何影响分布式ML系统。在这里，我们将重点讨论现有的分布式ML系统的一些设计和实现选择是如何使它们难以支持大型、稀疏或不规则的模型的。

用于训练最先进的SPMD模型的分布式ML系统通常采用多控制器架构，其中同一客户机可执行文件直接运行在系统中的所有主机上，在程序执行期间独占这些主机上的资源。这种架构的例子包括MPI (Clarke et al.， 1994)、PyTorch (Paszke et al.， 2019)、JAX (Bradbury et al.， 2018)和TensorFlow的最新配置(Shazeer et al.， 2018;Agrawal等人，2019)。这种架构的主要优点是调度加速器计算的低延迟(见图1a)，因为在每个加速器主机上运行相同的用户代码副本，而调度只涉及(相对)快速的PCIe链路上的通信。所有其他主机间的通信只通过使用专用互连的集体进行，如NVLink (Foley and Danskin, 2017)和ICI (Jouppi等人，2020)，而不通过主机内存。然而，这种体系结构并不适合使用管道或计算稀疏性的现代ML工作负载。在多控制器系统中，任何超出标准集合的通信都需要用户实现自己的协调原语。多控制器方法通常还假定硬件资源的独占所有权。这不仅将确保昂贵的加速器的高利用率的责任转移到用户身上，而且还使构建高效集群范围ML基础设施所需的资源虚拟化和多路复用等特性的设计变得复杂。

单控制器系统，如TensorFlow v1 (Abadi et al.， 2016)提供了一个非常通用的分布式数据流模型，包括优化的图内控制流(Yu et al.， 2018)。TensorFlow (TF) Python客户端构建一个计算图，并将其传递给一个协调运行时，协调运行时将图划分为每个工作人员的子图，并将子图的执行委托给工作人员的本地运行时。工作者之间的协调是通过数据中心网络(DCN)上的数据和控制边传递消息来完成的。虽然单控制器设计提供了灵活的编程模型和资源虚拟化，但它也带来了实现方面的挑战。

首先，虽然多控制器系统只需要通过PCIe通信来调度加速计算(图1a)，但单控制器系统中的客户端距离更远，而且调度延迟涉及到DCN通信，通常比PCIe慢一个数量级(图1b)。其次，支持MPMD程序与SPMD子计算并行执行，每跨越一个共享集群的加速器子集，运行时必须有某种机制来支持加速器计算的批量调度。在tpu的情况下，群调度是必不可少的，因为它们是单线程的，并且只运行不可抢占的内核，所以如果通信计算没有按照一致的顺序排队，系统将会死锁。即使是gpu或其他能够执行并发计算的加速器，群调度也能够更有效地执行集体(Feitelson和Rudolph, 1992)。因此，ML的单控制器系统需要一个分布式调度机制，以顺序的计算排队代表不同的程序。最后，现代ML工作负载的系统必须设计为运行分布在数千个加速器上的计算，并对分片表示和数据结构提供一流的支持。例如，一个简单的数据流图表示M路分片计算和N路分片计算之间的一条边，将需要M + N个节点和M N条边，这很快就变得难以处理。

TF v1所做的实现选择过于专业化，以至于假设了一个单独的、较小的、独占的加速器岛。这种过度专门化使得将TF用于当前或未来的ML工作负载实际上是不可行的。虽然TF可以通过send和recv ops运行需要跨主机协调或数据传输的计算(图1c)，但目的地的主机侧工作(如调度加速器计算)只有在传输完成后才会触发。在涉及许多跨主机传输的程序中，例如具有大量阶段的流水线模型，这些分派延迟会累积起来，导致低效的加速器利用率。虽然TF v1用户可以(低效地)通过使用控制边，在单个程序中对群调度执行一致的顺序，但在像TF v1这样的单控制器系统中，缺乏集中的调度程序，因此不可能确保跨程序计算之间的一致顺序。TF还实现了完整的分片计算图，当分片的数量达到数千时，在图序列化和执行中都会引入大量开销，导致子计算之间有数百万条图边。

PATHWAYS结合了单控制器框架的灵活性和多控制器的性能。我们采用了单控制器模型，因为我们相信，通过利用计算稀疏性和异构性，以及通过启用促进资源共享和虚拟化的集群管理系统，单控制器模型为新型和高效的ML计算提供了比多控制器更好的机会。我们的设计不同于以往的单控制器ML系统，它使用异步调度来匹配多控制器系统的性能，支持集中的资源管理和调度，一流的支持SPMD加速器分组计算，并使用分片的数据流系统来高效协调。

3. Pathway编程模型

我们已经实现了对用TensorFlow和JAX编写的源程序的目标路径的支持，但在本文中，我们主要关注JAX的评估。JAX用户可以显式地用装饰器包装标准Python代码，以指示应该编译成(可能是SPMD) XLA计算的片段。这些XLA计算通常以已知的输入和输出类型和形状、有界循环为特征，并带有很少(如果有的话)的条件(请参阅附录B了解更多细节)，这使得提前估计计算的资源需求成为可能。我们将这些已知资源需求的计算称为编译函数。每个这样的函数映射到路径程序中的单个(分片)计算节点。

今天，JAX不能扩展到单个TPU pod之外，因为在多控制器配置中运行的JAX程序使用XLA集合传输所有数据，而这些集合目前只能通过TPU上的ICI使用。路径可以作为JAX后端插件的替代品，允许JAX代码不加修改地运行，除非SPMD计算现在不仅可以访问本地连接的TPU核，还可以访问系统中提供的任意数量的核。由于PATHWAYS可以通过ICI和DCN进行通信，它允许JAX程序第一次扩展到多个TPU pods，其中包含数千个TPU核。

运行未经修改的JAX代码的能力很方便，但并不能释放pathway的全部性能。pathway用户可以请求一组虚拟设备，在设备类型、位置或互连拓扑上有可选的约束，然后能够在这些设备上放置特定的编译功能(图2)。系统将自动处理所有的数据移动，并在依赖计算之间重新分片。

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

JAX支持跟踪代码转换的理念与我们想要探索的研究方向非常匹配。例如，JAX有一个名为FLAX的配套库(Heek等人，2020)，用于表达分层DNN模型，我们已经编写了一个库，自动将FLAX模型转换为流水线路径程序。此外，JAX支持向量化每个示例的Python函数的转换，从而产生高效的批处理代码，而且这种转换是探索依赖数据的向量化控制流新形式的良好基础，我们将在后面简要描述(6.3)

4. Pathway系统架构

PATHWAYS广泛地构建在以前的系统上，包括XLA (TensorFlow, 2019)来表示和执行TPU计算，TensorFlow图和执行器来表示和执行分布式CPU计算，以及Python编程框架，包括JAX (Bradbury et al.， 2018)和TensorFlow api。通过利用这些构建块，我们能够专注于新颖的PATHWAYS协调方面，同时能够以最小的代码更改运行现有的ML模型。

4.1资源管理器
pathway后端由一组加速器组成，这些加速器被分组成紧密耦合的islands，这些岛屿通过DCN相互连接(图3)。路径有一个资源管理器，负责集中管理所有岛屿上的设备。客户可以要求具有特定的2D或3D网格形状的虚拟岛屿切片，以适应他们的通信模式。每个虚拟片包含虚拟设备，允许客户端表达计算是如何在网格上布局的。资源管理器动态地为满足所需互连拓扑、内存容量等要求的虚拟设备分配物理设备。

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

路径允许后端计算资源动态地添加和删除，资源管理器跟踪可用的设备。虚拟和物理设备之间的间接层,通过我们single-controller设计,将使我们在未来的支持功能,如透明的挂起/恢复和迁移,在客户年代虚拟设备暂时收回或重新分配从用户程序不需要合作。

4.2 Client客户端

当用户想要运行一个跟踪程序时，它会调用PATHWAYS客户端库，该客户端库首先为任何以前没有运行过的计算分配虚拟设备，然后将计算注册到资源管理器，触发服务器在后台编译计算。然后，客户端为程序构建一个设备位置无关的路径中间表示(IR)，表示为一种自定义的MLIR (latner等人，2021年)方言。IR通过一系列标准编译器的传递逐步降低，最终输出包含物理设备位置的低级表示。这个低级程序考虑到物理设备之间的网络连通性，并包括操作将输出从源计算分片传输到目标分片的位置，包括需要数据交换时的分散和收集操作。在通常情况下，虚拟设备的位置不会改变，如果资源管理器改变了虚拟设备和物理设备之间的映射，可以重新降低程序，重复运行低级程序是有效的。

旧的单控制器系统中的客户端可能很快成为性能瓶颈，因为它协调了数千个单独的计算，以及分布在数千个加速器上的每个计算碎片对应的数据缓冲区。Pathway客户端使用一个分片缓冲区抽象来表示一个逻辑缓冲区，这个逻辑缓冲区可以分布在多个设备上。这种抽象通过在逻辑缓冲区的粒度上分摊簿记任务的成本(包括引用计数)，而不是在单个的分片粒度上，帮助客户端扩展。

4.3 Coordination实现

所有使用DCN的跨宿主协调都依赖于PLAQUE。PLAQUE是一个现有的(闭源)生产分片数据流系统，在谷歌中使用，用于许多需要高fanout或高fanin通信的面向客户的服务，并且可伸缩性和延迟都很重要。low-level的PATHWAYS IR直接转化为PLAQUE程序，用数据流图表示。pathway对其协调底层有严格的要求，所有这些都由PLAQUE满足。

首先,表示用于描述pathway IR必须为每一个分片计算包含一个节点,以确保一个紧凑的表示计算跨许多碎片,即链接执行2 a和B计算N计算碎片数据流中的每个应该有4个节点表示

协调运行时还必须支持沿着分片边缘的稀疏数据交换，其中消息可以在动态选择的分片子集之间发送，使用标准的进度跟踪机制(Akidau等，2013;Murray等人，2013)来检测一个分片的所有消息何时已被接收。为了避免DCN成为加速器上依赖数据的控制流的瓶颈，我们需要高效的稀疏通信，这是我们希望路径能够实现的关键功能之一。

协调基板用于发送DCN消息，这些消息位于传输调度消息和数据句柄的关键路径上(图4)，因此它必须发送低延迟的关键消息，以及在需要高吞吐量时发送到同一主机的批处理消息。

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

我们认为，利用Ray (Moritz et al.， 2018)等其他分布式框架，而不是PLAQUE，重新实施完整的pathway设计是可行的，以实现低层次的协调框架。在这样的实现中，路径执行器和调度器将被长期运行的Ray Actor所取代，它们将在底层的Ray集群调度之上实现路径调度，并且执行器可以使用PyTorch来进行GPU计算和集合。为了获得类似的性能(见5)，可能需要添加一些内容，因为Ray缺少HBM对象存储，或者通过GPU互连有效地传输远程对象的原语。

4.4 Gang-scheduled dynamic dispatch

如前所述(2)，在一组共享加速器上支持SPMD计算的一个需求是支持高效的组调度。PATHWAYS运行时为每个岛包括一个集中的调度器，该调度器一致地对岛上的所有计算排序。当PATHWAYS排队程序执行时，PLAQUE数据流程序负责(i)排队本地编译函数的执行在每个加速器，以buffer futures作为输入;(ii)通过函数执行将网络发送到远程加速器的缓冲期货输出排队;(iii)与调度程序通信，以确定在岛上运行的所有程序的函数执行的一致顺序。调度器必须实现以毫秒为时间尺度分配加速器的策略。我们目前的实现只是简单地按照FIFO顺序排列工作，但更复杂的调度器可能会根据估计的执行时间重新排序计算。

4.5 并行异步分发

当在加速器上运行计算时，系统可以利用异步api将计算与协调重叠。考虑图4a中的三节点图，其中的方块对应于三个节点A、B和C，它们运行在连接到主机A、B和C的加速器上。所有节点计算都是常规编译的函数。主机A向节点A排队，接收A的输出，并将输出发送给主机B。主机B分配B的输入，并将输入缓冲区地址发送给主机A，完成启动节点B功能的大部分准备工作。当节点A完成时，它的输出通过加速器互连直接发送到节点B的输入缓冲区，然后主机B启动节点B。一个节点完成和下一个节点启动之间的延迟可以被设置为比数据传输时间稍长一点

如果上一级节点的计算时间超过主机之间的调度、资源分配、协调时间，则可以采用上述设计。但如果计算时间过短，如图所示，则异步管道会停滞，主机端工作成为整个计算序列执行的关键瓶颈。考虑到编译后的函数都是正则的，实际上可以在前一个节点的计算进入队列之前计算出后续节点的输入形状

因此我们引入一种新颖的并行异步调度设计如图4b所示,利用常规编译函数中已知的静态资源使用来运行大多数日志并行计算的节点,而不是序列化为一个节点工作发生后其前辈已经加入队列。由于工作只能在功能正常的情况下并行调度，PATHWAYS将并行调度视为一种优化，当一个节点的资源需求直到一个前一个计算完成时才知道时(例如，由于依赖数据的控制流)，就会退回到传统模型。当计算的一个子图可以被静态调度时，程序向调度器发送一条消息(描述整个子图)，调度器能够将子图中所有活动分片的执行顺序排到后面。使用单个消息的目的是最小化网络流量，但不需要调度器将所有的子图碎片作为一个批处理队列:计算仍然可能与其他并发执行的程序提交的计算交叉。在5中，我们对不同调度机制的成本进行了评估。

4.6 数据管理

每个主机管理一个分片对象存储，类似于Ray的对象存储(Moritz等人，2018)，但也扩展到跟踪每个分片上的加速器HBM中保存的缓冲区。客户机程序可以保存对远程主机或加速器内存中的对象的引用，客户机和服务器使用不透明句柄引用它们，允许系统在需要时迁移它们。中间程序值也保存在对象存储中，例如，当系统等待在加速器之间传输它们或将它们传递给后续的计算时。这些对象都带有所有权标签，以便在程序或客户端失败时可以对其进行垃圾收集。如果计算不能分配内存，因为其他计算缓冲区正在临时占用HBM，我们可以使用简单的背压来停止计算。