[论文解读] Pathways: Asynchronous Distributed Dataflow for ML

Paul Barham, Aakanksha Chowdhery+14 more|arXiv (Cornell University)|Mar 23, 2022

Parallel Computing and Optimization Techniques被引用 37

一句话总结

Pathways 提出一个单控制器、异步分布式数据流系统，通过分片的数据流图和集中调度，在成千上万的加速器上实现高加速器利用率，并支持异质的、MPMD 模式。

ABSTRACT

We present the design of a new large scale orchestration layer for accelerators. Our system, Pathways, is explicitly designed to enable exploration of new systems and ML research ideas, while retaining state of the art performance for current models. Pathways uses a sharded dataflow graph of asynchronous operators that consume and produce futures, and efficiently gang-schedules heterogeneous parallel computations on thousands of accelerators while coordinating data transfers over their dedicated interconnects. Pathways makes use of a novel asynchronous distributed dataflow design that lets the control plane execute in parallel despite dependencies in the data plane. This design, with careful engineering, allows Pathways to adopt a single-controller model that makes it easier to express complex new parallelism patterns. We demonstrate that Pathways can achieve performance parity (~100% accelerator utilization) with state-of-the-art systems when running SPMD computations over 2048 TPUs, while also delivering throughput comparable to the SPMD case for Transformer models that are pipelined across 16 stages, or sharded across two islands of accelerators connected over a data center network.

研究动机与目标

• 需要一个可编程、可扩展的 ML 系统以支持异构且稀疏的工作负载，超越传统的 SPMD MPI 风格模型。
• 提出一种带有分片数据流和异步派发的单控制器架构，以匹配多控制器性能。
• 实现集中式资源管理和组调度，以提高利用率并支持新的并行模式。
• 提供一种编程模型，能够兼容非 SPMD 的计算并支持与 JAX/TensorFlow 等现有框架的轻松集成。

提出的方法

• 引入一个分片数据流图，其中每个节点表示一个编译后的函数，边表示数据传输。
• 实现一个集中资源管理器，在加速岛上分配虚拟设备切片并映射到物理设备。
• 使用 Plaque 异 host 协作从 Pathways IR 构建低级数据流程序，使得稀疏的高吞吐 DCN 通信成为可能。
• 结合并行异步派发以重叠主机侧工作、调度和数据传输，适用于常规（可预测）计算。
• 每个岛上使用一个中央调度器来排序执行，并实现跨众多分片的 SPMD 计算的组调度。
• 维护一个分片化对象存储，管理 CPU、GPU 和加速器内存之间的缓冲区，具备背压和 GC。

实验结果

研究问题

• RQ1 单控制器、异步数据流系统是否能够在大规模加速器上达到多控制器 ML 系统的性能？
• RQ2 集中式资源管理和组调度如何在成千上万的加速器上实现高效的 MPMD 和异构 workloads？
• RQ3 什么样的编程模型和 IR 设计能够在保留性能的同时支持非 SPMD 工作负载？
• RQ4 为了在 IPIs、DCN 和互连上实现高利用率，需要哪些数据移动与协同机制？
• RQ5 Pathways 在微基准和端到端 ML 工作负载下，与 JAX、TF、Ray 在各种配置下相比表现如何？

主要发现

• Pathways 在 2048 个 TPU 上运行 SPMD 计算时实现了大约 100% 的加速器利用率。
• Pathways 在数据中心网络连接的两个加速岛之间以两阶段管道或分片的 Transformer 模型，与 SPMD 基线相比具备可观的吞吐量。
• 并行异步派发减少主机端瓶颈，随着管道阶段增多，吞吐量优于顺序派发。
• Pathways 在足够大规模的计算上与 JAX 吞吐量相当，并扩展到成千上万的加速器，在微基准上优于某些单控制器框架。
• 该系统支持多租户，能够实现高综合吞吐量的并发程序执行和高效的资源共享。