谈及当前AI系统面临的问题时，低效是经常提到的。谷歌人工智能总监JeffDean曾在博文中写道，“今天的人工智能系统总是从头开始学习新问题——数学模型的参数从随机数开始。这就像每次学习一项新技能（比如跳绳），你总是忘记之前学过的一切，包括如何平衡，如何跳跃，如何协调手部动作等等，然后从头开始重新学习。这或多或少是我们训练大多数机器学习模型的方式今天如何：我们不是扩展现有模型来学习新任务，而是从头开始训练一个新模型来做一件事（或者我们有时为特定任务专门化一个通用模型）。结果是我们最终得到了数千个为任务开发了数千个模型。以这种方式学习每个新任务不仅需要更长的时间，而且需要更多的数据。”为了改变这种情况，JeffDean等人在去年提出了一种叫做“Pathways”的通用AI架构，他说Pathways旨在用一个架构同时处理多个任务，并且具有快速学习的能力这个架构的特点可以概括为：它可以训练一个模型做几千件事；目前的模型只关注一种感觉，而Pathways可以做多种事情；目前的模型是密集的且效率低下，而Pathways会让模型变得稀疏而高效。JeffDean在发表想法半年多后，终于发表了Pathways论文，其中包含了很多技术细节。论文链接https://arxiv.org/pdf/2203.12533.pdf论文中写道PATHWAYS使用异步算子的shardeddataflowgraph，消费和生成futures，并在数千个加速器上高效地gang-schedule异构并行计算ators，同时协调其专用互连上的数据传输。PATHWAYS使用新的异步分布式数据流设计，允许控制平面并行执行，尽管数据平面存在依赖关系。这种设计让PATHWAYS采用单控制器模型，更容易表达复杂的新并行模式。实验结果表明，在2048个TPU上运行SPMD（单程序多数据）计算时，PATHWAYS（加速器利用率接近100%）的性能与SOTA系统相当，吞吐量与跨越16个阶段或拆分相当进入由数据中心网络连接的两个加速器岛的Transformer模型的SPMD示例。以下为论文详情：研究背景近十年来，深度学习在图像理解、自然语言处理等诸多领域取得了显着进步。连接两者的是ML模型、加速器硬件和软件系统。共同进化的结果。这种协同进化带来的隐患是，深度学习系统可能过度关注当前的工作负载而无法预测未来的需求。PATHWAYS是为分布式ML构建的新系统，针对未来ML工作负载所需的特定功能。目前，这些工作负载缺乏SOTA系统的支持。例如，当今的SOTAML工作负载大多使用受MPI启发的单程序多数据(SPMD)模型，其中所有加速器同步运行相同的计算，加速器之间的通信由AllReduce等集合体描述。但近年来，研究人员开始在ML计算中受到SPMD的阻碍。大规模语言模型已经使用流水线并行而不是纯数据并行进行了扩展；混合专家（MoE）等模型已经开始探索计算稀疏性，其最自然的表现是使用细粒度控制流和跨加速器的异构计算；系统设计人员已经开始采用巧妙的技术在MPI风格的系统上实现流水线式、同构的MoE模型，但是MPI编程模型对用户和底层系统的限制都太大了。另一方面，ML中的异构环境在每一代新一代加速器中变得越来越普遍。提供对同类加速器（通过高带宽互连连接）的大型“孤岛”的独占访问是昂贵的并且通常是浪费的，因为单个用户程序必须努力保持所有加速器始终忙碌。这种限制进一步推动研究人员转向“多程序、多数据”(MPMD)计算。MPMD通过将整个计算的子部分映射到一组更小、更易于访问的加速器岛来实现更大的灵活性。为了提高利用率，一些ML硬件资源管理研究人员以细粒度的方式在工作负载之间重用硬件，实现工作负载弹性，并提高容错能力。最后，研究人员着手标准化一组基础模型，这些模型在大数据上进行大规模训练，可以适应各种下游任务。此类模型的训练和推理提供了通过跨多个任务多路复用资源并在它们之间有效共享状态来提高集群利用率的机会。例如，多个研究人员可能会同时针对不同的任务微调基础模型，使用相同的加速器来保持基础模型层的固定。共享子模型的训练或推理可以受益于允许将来自不同任务的示例组合在矢量化批处理中以提高加速器利用率的技术。本文提出的PATHWAYS在功能和性能方面可与SOTAML系统相媲美，同时提供支持未来ML工作负载所需的功能。它使用客户端-服务器架构，使PATHWAYS的运行时能够代表系统管理计算岛上的许多客户端执行程序。PATHWAYS是第一个旨在跨多个TPUpod透明高效地执行程序的系统。它可以通过采用新的数据流执行模型扩展到数千个加速器。PATHWAYS的编程模型使得表达非SPMD计算变得容易，并支持集中资源管理和虚拟化以提高加速器利用率。PATHWAYS系统架构PATHWAYS建立在以前的系统之上，包括用于表征和执行TPU计算的XLA（TensorFlow，2019），用于表征和执行分布式CPU计算的TensorFlow图和执行器（Abadi等人，2016），以及Python编程框架（Bradbury等人）.,2018)和TensorFlowAPIs包括JAX(Bradburyetal.,2016)。使用这些构建块，PATHWAYS能够以最少的代码更改运行现有的ML模型，同时保持协调。资源管理器PATHWAYS的后端由一组加速器组成，这些加速器被分组为紧密耦合的岛，这些岛又通过DCN相互连接，如上图3所示。PATHWAYS有一个“资源管理器”，负责集中管理岛上的所有设备。客户端可能会请求岛的“虚拟切片”具有适合其通信模式的特定2D或3D网格形状。每个虚拟切片都包含“虚拟设备”，允许客户端表达计算在网格上的布局方式。资源管理器动态地将物理设备分配给满足所需互连拓扑、内存容量等的虚拟设备。最初的资源管理器使用简单的启发式方法实现了这一点，试图通过在所有可用设备上分散计算来静态平衡负载，并维护一个虚拟设备和物理设备之间的一对一映射。如果未来的工作负载需要它，可以采用更复杂的分配算法，例如考虑所有客户端计算的资源需求和系统的当前状态来近似物理设备的最佳分配。PATHWAYS允许动态添加和删除后端计算资源，资源管理器跟踪可用设备。单控制器设计启用的虚拟和物理设备之间的间接层将允许未来支持透明挂起/恢复和迁移等功能，其中客户端的虚拟设备可以在没有用户程序协助的情况下临时回收资源或重新分配。client当用户想要运行一个被跟踪的程序时，可以调用PATHWAYS客户端库。它首先将虚拟设备分配给之前未运行过的任何计算，向资源管理器注册计算，并触发服务器在后台编译计算。然后，客户端为程序构建一个与设备位置无关的PATHWAYS中间表示(IR)，表示为自定义MLIR(Lattneretal.,2021)方言。IR通过一系列标准编译器传递逐渐降低，最终输出包含物理设备位置的低级表示。这个低级程序考虑了物理设备之间的网络连接，并包含将输出从源计算分片传输到其目标分片位置的操作，包括需要数据交换时的分散和收集操作。在虚拟设备的位置没有改变的通常情况下，重新运行低级程序是有效的。如果资源管理器改变了虚拟设备和物理设备的映射关系，它可以重新降低程序。较旧的单控制器系统中的客户端很快就会成为性能瓶颈，因为它负责协调数千个单独的计算，以及协调与分布在数千个加速器上的计算分片相对应的数据缓冲区。PATHWAYS客户端使用切片缓冲区抽象来表示可能分布在多个设备上的逻辑缓冲区。这种抽象通过以逻辑缓冲区而不是单个分片的粒度分摊簿记任务（包括引用计数）的成本来帮助客户端扩展。协调实施PATHWAYS依赖于PLAQUE来使用DCN进行所有跨主机协调。PLAQUE是一个现有的（闭源）生产分片数据流系统，谷歌将其用于许多需要高扇出或高扇入通信且可扩展性和延迟很重要的面向客户的服务。Low-levelPATHWAYSIR被直接翻译成PLAQUE程序并表示为数据流图。PATHWAYS对其配位底物有严格要求，而PLAQUE则满足所有要求。首先，用于描述PATHWAYSIR的表示必须包含每个分片计算的单个节点，以确保跨多个分片的计算的紧凑表示，即具有N个计算分片的2个计算A和B的链无论N是什么，每个计算片都有数据流表示中的4个节点：Arg→Compute(A)→Compute(B)→Result。在PLAQUE运行时实现中，每个节点都会产生一个输出数据元组，并标记目标分片，因此在执行数据并行执行时，N个数据元组将在每对相邻的IR节点之间流动。协调运行时还必须支持沿分片边缘的稀疏数据交换，其中可以在动态选择的分片子集之间发送消息，使用标准进度跟踪机制（Akidau等人，2013年；Murray等人，2013年）检测所有消息何时发送因为已收到碎片。高效的稀疏梳理避免了DCN成为加速器上依赖数据的控制流瓶颈，这是PATHWAYS实现的关键特性之一。如下图4所示，coordinationsubstrate用于在传输调度消息和数据句柄的关键路径上发送DCN消息，因此它必须低延迟发送关键消息，并在高吞吐量时批量发送消息到同一主机必需的。使用可扩展的通用数据流引擎来处理DCN通信也很方便，因为这意味着PATHWAYS也可以将其用于后台管理任务，例如分发配置信息、监控程序、清理程序、在发生故障时引发错误等.谷歌认为，使用Ray（Moritz等人，2018年）等其他分布式框架代替PLAQUE进行低级协调框架来重新实现完整的PATHWAYS设计是可行的。在此实现中，PATHWAYS执行器和调度器将被长期运行的RayActors取代，这些RayActors在底层Ray集群调度之上实现PATHWAYS调度，并且执行器可以使用PyTorch进行GPU计算和聚合。组调度动态调度如前所述，在一组共享加速器上支持SPMD计算的一个要求是支持高效的组调度。PATHWAYS运行时包括一个每个岛的集中式调度程序，该调度程序始终对岛上的所有计算进行排序。当PATHWAYS排队执行一个程序时，PLAQUE数据流程序负责：排队网络发送到远程加速器队列以获得函数执行输出的缓冲期货；与调度程序通信，以确定岛上运行的所有程序的一致功能执行顺序。调度程序必须实施以毫秒为单位分配加速器的策略。然而，当前的实现只是按FIFO顺序排列工作，但更复杂的调度程序可能会根据估计的执行时间重新排序计算。并行异步调度在加速器上运行计算时，系统可以利用异步API来重叠计算和协调。如下图4a的三节点图所示，方块分别对应三个节点A、B、C，分别运行在连接主机A、B、C的加速器上，所有节点计算都是常规的编译函数。主机A使节点A入队，接收A的未来输出并将其传输到主机B。主机B对节点B的输入进行分类，将输入缓冲区地址传输到节点A，并执行大部分准备工作以启动节点B的功能。当节点A完成时，其输出直接通过加速器互连发送到节点B的输入缓冲区，主机B启动节点B。一个节点完成和另一个节点启动之间的延迟比数据传输时间长。当前驱节点的计算时间超过用于调度、资源排序和主机间协调的时间时，上述设计效果很好。但如果计算时间太短，异步流水线就会停滞，主机端的工作成为整个计算序列执行的关键瓶颈。考虑到编译后的函数都是规则的，后续节点的输入形状实际上可以在前一个计算入队之前计算出来。因此，谷歌推出了一种新的并行异步调度设计方案，如下图4b所示。该方案利用常规编译函数的静态已知资源在计算节点上并行运行主机端工作，而不是在前任排队后序列化节点工作。考虑到正则函数只能并行调度工作，PATHWAYS使用并行调度作为优化方法，在前驱计算完成之前节点资源需求未知时回退到传统模型。当一个计算子图是静态可调度的时，程序将描述整个子图的单个消息发送到一个调度器，该调度器能够对子图中所有活动分片的执行进行背靠背排序。单个消息旨在最小化网络流量，但不需要调度程序将所有子图分片作为批次排队：计算可能仍会与其他并发执行程序提交的计算交错。三节点程序的顺序调度(a)与并行调度(b)。实验结果Google展示了PATHWAYS在训练真实机器学习模型（可以表示为SPMD程序）方面的性能。首先与使用编码器-解码器架构运行Transformer模型的JAX多控制器进行比较。下面的表1显示了在不同数量的加速器上训练时，不同大小的Text-to-TextTransformer模型的训练吞吐量（令牌/秒）。正如预期的那样，由于模型代码相同，在JAX和PATHWAYS上训练的模型以相同的步骤数实现了相同的困惑度。接下来，Google比较了PATHWAYS在仅使用解码器架构训练Transformer语言模型时在不同配置上的表现。如表2所示，PATHWAYS的训练吞吐量与每个流水线阶段的TPU核数成比例增加，这与其他系统一致。上述结果与下面的图5一致，表明PATHWAYS的吞吐量与主机数量呈线性关系。增加流水线阶段数会增加最小开销，当阶段数从4增加到16时，吞吐量从133.7k令牌/秒减少到131.4k令牌/秒。Google将流水线模型的性能与使用SPMD的等效模型进行了比较，观察到至少在这种情况下，流水线的性能与SPMD相当，因为SPMD计算内部聚合通信的开销高于管道气泡（bubble）更昂贵。
