解读谷歌Pathways架构（二）：向前一步是OneFlow

OneFlow深度学习框架

已于 2022-04-13 21:53:07 修改

阅读量5k

点赞数 1

分类专栏：业界观点文章标签：分布式人工智能大数据深度学习机器学习

于 2022-04-11 09:06:23 首次发布

本文链接：https://blog.csdn.net/OneFlow_Official/article/details/124113864

版权

本文深入探讨了谷歌的Pathways架构及其与OneFlow的关系，重点解析了Pathways的数据流执行引擎PLAQUE，强调了OneFlow在分布式深度学习框架中的独特优势。Pathways通过sharded dataflow graph实现异步分布式计算，而OneFlow在静态调度、内存管理和并行异步分发方面展现出更高的效率。文章还指出，OneFlow的设计理念和实现方式在某些方面超越了Pathways，暗示OneFlow可能代表着更先进的AI架构方向。

摘要由CSDN通过智能技术生成

撰文｜袁进辉

本文目录如下：

下一代 AI 架构是否成立？
问题设定
Pathways 如何工作？
数据流执行引擎 PLAQUE
关于 PLAQUE 的一些问题
从 Pathways 的视角理解 OneFlow SBP
从 Pathways 的视角理解 OneFlow Actor
OneFlow 可以向 Pathways 学习什么？
Ray 距离Pathways 有多远？
总结

下一代 AI 架构是否成立？

多数情况下，在学术界提“下一代”、“重新思考”、“重新设计” 这样的字眼是一件很讨人厌的事，其令人反感的程度堪比“我不是针对谁，我是说在座的都不如我厉害”，归根结底，不是谁都可以站出来定义真正的“下一代”。

最近几年，在分布式计算系统领域，印象中有 3 篇论文这么做过：

第一篇是《Ray: a distributed framework for emgerging AI applications（https://www.usenix.org/system/files/osdi18-moritz.pdf）》，主张面向未来的 AI 应用必须重新思考分布式计算框架，结果就是 Ray。这篇论文据说也被拒过，不过最终被 OSDI 录用。其作者是伯克利 AMP Lab 的 Ion Stoica，其神奇履历见访谈《Ion Stoica: 做成 Spark 和 Ray 两个明星项目的秘笈》。虽然 Ray 以强化学习场景为目标，但我所知道的大规模的强化学习系统都不是用它实现的，原因是用 Ray 时遇到无法搞定的问题。我个人认为 Ray 已经很有影响力，但仍未解决自己的定位问题，它好似介于 K8S 和深度学习框架之间的位置，但在这两个领地都无法真正实现取而代之。

第二篇是OneFlow团队写的《OneFlow: Redesign the Distributed Deep Learning Framework from Scratch（https://arxiv.org/pdf/2110.15032.pdf）》，论文称，为了灵活、高效的支持近些年的大模型需求，不能在已有深度学习框架基础上修修补补，必须另起炉灶，相对于在已有框架上做增量式改进，从零开始设计的 OneFlow 更有优势。不幸的是，这篇论文被 OSDI, SOSP, MLSys 拒稿。为什么要重新设计？在原有系统上改造不香吗？你能证明在原有系统上增量修改不能实现一样的目标吗？

第三篇就是本文要进一步讨论的《Pathways: asynchronous distributed dataflow for ML》，虽然论文标题没有写下一代 AI 架构，但作者之一 Jeff Dean 投稿时在 Google 官网发了一篇震动世界的博客《Pathways：下一代 AI 架构》。如果从 SPMD 的需求来看框架，真是再简单不过了，但是，MPMD 是由多个 SPMD 组成。Pathways 要思考为 MPMD 服务的最优架构，就要跳出传统深度学习框架的藩篱，在更高层次进行设计。如果把之前的 PyTorch 和 JAX 等为 SPMD 场景设计的深度学习框架理解成狭义的框架，那么 Pathways 之所以敢自称下一代 AI 架构，是因为它可以被理解成框架的框架。

一般来说，这样的论文立意宏大，难写。在论文里论证研发下一代架构的必要性也很难，必须要有理有据的证明。强如 Google Brain 这样的团队也只能说，尽管 TF v1 可以做到这个以及那个，但它有这样及那样的问题，需求发生变化了，我们需要重新思考这个问题。

The implementation choices made by TF v1 were over-specialized to assume a single, smallish, exclusively-owned island of accelerators. This over-specialization makes it practically infeasible to use TF for contemporary or future ML workloads.

但这个论证并不能以理服人。毕竟，为了训练 GPT-3， TensorFlow 团队还是研发了Mesh-tensorflow、GPipe、GShard、GSPMD，虽然 PyTorch 还没有解决这些问题，但英伟达在其基础上做了Megatron-LM，微软做了 DeepSpeed 都还可以训练大模型和 MoE，用户也不少，你怎么能说人家这些增量式的改进行不通？

这样的论文也难懂。在上一篇《解读Pathways（一）：Single-controller与Multi-controller》中，我说 Pathways 的论文不好理解，因为绝大部分人思维的舒适区还是在 SPMD，理解起来 MPMD 不那么顺理成章。

所以，写文章解读 Pathways 也有一定的价值吧。

问题设定

为便于理解 Pathways 的设计理念，我们先简要回顾一下它的设计目标是什么。Pathways 认为未来的 AI 模型不是 SPMD 所能描述的，得用MPMD（譬如流水并行这样的任务）。简单来说， MPMD 是由多个 SPMD 构成的，每个 SPMD 是同一个 computation 在一组设备上通过 GShard 或 GSPMD 的方式来均匀、对称的划分。我们用 Pathways 论文里的一个例子来说明。

上图展示了 Pathways 可以运行的 MPMD 类型的 program，这个 program 由 a, b, c, a 四个 computation 构成，每个 computation 都是一个 SPMD 类型的任务，每个 computation 需要被 sharding 到若干设备上去，例如 a，b，c 都被分配了 2 个设备。

上图则展示了 Pathways 执行 Figure 2 代码时内部发生的事情。需要注意的是，尽管 A, B, C 每一个 computation 在多个设备上并行执行，数据和代码都对应着多份，但最左边的数据流图仍把每个 computation 抽象成数据流图的一个节点来表示。

集群中设备被分成很多组，每个组包含一批互联成 mesh 拓扑的同质设备，称为一个 island （一般表示一个 POD），一个 island 内部的设备可以通过称为 ICI 的高速连接通信，island 之间则需要通过带宽低一些的DCN（RDMA）来通信。每个 island 上有一个中心化的 scheduler 来实现该 island 上的群调度，一个island 内部每个设备都有一个运行在对应 host (CPU）上的 executor，同一个 island 上的 SPMD 内部会发生集群通信，相邻的 SPMD 之间也会通过 ICI 传输中间结果。

Pathways 如何工作？

首先对 Pathways 系统做一个鸟瞰式介绍，再去深入讨论每一个具体问题。

第①步，用户想运行一个由多个 computations（SPMD）构成的 traced program（MPMD）时，通过调用 Client library 来使用 Pathways 系统。

第②步，Client 为之前没有运行过的 computation 分配虚拟设备（virtual devices），并向 Resource Manager 注册这个 computation，也就是获得虚拟设备对应的物理设备（physical devices）。

第③步，client 触发后台服务器对 computations 进行编译：首先，基于 MLIR 的 dialect 构造一个设备无感（location-agnostic）的中间表示（IR），这个 IR 通过一系列标准的编译（pass）被逐步降低（lowered）成一个感知设备物理位置的底层表示。值得一提的是，这个底层表示已经考虑了物理设备之间的网络连