机器之心报道
机器之心编辑部

「当前的 AI 模型只做一件事 。 Pathways 使我们能够训练一个模型 ， 做成千上万件事情 。 」

在谈及当前的 AI 系统所面临的问题时 ， 低效是经常被提及的一个 。
谷歌人工智能主管 Jeff Dean 曾在一篇博文中写道 ， 「今天的人工智能系统总是从头开始学习新问题 —— 数学模型的参数从随机数开始 。 就像每次学习一项新技能（例如跳绳） ， 你总会忘记之前所学的一切 ， 包括如何平衡、如何跳跃、如何协调手的运动等 ， 然后从无到有重新学习 。 这或多或少是我们今天训练大多数机器学习模型的方式：我们不是扩展现有模型来学习新任务 ， 而是从无到有训练新模型来做一件事（或者我们有时将通用模型专门用于特定任务） 。 结果是我们最终为数千个单独的任务开发了数千个模型 。 以这种方式学习每项新任务不仅需要更长的时间 ， 而且还需要更多的数据 。 」
为了改变这种局面 ， Jeff Dean 等人去年提出了一种名叫「Pathways」的通用 AI 架构 。 他介绍说 ， Pathways 旨在用一个架构同时处理多项任务 ， 并且拥有快速学习新任务、更好地理解世界的能力 。

文章图片

该架构的特点可以概括为：

• 能够训练一个模型来做成千上万件事情；
• 当前模型只注重一种感官 ， Pathways 可做到多种；
• 当前模型密集且效率低下 ， Pathways 会把模型变得稀疏而高效 。

在发布想法大半年之后 ， Jeff Dean 终于公布了 Pathways 的论文 ， 其中包含很多技术细节 。

文章图片

论文链接 https://arxiv.org/pdf/2203.12533.pdf
论文写道 ， PATHWAYS 使用了异步算子的一个分片数据流图（sharded dataflow graph） ， 这些算子消耗并产生 futures ， 并在数千个加速器上高效地对异构并行计算进行 gang-schedule ， 同时在它们专用的 interconnect 上协调数据传输 。 PATHWAYS 使用了一种新的异步分布式数据流设计 ， 它允许控制平面并行执行 ， 尽管数据平面中存在依赖关系 。 这种设计允许 PATHWAYS 采用单控制器模型 ， 从而更容易表达复杂的新并行模式 。
实验结果表明 ， 当在 2048 个 TPU 上运行 SPMD（single program multiple data）计算时 ， PATHWAYS 的性能（加速器利用率接近 100%）可以媲美 SOTA 系统 ， 同时吞吐量可媲美跨越 16 个 stage 或者被分割成两个通过数据中心网络连接的加速器岛的 Transformer 模型的 SPMD 案例 。
以下是论文细节：
研究背景
在过去的十年里 ， 深度学习在图像理解、自然语言处理等多个领域取得了显著的进展 ， 这是 ML 模型、加速器硬件以及将两者联系在一起的软件系统协同进化的结果 。 这种协同进化带来的隐患是：深度学习系统可能过度专注于当前的工作负载 ， 无法预测未来的需求 。
PATHWAYS 是一个为分布式 ML 构建的新系统 ， 剑指未来 ML 工作负载将需要的特定能力 。 当前 ， 这些工作负载缺乏 SOTA 系统的支持 。
例如 ， 当今 SOTA ML 工作负载大多使用单程序多数据（SPMD）模型 ， 该模型受到了 MPI 的启发 ， 其中所有加速器都在同步运行相同的计算 ， 加速器之间的通信由 AllReduce 等集体来描述 。
但近年来 ， 研究人员开始在 ML 计算中被 SPMD 掣肘 。 大型语言模型已经使用流水线并行而不是纯粹的数据并行来扩展；混合专家（MoE）等模型已经开始探索计算稀疏性 ， 其最自然的表达方式是使用细粒度控制流和跨加速器的异构计算；系统设计者们已经开始采用巧妙的技术来在 MPI 风格的系统上执行流水线（pipelined）、同构 MoE 模型 ， 但是 ， MPI 编程模型对于用户和底层系统来说都太受限制了 。