PATHWAYS: 用于机器学习的异步分布式数据流

摘要

我们提出了一种新的大规模加速器调度层的设计。我们的系统，PATHWAYS，专门设计用于启用新系统和机器学习研究思想的探索，同时保持当前模型的最先进性能。PATHWAYS 使用一个分片的数据流图，包含异步操作符，这些操作符消耗和生成未来值，并高效地在成千上万的加速器上进行并行计算，同时协调通过其专用互连的数据传输。PATHWAYS 利用一种新颖的异步分布式数据流设计，使得控制平面能够并行执行，尽管数据平面存在依赖关系。通过精心的工程设计，这种设计使得 PATHWAYS 采用单控制器模型，从而更容易表达复杂的新并行模式。我们展示了，PATHWAYS 在运行 SPMD 计算时，能够在 2048 个 TPUs 上实现与最先进系统相当的性能（大约 100% 的加速器利用率），同时在 Transformer 模型的情况下，提供与 SPMD 相当的吞吐量，这些模型通过 16 个阶段进行流水线处理，或跨两个加速器岛屿进行分片，并通过数据中心网络进行连接。

1 引言

在过去的十年里，深度学习在各个领域取得了显著的成就，从图像理解（Krizhevsky 等，2012；He 等，2016）到自然语言处理（Devlin 等，2019；Brown 等，2020）。机器学习（ML）的这一快速进展，得益于 ML 模型、加速器硬件和将两者联系起来的软件系统的共同演进。这种共同演进带来了一定的风险，即系统可能会过度专业化，适应当前的工作负载，而无法预见未来的需求。在本文中，我们描述了 PATHWAYS，这是一个为分布式 ML 构建的新系统。PATHWAYS 旨在满足我们认为未来 ML 工作负载所需的特定功能（Dean，2021），因此这些功能也正是今天支持这些工作负载研究所需要的，但目前的先进系统对其支持不足。

例如，目前大多数最先进的 ML 工作负载使用“单程序多数据”（SPMD）模型，这一模型受到了 MPI（Clarke 等，1994）启发，其中所有加速器以同步的方式执行相同的计算，加速器之间的通信通过像 AllReduce 这样的集体操作来描述。最近，研究人员开始遇到 SPMD 在 ML 计算中的局限性。非常大的语言模型已通过流水线技术进行扩展，而不是纯粹的数据并行（Narayanan 等，2019；Rasley 等，2020；Narayanan 等，2021），而像 Mixture of Experts（MoE）（Shazeer 等，2017）这样的模型开始探索计算稀疏性，最自然的表达方式是使用细粒度的控制流和加速器之间的异构计算。系统设计者已经采用巧妙的技术来执行流水线（Narayanan 等，2021；Rasley 等，2020；Narayanan 等，2019；Huang 等，2019）和同质化的 MoE（Lepikhin 等，2020；Fedus 等，2021）模型，在 MPI 风格的系统上进行计算，但正如我们在后文详细论述的那样，MPI 编程模型对于用户和底层系统来说都过于限制。

另一方面，随着每一代新加速器的出现，ML 集群变得越来越异构（Jeon 等，2019；Chaudhary 等，2020；Weng 等，2022）。为大型“岛屿”上的同质加速器提供独占访问权限，且这些加速器通过高带宽互联连接是非常昂贵的，而且通常会浪费，因为单个用户程序必须试图保持所有加速器持续繁忙。这种约束进一步推动了研究人员朝着“多程序多数据”（MPMD）计算发展，这使得通过将整体计算的子部分映射到更多现成的较小加速器岛屿上，从而提供更多灵活性。为了提高利用率，一些 ML 硬件资源管理的研究人员（Xiao 等，2020；Bai 等，2020；Yu 和 Chowdhury，2020；Wang 等，2021；Lim 等，2021；Zhao 等，2022；Weng 等，2022）采用细粒度的硬件多路复用技术在不同工作负载之间分配资源，从而实现工作负载弹性并提高容错性。

最后，研究人员开始对一组基础模型（Bommasani 等，2021；Dean，2021）进行标准化，这些模型在大规模数据上进行训练，并能够适应多个下游任务。对于这样的模型的训练和推理提供了提高集群利用率的机会，通过跨多个任务多路复用资源，并有效地共享其状态。例如，几位研究人员可能会并行地微调（Houlsby 等，2019；Zhang 等，2021）同一基础模型的不同任务，使用相同的加速器来存储固定的基础模型层。跨共享子模型的训练或推理可以受益于允许不同任务的示例在一个单一的矢量化批次中组合的技术，从而获得更好的加速器利用率（Crankshaw 等，2017）。

本文描述了我们的系统 PATHWAYS，该系统在提供与当前最先进 ML 系统相匹配的功能和性能的同时，还提供了支持未来 ML 工作负载所需的功能。PATHWAYS 使用客户端-服务器架构，使得 PATHWAYS 的运行时可以代表多个客户端在系统管理的计算岛屿上执行程序。PATHWAYS 是第一个设计用来透明且高效地执行跨多个“TPU pod”（Google，2021）程序的系统，它通过采用新的数据流执行模型，扩展到数千个加速器。PATHWAYS 的编程模型使得表达非 SPMD 计算变得更加简单，并且能够实现集中式资源管理和虚拟化，从而提高加速器的利用率。

在本文的其余部分，我们首先讨论当前分布式 ML 系统的局限性，并阐述我们对 PATHWAYS 设计选择的动机（§2），接着描述 PATHWAYS 支持的灵活编程模型（§3）。我们介绍 PATHWAYS 的架构（§4），重点说明我们如何通过使用分片的数据流模型和异步群体调度，解决传统客户端-服务器 ML 系统的关键限制。我们展示了微基准测试和使用真实 ML 模型的端到端评估，证明我们已实现与当前最先进的多控制器系统在现实工作负载上的性能匹配（§5），并验证 PATHWAYS 的机制非常适合支持研究和部署新颖且高效的 ML 方法所需的特性。

2 设计动机

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

用于训练最先进 SPMD 模型的分布式 ML 系统通常采用多控制器架构，其中相同的客户端可执行程序直接在系统中的所有主机上运行，在程序执行期间独占这些主机上的资源。这种架构的例子包括 MPI（Clarke 等，1994）、PyTorch（Paszke 等，2019）、JAX（Bradbury 等，2018）和 TensorFlow（Shazeer 等，2018；Agrawal 等，2019）等更现代的配置。这种架构的主要优势在于调度加速器计算的低延迟（见图 1a），因为用户的代码副本在每个加速器主机上运行，并且调度仅涉及通过（相对）快速的 PCIe 链路进行通信。主机之间的所有其他通信仅通过使用专用互联如 NVLink（Foley 和 Danskin，2017）和 ICI（Jouppi 等，2020）的集体操作来实现，而无需通过主机内存。然而，这种架构与现代 ML 工作负载（使用流水线或计算稀疏性）不匹配。多控制器系统中任何超出标准集体操作的通信，都要求用户实现自己的协调原语。多控制器方法通常也假设硬件资源的独占所有权。这不仅将确保昂贵加速器高效利用的责任转移给用户，而且还使得设计诸如资源虚拟化和多路复用等特性变得更加复杂，而这些特性是构建高效的集群级 ML 基础设施所必需的。

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

首先，尽管多控制器系统只需要通过 PCIe 来调度加速器计算（图 1a），但单控制器系统中的客户端“距离”更远，调度延迟涉及通过 DCN 进行通信，通常比 PCIe 慢一个数量级（图 1b）。其次，为了支持具有 SPMD 子计算的 MPMD 程序的并发执行，且每个程序跨越从共享集群中抽取的一部分加速器，运行时必须有某种机制来支持加速器计算的群体调度。在 TPUs 的情况下，群体调度是必不可少的，因为它们是单线程的，只运行不可抢占的内核，因此如果通信计算未按一致顺序入队，系统将死锁。即使对于可以执行并发计算的 GPU 或其他加速器，群体调度也能更高效地执行集体操作（Feitelson 和 Rudolph，1992）。因此，ML 的单控制器系统需要一种分布式调度机制，以便为不同程序代表的计算队列排序。最后，现代 ML 工作负载的系统必须设计成能够在数千个加速器上分布式运行计算，并为分片表示和数据结构提供一流的支持。例如，表示 M-way 分片计算和 N-way 分片计算之间边的一个天真数据流图，将需要 $M+N$ 个节点和 $M\times N$条边，迅速变得不堪重负。

TensorFlow v1 的实现选择过于专门化，假设只有一个小型的、独占的加速器岛屿。这种过度专业化使得使用 TF 处理当代或未来的 ML 工作负载几乎不可行。虽然 TF 可以运行需要跨主机协调或通过 send 和 recv 操作进行数据传输的计算（图 1c），但目标主机上的工作（如调度加速器计算）只有在传输完成后才会触发。在涉及许多跨主机传输的程序中，例如具有大量阶段的流水线模型，这些调度延迟会累积，导致加速器利用率低效。虽然 TF v1 用户可以（低效地）通过使用控制边缘来强制执行程序内的一致调度顺序，但由于像 TF v1 这样的单控制器系统缺少集中式调度程序，无法确保跨程序计算之间的一致顺序。TF 还会物化完整的分片计算图，这在图序列化和执行时引入了巨大的开销，特别是当分片数量达到千级时，会导致子计算之间有数百万条图边。PATHWAYS 将单控制器框架的灵活性与多控制器的性能相结合。我们采用单控制器模型，因为我们认为它比多控制器更能为新颖且高效的 ML 计算提供更好的机会，既能利用计算稀疏性和异构性，也能支持促进共享和虚拟化资源的集群管理系统。我们的设计与旧有的单控制器 ML 系统不同，它使用异步调度来匹配多控制器系统的性能，支持集中式资源管理和调度，优先支持 SPMD 加速器计算的群体调度，并使用分片数据流系统来实现高效的协调。

3 PATHWAYS 编程模型

我们已实现对从 TensorFlow 和 JAX 编写的源程序的 PATHWAYS 支持，但本文的评估集中于 JAX。JAX 用户可以显式地使用装饰器包装标准 Python 代码，以指示应该编译成（可能是 SPMD）XLA 计算的代码片段。这些 XLA 计算通常具有已知的输入和输出类型与形状、有限的循环，并且几乎没有条件判断（更多细节见附录 B），因此可以提前估算计算所需的资源。我们将这些已知资源需求的计算称为“已编译函数”。每个这样的函数映射到 PATHWAYS 程序中的一个单一（分片）计算节点。
目前，JAX 不能扩展到超过一个 TPU pod，因为运行在多控制器配置中的 JAX 程序通过 XLA 集体操作传输所有数据，而这些操作当前仅在 TPU 上通过 ICI 可用。PATHWAYS 可以作为 JAX 后端的插件替代，使得 JAX 代码能够在不修改的情况下运行，只不过 SPMD 计算现在不仅可以访问本地连接的 TPU 核心，还可以访问系统中配置的所有 TPU 核心。由于 PATHWAYS 可以通过 ICI 和 DCN 进行通信，它使得 JAX 程序首次能够扩展到多个 TPU pod，涵盖成千上万的 TPU 核心。

运行未修改的 JAX 代码虽然方便，但并不能完全释放 PATHWAYS 的性能潜力。PATHWAYS 用户可以请求一组“虚拟设备”，并可以对设备类型、位置或互联拓扑设置可选约束，之后可以将特定的已编译函数部署到这些设备上（图 2）。系统会自动处理所有依赖计算之间的数据移动和重新分片。

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

JAX 支持对追踪代码的转换，其理念与我们希望探索的研究方向非常契合。例如，JAX 有一个名为 FLAX 的配套库（Heek 等，2020），用于表达分层的 DNN 模型，我们编写了一个库，自动将 FLAX 模型转换为流水线化的 PATHWAYS 程序。此外，JAX 支持将“每个示例” Python 函数向量化的转换，生成高效的批量代码，这些转换为探索新的基于数据的向量化控制流提供了良好的基础，正如我们稍后简要描述的那样 (§6.3)。

4. PATHWAYS 系统架构

PATHWAYS 在现有系统的基础上进行了大量构建，包括 XLA（TensorFlow，2019）用于表示和执行 TPU 计算，TensorFlow 图和执行器（Abadi 等，2016）用于表示和执行分布式 CPU 计算，以及包括 JAX（Bradbury 等，2018）和 TensorFlow API 在内的 Python 编程框架。通过利用这些构建块，我们能够集中精力于 PATHWAYS 的新型协调方面，同时能够在最小的代码修改下运行现有的 ML 模型。

4.1 资源管理器

PATHWAYS 后端由一组加速器组成，这些加速器被分组到紧密耦合的岛屿中，这些岛屿通过数据中心网络（DCN）相互连接（图 3）。PATHWAYS 具有一个“资源管理器”，负责对所有岛屿上的设备进行集中管理。客户端可以请求具有特定 2D 或 3D 网格形状的“虚拟切片”，以适应其通信模式。每个虚拟切片包含“虚拟设备”，允许客户端表达如何在网格上布局计算。资源管理器会动态分配物理设备给虚拟设备，以满足所需的互联拓扑、内存容量等要求。

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

PATHWAYS 允许后端计算资源动态地添加和移除，资源管理器会跟踪可用的设备。通过我们单控制器设计启用的虚拟设备与物理设备之间的间接层，未来我们将能够支持透明挂起/恢复和迁移等特性，在这种情况下，客户端的虚拟设备会被暂时回收或重新分配，而无需用户程序的合作。

4.2 客户端

当用户希望运行一个跟踪程序时，它会调用 PATHWAYS 客户端库，首先为所有尚未运行过的计算分配虚拟设备，并将这些计算注册到资源管理器中，触发服务器在后台编译计算。客户端随后构建一个设备位置无关的 PATHWAYS 中间表示（IR），该表示以自定义的 MLIR（Lattner 等，2021）方言表达。IR 会通过一系列标准的编译器转换逐步“降级”，最终输出一个低级表示，其中包含物理设备的位置。这个低级程序会考虑物理设备之间的网络连接，并包括将源计算片段的输出传输到目标片段位置的操作，当需要数据交换时，还包括 scatter 和 gather 操作。

在虚拟设备位置没有变化的常见情况下，反复运行低级程序是高效的，如果资源管理器改变了虚拟设备与物理设备之间的映射，程序可以重新降级。

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

4.3 协调实现

PATHWAYS 依赖于 PLAQUE 进行所有使用数据中心网络（DCN）的跨主机协调。PLAQUE 是一个现有的（闭源）生产级分片数据流系统，Google 用于许多客户面向服务的场景，尤其是在需要高扩展性和低延迟的高广度或高汇聚通信中。PATHWAYS 的低级 IR 直接转换为一个 PLAQUE 程序，表示为数据流图。PATHWAYS 对其协调基础设施有严格的要求，而这些要求都由 PLAQUE 满足。

首先，用于描述 PATHWAYS IR 的表示必须为每个分片计算包含一个单一的节点，以确保跨多个分片的计算具有紧凑的表示。也就是说，执行链式计算的两个计算 A 和 B，各自有 N 个计算片段时，数据流表示应该有 4 个节点：$Arg\to Compute(A)\to Compute(B)\to Result$，无论 N 的值如何。在 PLAQUE 运行时实现中，每个节点会生成带有目标分片标记的输出数据元组，因此，在执行数据并行计算时，N 个数据元组会流动，每个数据元组在相邻的 IR 节点之间传递。

协调运行时还必须支持沿着分片边的稀疏数据交换，即可以在动态选择的分片子集之间发送消息，并使用标准的进度跟踪机制（Akidau 等，2013；Murray 等，2013）来检测某个分片的所有消息是否已被接收。高效的稀疏通信是避免数据中心网络成为加速器数据依赖控制流瓶颈的必要条件，这是 PATHWAYS 要实现的关键能力之一。

协调基础设施用于发送处于关键路径中的 DCN 消息，这些消息用于传输调度消息和数据句柄（图 4）。因此，它必须以低延迟发送关键消息，并且在需要高吞吐量时，将目标主机的消息批量发送。

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

我们认为，使用 Ray（Moritz 等，2018）等其他分布式框架来重新实现完整的 PATHWAYS 设计是可行的，而不是依赖 PLAQUE 来实现低级协调框架。在这种实现中，PATHWAYS 执行器和调度器将被长时间运行的 Ray actors 替代，这些 actors 将在底层 Ray 集群调度之上实现 PATHWAYS 调度，而执行器可以使用 PyTorch 进行 GPU 计算和集体操作。为了实现类似的性能，可能需要进行一些补充（见 §5），因为 Ray 缺少例如 HBM 对象存储或高效传输远程对象的 GPU 互联原语等特性。

4.4 群组调度的动态调度

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

4.5 并行异步调度

在加速器上运行计算时，系统可以利用异步 API 来重叠计算和协调（Kwon et al., 2020）。考虑图 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，我们可以使用简单的背压机制来暂停该计算。

5 评估

在评估 JAX、PATHWAYS 和 TensorFlow 在 TPU 上的表现时，我们使用三种不同的配置。配置 (A) 每个主机有 4 个 TPU，最大的实例报告包含 512 个主机，共有 2048 个 TPU 通过 ICI 连接。配置 (B) 每个主机有 8 个 TPU，最大的实例报告包含 64 个主机，共有 512 个 TPU。配置 © 使用四个 TPU 岛屿，每个岛屿有 4 个主机和 32 个 TPU。我们会在文中注明实验使用了某一配置的部分 TPU。

在评估 Ray 在 GPU 上的表现时，我们使用 Ray v1.3 和 PyTorch 1.8.1，运行在 p3.2xlarge 虚拟机上，主机通过 DCN 连接并使用 Amazon Placement Groups 进行调度。

我们主要将 PATHWAYS 与多控制器 JAX 进行比较，因为 JAX 在工业标准基准测试中已展示了领先的性能（Mattson 等人，2020），且我们可以在相同的硬件配置上轻松运行 JAX 和 PATHWAYS（PW）。我们还将与 TensorFlow（TF）和 Ray 进行微基准测试比较，以检验 PATHWAYS 分布式系统性能的特定方面，并展示在 PATHWAYS 上运行的 TF 模型的流水线性能。

5.1 单控制器调度开销

我们的第一个实验是一个微基准测试，用于比较 JAX 多控制器和单控制器框架的开销。我们构建了一个程序，反复执行一个简单的 gang-scheduled 计算，其中包含一个标量的 AllReduce 操作，随后是一个标量加法，将一个计算的输出作为下一个计算的输入。我们测量吞吐量：每秒钟在加速器上执行的计算次数。我们比较了三种用户代码排队计算的方式：

• OpByOp (-O)：用户代码中为每次执行计算包含一个单独的调用。
• Chained (-C)：用户代码中包含一系列调用，每个调用执行一个包含 128 个节点的链式计算，每个节点执行一次计算。系统在响应单个客户端调用时执行整个链式操作。
• Fused (-F)：用户代码中包含一系列调用，每个调用执行一个单一计算节点，节点内包含 128 次计算。

对于 JAX 多控制器，OpByOp 意味着 JIT 编译一个包含单个计算的函数，并从 Python 中反复调用它；Fused 意味着 JIT 编译一个包含计算链的函数。对于多控制器来说，没有类似 Chained 的操作。对于 PATHWAYS，OpByOp 和 Fused 使用与多控制器相同的 JAX 源代码，而 Chained 则使用 PATHWAYS 程序跟踪器来构建一个多节点程序，每个节点包含一个简单计算。TensorFlow（TF）与 PATHWAYS 类似，我们构建相同的 TPU 计算并使用 TF 图执行，而不是 PATHWAYS。对于 Ray，OpByOp 意味着为每次计算执行一个单独的 actor 方法，该方法执行一个 PyTorch AllReduce 操作；Chained 意味着通过传递 Ray futures 链接一系列 actor 方法，每个方法执行一个单独的 PyTorch AllReduce；Fused 意味着执行一个单独的 actor 方法，该方法在一个循环中运行一系列 PyTorch AllReduce 命令。

图表显示了结果。需要注意的是，OpByOp 是一个最坏情况实验，在任何框架中都不是惯用的做法，它仅用于强调底层系统的性能瓶颈。正如预期的那样，对于 OpByOp，JAX 多控制器的吞吐量比单控制器系统要好得多，特别是当加速器数量增加时。PATHWAYS 的大部分开销来自于客户端在协调器排队完成一个计算并返回输出句柄后，才开始排队下一个计算。我们可以通过允许用户代码与排队的 RPC 并行进行，从而消除大部分开销，同时通过批量处理多个小计算来生成一个单一的 PATHWAYS 程序来进一步优化这一点。我们没有特别关注优化非常小的计算的开销，因为在实际模型中，涉及多个标量的计算，PATHWAYS 已经能够匹配 JAX 多控制器的性能（见 §5.3）。一旦足够的工作被融合到一个单一节点，PATHWAYS 的性能就能与 JAX 匹敌，甚至在 1000 个 TPU 核心的情况下也是如此。此外，PATHWAYS 的 Chained 方法在 256 个核心时优于 JAX 的 OpByOp，因为 PATHWAYS 可以直接从 C++ 执行连续的加速器计算，而 JAX 的 OpByOp 每次计算都需要切换到 Python。

TensorFlow 和 Ray 都因缺乏设备对象存储而遭受性能瓶颈：Ray 必须将计算结果从 GPU 转移到 DRAM，然后将对象句柄返回给客户端，而 TensorFlow 则需要将数据传输回客户端。这种开销影响了它们的 OpByOp 性能，但对于 Chained 和 Fused，这一开销大多会被摊销。由于 Ray 和 PATHWAYS 使用不同的硬件，二者的性能不具可比性，但我们认为，如果完整的 PATHWAYS 设计被实施，并替代 Ray 来实现 PLAQUE，应该可以达到类似的性能。Ray 的 OpByOp 性能比 PATHWAYS 慢一个数量级是可以理解的，因为 Ray 可以执行通用的 Python actor，而 PATHWAYS 是专门为从 C++ 启动的 TPU 计算进行优化的。通过精心设计，可能有可能为 Ray 添加快速路径，比如一个基于 GPU 的对象存储和有效的 GPU 互联传输原语，来消除大部分额外的开销。TensorFlow 在多核心运行时表现较慢，因为它使用了集中式的 barrier，通过控制边缘来序列化 gang-scheduled 计算。

图 6 变化了每次计算的时间，以找出 PATHWAYS 匹配 JAX 吞吐量的最小计算时间。对于配置 (B) 中的 16 个主机和 128 个 TPU，仅需 2:3 毫秒就能实现平衡，而对于配置 (A) 中的 512 个主机和 2048 个 TPU，至少需要 35 毫秒的计算时间才能掩盖 PATHWAYS 的单控制器开销。

我们的下一个微基准测试也在配置 (B) 上，评估 §4.5 中描述的并行异步调度机制的优势。我们构建了一个更现实的流水线基准，其中之前实验中的简单计算再次串联在一起，但现在每个计算都在不同的 4 个 TPU 核心上运行，并且每个计算都运行在不同的主机上。一个计算的输出数据必须通过 ICI 发送，才能让下一个计算执行。图 7 显示了三个“阶段”：首先，随着主机数量的增加，固定客户端开销逐渐摊销；然后，增加更多阶段时传输成本开始占主导地位；最后，系统开始摊销固定的调度开销。最终，我们预计传输开销将再次占主导地位。为了比较，我们还展示了当我们强制 PATHWAYS 数据流执行使用顺序异步调度时的性能，在这种情况下需要等待一个计算被排队后才能排队下一个计算，从而衡量并行异步调度带来的收益。

5.2 多租户

我们在图 8 中验证了 PATHWAYS 能够在并发程序之间对加速器进行时间复用（在配置 (B) 上执行）。当多个客户端同时提交不同的 PATHWAYS 程序时，PATHWAYS 可以实现至少与 JAX 相同的总吞吐量，即在资源能够同时适应 HBM 时，不会出现不同客户端之间程序上下文切换的开销（附录 D 中的跟踪）。正如图 6 中所展示的那样，匹配吞吐量所需的并发度在计算规模较大时较低，因为 TPU 核心很快就能达到完全利用。值得注意的是，对于非常小的计算，PATHWAYS 的最大吞吐量超过了 JAX，能够实现更高的 TPU 利用率。这是因为 PATHWAYS 工作线程可以接受来自远程客户端的更多计算，而 JAX 使用本地 Python 调度时，无法达到相同的调度能力。

图 9 显示了在 PATHWAYS 上运行上述工作负载时，128 个核心的跟踪结果。这个实验突出显示了 PATHWAYS 对 4 个独立客户端提交的程序进行帮调度，同时控制加速器时间的分配以保证公平性；例如，调度器可以在多租户设置中强制执行按比例共享。

5.3 大规模模型性能

最后，我们展示了 PATHWAYS 在训练可以表达为 SPMD 程序的实际机器学习模型中的性能。我们将 JAX 和 TF 模型在其本地系统上运行与在 PATHWAYS 上运行的相同模型进行了比较，并验证了数值结果完全一致，因此我们只关注性能。

首先，我们比较了 JAX 多控制器在训练用于多种文本到文本自然语言处理任务的 Transformer 模型（具有编码器-解码器架构）时的表现。我们使用了 (Raffel et al., 2019) 的模型配置，并在每个加速器具有 16GB 内存的 TPUv3 上运行实验。表 1 显示了不同模型大小（最多 110 亿个参数）在不同数量的加速器上训练的文本到文本 Transformer 模型的训练吞吐量（每秒处理的 token 数）。如预期的那样，由于模型代码相同，JAX 和 PATHWAYS 上训练的模型在相同的步骤数下实现了相同的困惑度。在所有测试的模型大小中，由于现实中的计算足够大，能够掩盖单控制器的开销，因此这两个系统表现出相同的性能。尽管我们没有报告详细结果，但我们有大量经验在 PATHWAYS 上运行 JAX 模型，这进一步证实了两个系统在广泛设置下的性能是相当的。

接下来，我们比较了 PATHWAYS 在训练基于 Transformer 的语言模型（采用仅解码器架构）时，在配置 (B) 和 © 上的表现。对于这个实验，我们使用了一个通过 TensorFlow 用 Python 表达的模型。该模型包含 62 层 Transformer，每层的模型维度为 2048，隐藏层维度为 8192，总共有 30 亿个参数。我们将 SPMD 配置与使用 GPipe 类似的调度方式（Huang et al., 2019）下的流水线模型进行了比较。流水线模型被分成多个阶段，每个阶段的计算量是平衡的。由于第一个阶段有额外的嵌入查找层，而最后一个阶段有额外的 softmax 层，因此我们从第一个和最后一个阶段中各取出一个 Transformer 层，以平衡每个阶段的计算量。每个阶段被分配到一组不同的加速器，跨越多个主机。

表 2 显示了在不同阶段数 (S) 和微批次数 (M) 下的训练吞吐量，同时保持全局批次大小和训练超参数不变。每个微批次的示例数对于所有情况固定为 4，因此对于 128 核配置，全局批次大小每步为 2048（对于 512 核配置为 8192）。

PATHWAYS 的训练吞吐量与流水线阶段中 TPU 核的数量成比例增加（见表 2），这一结果与其他系统（Rasley 等，2020；Narayanan 等，2021）一致。这一结果也与图 5 中显示的 PATHWAYS 吞吐量随着主机数量线性扩展的情况相符。增加流水线阶段的数量带来的开销非常小，当阶段数从 4 增加到 16 时，吞吐量从 133.7k tokens/sec 减少到 131.4k tokens/sec。我们将流水线模型的表现与使用 SPMD 表达的等效模型进行比较，并观察到，至少在此实例中，流水线模型与 SPMD 的性能相当，因为 SPMD 计算中的集体通信引入的开销大于流水线中的“气泡”开销。

我们还展示了 PATHWAYS 如何高效地训练通过 DCN 连接的 TPU 岛屿上的模型。在 S = 16，M = 64 配置下，使用 128 核的单一岛屿与使用 4 个每个 32 核的岛屿（配置 ©）时，吞吐量相同，均为 131.4k tokens/sec。图 10 显示了当流水线阶段被划分到岛屿时的一部分核的执行追踪。DCN 传输发生在每 8 行的追踪组之间，并未在追踪中显现，因为通信时间已有效与计算重叠。

最后，我们将大规模 Decoder-only Transformer 模型的训练扩展到 64B 和 136B 参数，使用两个岛屿的加速器进行训练。在 PATHWAYS 上使用通过 DCN 连接的两个岛屿进行训练时，PATHWAYS 达到了单个岛屿设备数量的 97% 的吞吐量。对于 136B（64B）语言模型，我们在两个岛屿上分别使用 1024（512）核，利用岛屿内部快速 ICI 进行归约，然后通过 DCN 进行岛屿间传输（执行追踪可见于附录 D），全局归约所需的数据量分别为 1030GB（457GB）。

6 讨论

6.1 PATHWAYS 设计与实现

PATHWAYS 的设计目标是面向大规模 TPU 加速器的使用。由于 TPU 和 GPU 的硬件架构差异，许多低级设计决策受到了影响。TPU 和 GPU 最大的区别在于，长时间运行和更复杂的计算可以融合到一个单一的 TPU 核心中执行，这是因为 TPU 支持丰富的控制流和通信原语，而 GPU 系统需要由驱动代码来执行这些原语。与此相反，GPU 与主机内存系统和 DCN（NVIDIA, 2021）有更紧密的集成（更多细节见附录 A.5）。TPU 适合 PATHWAYS，因为 XLA 可以编译包含融合集体操作的高效函数，而且大规模的高性能 TPU 互联允许灵活调度不同大小的计算。

尽管如此，我们认为我们在 PATHWAYS 中做出的许多高级架构选择，及本论文中描述的设计思路，对于大规模的 GPU 系统同样适用。

6.2 资源管理

PATHWAYS 旨在支持多种精细化的动态资源管理策略。我们最初的研究集中在 TPU 计算的高效动态时间复用上。对于未来更复杂的多租户用例，PATHWAYS 将需要处理更广泛的资源类型，包括但不限于设备和主机内存、ICI、DCN 和 PCIe 带宽。PATHWAYS 的单控制器模型使系统能够广泛地跟踪可用资源，并在大规模下进行资源分配。我们计划探索常见的多租户需求，如优先级、性能隔离、访问控制和资源计量，但这将在比以往更小的时间尺度下进行，并且资源池的规模将大幅扩大（例如，数千个核心和 TB 级加速器内存）。

6.3 数据依赖的向量化控制流

目前几乎所有的机器学习模型都会根据每个训练步骤中的每个训练样本更新所有的模型权重。我们希望能够支持使用细粒度的控制流进行研究，以便能够根据每个示例甚至子示例（图像的补丁或句子的单词）来更新不同的模型权重。像专家混合模型（Mixture of Experts, MoE）（Shazeer et al., 2017）和路由胶囊网络（routed capsule networks）（Hinton et al., 2018; Barham and Isard, 2019）这样的模型通过根据随着训练进展不断更新的学习函数“路由”不同的（子）示例到托管不同子集模型权重的加速器，从而利用计算稀疏性。这种路由需要节点之间进行细粒度的数据依赖交换。我们的机器学习研究同事告诉我们，他们希望在训练越来越大的模型、承担越来越多任务时，能够更有效地利用稀疏性，但当前的框架限制了他们在实验新型模型架构方面的能力。支持数据依赖的向量化控制流将是未来研究的主题，它不仅需要一个干净的编程模型，还需要良好的性能。

7 相关工作

我们在 §2 中详细讨论了与本工作密切相关的研究。本节进一步扩展了那些处理需要超出 SPMD 多控制器提供的能力的机器学习工作负载的相关研究，并验证了我们 PATHWAYS 设计选择的有效性。

在多个任务之间共享加速器对于实现高资源利用率至关重要。传统的资源共享通常是粗粒度的。例如，通用虚拟化使得云应用能够有效地共享多租户资源并保证性能隔离（Angel et al., 2014; Wentzlaff et al., 2010; Shahrad and Wentzlaff, 2016; Baumann et al., 2009），但云提供商会将加速器分配给单个用户。集群调度器优化 ML 工作负载的异构性（Narayanan et al., 2020）和多作业、多用户的公平性与性能（Xiao et al., 2018; Ren et al., 2015; Mahajan et al., 2020; Jeon et al., 2018），但资源仍然是长期为单个作业独占。

近期的研究表明，细粒度的共享可以进一步提高资源效率：加速器虚拟化（Yu et al., 2020; Gupta et al., 2011; Vijaykumar et al., 2016）避免将整个加速器专门分配给单个用户。大规模模型（Brown et al., 2020）可能受到可用加速器内存的限制，因此需要 GPU 内存虚拟化（Rhu et al., 2016; Ausavarungnirun et al., 2018）或 DRAM 卸载（Rajbhandari et al., 2021）。并行（时间复用或重叠）执行机器学习任务（Gupta et al., 2018; Xiao et al., 2020; Bai et al., 2020; Yu and Chowdhury, 2020; Wang et al., 2021; Lim et al., 2021）有助于挖掘加速器中的空闲资源。这些细粒度共享技术展示了在没有像 PATHWAYS 这样单控制器系统的情况下，难以在大规模下有效利用的加速器共享机会。许多研究表明，偏离 SPMD 计算模型可以提高大规模工作负载的效率：流水线（Huang et al., 2019; Narayanan et al., 2019; Yang et al., 2021）将 ML 模型划分为跨加速器的静态异构计算。图神经网络训练（Jia et al., 2020）、神经架构搜索（Pham et al., 2018）和多模态多任务学习系统（Ma et al., 2018; Lepikhin et al., 2020; Zhao et al., 2019）是一些典型的异构和动态任务实例，它们不适合直接使用 SPMD 模型。我们预见到，未来的大规模高效 ML 模型可能由共享层和专属层组成（Bommasani et al., 2021），这可以自然地表示为 MPMD（Multiple Program Multiple Data）。

8 结论

PATHWAYS 在当前的单租户 SPMD 模型中与最先进的多控制器系统性能相当。我们已确保与多控制器 JAX 的严格兼容性，正如我们在 §5 中展示的，PATHWAYS 在大规模系统上与 JAX 的性能一致，除了在计算量最小的情况下。

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

我们展示了精心的系统设计和工程如何让我们“获得双重优势”，在今天的 ML 模型上匹配性能，同时为未来的模型编写提供所需的功能。

附录

A 加速器设计考虑

硬件加速对现代深度学习至关重要；然而，利用加速器实现高性能是一个非简单的系统工程。以下子章节列出了深度学习系统中常用的已建立的技术，用以实现良好的性能。

A.1 批处理

随着 Dennard 定律的结束，加速器实现了硬件并行性，通常采用 SIMT（Kirk, 2007）或脉动阵列（Jouppi et al., 2020）设计。虽然这些硬件架构消除了算术瓶颈，但内存带宽很快成为关键资源，迫使使用高带宽内存（HBM），这是一种昂贵且容量有限的内存技术。现代神经网络的训练方案利用批处理来解锁并行性（有助于并行 ALU 的输入）并实现内存重用（一个浮动值从内存中读取一次，并用于多次计算，从而大大减少计算的内存带宽需求）。然而，批处理并不是万灵药：它对有限的 HBM 内存容量施加压力，且非常大的批量大小可能会减缓模型的收敛速度（Shallue et al., 2018; You et al., 2017; Lanchantin et al., 2020; Anil et al., 2021）。虽然现代 GPU 支持统一内存——能够透明地在加速器之间或从 HBM 到主机的 DRAM 之间分页——如果用户不小心，受 HBM 带宽限制的计算可能会降速至 PCIe 带宽，从而使加速器的利用率下降一个数量级（Lim et al., 2021）。

A.2 异步编程

加速器抽象依赖于异步编程模型来实现性能；同步抽象会浪费太多加速器计算资源，例如 PCIe 延迟、内核调度开销和中断延迟。计算任务被排队到流中，以便在未来某个时刻在加速器上执行。这种异步抽象有效地隐藏了小操作的调度延迟，只要维持足够大的工作流水线。

A.3 高性能互连

现代深度神经网络的规模远远超过加速器（HBM）内存的容量（Lepikhin et al., 2020; Huang et al., 2019）。这些神经网络中的并行性适合于在多个加速器上同时进行分片计算，然而，加速器之间的高速互连在此时变得至关重要。GPU 使用如 NVLink 等互连技术，实现小范围主机之间加速器“岛屿”的高速通信（Naumov et al., 2020），并通过以太网和 Infiniband NIC 的 RDMA 功能（GPUDirect）在岛屿之间快速传输数据。TPU 则在芯片中内置了定制的网格网络，芯片可以直接通信，而无需涉及主机或数据中心网络。专用的 GPU 和 TPU 互连通常通过 30 年前的 MPI 原语（例如，AllReduce）暴露给应用，这些原语必须进行集群调度，以确保每个程序同时进入相同的原语。随着更大计算任务的运行（例如，训练一个更大的神经网络，或者通过一种称为数据并行扩展的弱扩展方式，训练固定大小的神经网络以利用更多的加速器），需要更快速的集体操作，因此网络带宽也需要得到保证，以维持集群资源的高效利用。这促使了大量对替代芯片网络拓扑的实验，包括超立方体、二维和三维网格环网（Naumov et al., 2020）。

A.4 单租户

与计算机中的大多数资源不同，加速器通常不会被多个程序同时共享。深度学习模型可以通过增加参数数量或批量大小来轻松扩展，以使用更多的内存，因此实际中程序消耗了大部分可用的加速器（HBM）内存。PCIe 带宽远小于 HBM 或加速器互连带宽。这意味着细粒度的上下文切换（其中大量 HBM 中的数据被通过 PCIe 页面调出到主机 DRAM）会浪费大量的加速器周期。因此，当主机程序没有充分利用加速器时，计算资源就会被闲置，无法有效使用。此外，实际中加速器资源的抢占最小化，导致在为异构工作负载提供服务的大型共享集群中资源调度不尽优化；很难分配大量物理上接近的设备来利用网络局部性。

A.5 比较 GPU 和 TPU

虽然 GPU 和 TPU 之间有许多相似之处，但也存在一些重要的差异。GPU 系统往往拥有小型的通过 NVLink 连接的设备岛屿（例如，一台主机内的 8 个 GPU），更大的设备聚合通过 infiniband 或数据中心网络技术连接。GPU 通常通过调度许多小的预编译“内核”来编程加速器，因为它们是预编译的，内核必须支持动态形状。无论是通过 NVLink 还是通过 DCN，GPU 之间的任何通信都是通过 NCCL 库进行的，并由主机发起。

TPU 系统拥有数千个设备相互连接，每个“岛屿”中有数百个主机（见图 3 中部）。TPU 包含一个强大的“标量核心”，该核心协调 TPU 的向量计算单元，允许 TPU 执行由 XLA（TensorFlow, 2019）编写的长时间运行的函数，而无需主机交互，这些函数可能包括通过专用的 ICI 网络进行的集体通信。因此，在 TPU 上，ML 框架通常构建一个大型的 XLA 程序，该程序经过即时（JIT）编译并分派到加速器上。由于单个 XLA 计算的运行时间通常比 GPU 内核长几个数量级，因此编译器会加大优化力度，例如静态缓冲区分配和中间程序值的自动再物化（节省内存容量）。由于静态缓冲区分配的结果，TPU 对动态形状的支持有限，这使其非常适合 PATHWAYS 概念中的常规编译函数。

TPU 限制一次只能运行一个程序，不支持本地抢占，主要是因为它们高性能的 RDMA 通信实现使得在没有分布式协调的情况下进行安全的抢占变得困难。由于计算无法抢占，因此必须按一致的顺序将通信计算排队到设备上，否则系统将发生死锁。这个要求意味着 PATHWAYS 必须执行集中式的 gang-scheduling（批量调度）。然而，正如本文主文所指出的，gang-scheduling 对 GPU 效率也非常有利。对于一个优先考虑 ML 训练工作负载的集群，其中吞吐量比延迟更为重要， dedicating 整个 GPU 或 GPU 的静态部分给一个精心调整的计算任务，比允许 GPU 驱动程序和硬件运行时在多个竞争的并发计算任务中动态多路复用计算资源更为高效。因此，即使 GPU 可以在没有集中调度的情况下执行并发程序，采用像 PATHWAYS 这样的设计仍然能够更高效地利用资源。

B. 典型 ML 程序的结构

本小节描述了典型的当代 ML 计算，从高层结构角度映射子计算到加速器，并将子计算降级为加速器内核。

执行 ML 工作负载的加速器上的计算主要由我们称之为“编译函数”的子计算主导。这些编译函数具有以下特征：

• 输入和输出类型，以及任何输入/输出张量的形状，在输入数据计算之前是已知的。
• 任何循环的边界要么在节点计算调度时已知，要么指定为最大迭代次数，并有可能提前终止。
• 条件语句是“函数式的”，即两个分支具有相同的输出类型，并且提前分配足够的资源以满足任一分支的需求。

编译函数的约束主要来源于 ML 模型与硬件的共同发展，详见 §A。在此，我们讨论编译函数的资源需求已知这一事实的某些影响。

当今几乎所有高性能 ML 计算都是通过长时间执行的编译函数来表示，且只有在极少数情况下（如果有的话）才会基于由编译函数计算得出的数据进行分支。由于系统可以提前为编译函数进行资源分配，当代的 ML 框架通过异步排队编译函数，允许在其前驱计算未完成之前执行，从而使主机端的工作可以与加速器计算并行进行（Bradbury et al., 2018；Paszke et al., 2019）。框架尽可能将编译函数图提交给“即时” (JIT) 编译器（Chen et al., 2018；TensorFlow, 2019），该编译器能够利用布局分配和融合等优化，从而大幅提升最终加速器代码的效率。

为了优化编译函数图以实现加速器的峰值性能，框架通常会追踪可以降级为编译函数的高层（Python）代码片段的执行。因此，即使客户端代码可能是用高级语言编写的，并且与主机上的运行上下文绑定的复杂状态，性能敏感的节点计算通常会降级为内部表示（IR），该表示可以序列化并相对容易地发送到远程主机执行。

C. 输入数据处理

JAX 故意避免重新实现数据加载管道，tensorflow/datasets（TensorFlow, 2021）通常用于 JAX 的输入处理，因此将 JAX 程序适配以将输入处理卸载到在 PATHWAYS 工作节点上运行的基于 CPU 的 TensorFlow 执行器并不困难。PATHWAYS 在每个主机上实例化一个基于 CPU 的 TensorFlow 执行器，使得用户程序可以将输入处理序列化为 TensorFlow 图，并将其分发到各个工作节点。我们计划支持流数据协议，以便可以在独立管理的服务器集群上执行基于 CPU 的计算，从而将昂贵的 TPU 连接主机与可用于输入处理的 CPU 资源解耦。

D. 评估工作负载跟踪

图11展示了图8中工作负载的跟踪，客户端并发提交程序的数量不同（§5.2）。单个客户端使用每个程序非常小的计算时间0.33毫秒，这不足以饱和加速器。在PATHWAYS的多租户支持下，使用多个客户端可以将设备利用率提高到大约100%。所有客户端程序在所有核心上进行并行调度，并且以毫秒级别或更小的时间间隔交错执行，几乎没有上下文切换的开销。

图12显示了多个训练步骤的跟踪剖面，当64B解码器的Transformer模型在两个加速器岛上（每个岛上有512个芯片）通过数据并行训练时（§5.3）。前八行（蓝色）对应于第一个岛上主机上的TPU计算，接下来的八行（绿色）对应于第二个岛上主机上的TPU计算。在这种情况下，每个岛计算梯度，然后将这些梯度的传输任务排入队列，传输完成后，DCN会结束，每个岛应用接收到的梯度并开始下一轮训练步骤。即使在128个主机对的规模下，DCN传输也几乎没有开销，导致与使用ICI通信的SPMD配置相比，训练吞吐量达到97.2%。

论文名称：PATHWAYS: ASYNCHRONOUS DISTRIBUTED DATAFLOW FOR ML

论文地址：https://proceedings.mlsys.org/paper_files/paper/2022/file/37385144cac01dff38247ab11c119e3c-Paper.pdf