[論文レビュー] 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.

研究の動機と目的

• 従来の SPMD MPIスタイルのモデルを超える、異種かつスパースなワークロードをサポートする、プログラム可能でスケーラブルな ML システムの必要性を動機づける。
• シャード化されたデータフローと非同期ディスパッチを備えた単一コントローラアーキテクチャを提案し、マルチコントローラの性能に匹敵させる。
• 利用率を向上させ、新規の並列性パターンをサポートするために、集中型リソース管理と gang scheduling を有効にする。
• 非-SPMD計算を受け入れ、JAX/TensorFlow など既存フレームワークと容易に統合できるプログラミングモデルを提供する。

提案手法

• 各ノードがコンパイル済み関数を表し、エッジがデータ転送を表すシャード化されたデータフローグラフを導入する。
• アクセラレータ島上に仮想デバイススライスを割り当て、それを物理デバイスへマッピングする集中型リソースマネージャを実装する。
• Pathways IR から疎結合で高スループットな DCN 通信を可能にする、 Plaque ベースのホスト間協調を使用して低レベルのデータフロー プログラムを形成する。
• 定期的（予測可能な）計算のために、ホスト側の作業とスケジューリングおよびデータ転送を重ね合わせるための parallel asynchronous dispatch を組み込む。
• 多くのシャードに跨る SPMD 計算の実行順序を決定し、gang-scheduling を有効にするために、島ごとに中央スケジューラを配置する。
• CPU、GPU、アクセラレータメモリ全体のバッファを管理するシャード化されたオブジェクトストアを維持し、バックプレッシャーと GC を伴う。

実験結果

リサーチクエスチョン

• RQ1 単一コントローラ型の非同期データフローシステムは、大規模なアクセラレータ上でマルチコントローラ ML システムの性能に匹敵できるのか？
• RQ2 集中型リソース管理と gang scheduling により、数千台のアクセラレータにわたる効率的な MPMD および異種ワークロードをどう実現できるか？
• RQ3 パフォーマンスを維持しつつ、非-SPMD ワークロードをサポートするプログラミングモデルと IR 設計とは何か？
• RQ4 IPI、DCN、インターコネクト全体で高い利用率を実現するために必要なデータ移動と調整メカニズムは何か？
• RQ5 さまざまな構成で、マイクロベンチマークおよびエンドツーエンドの ML ワークロードにおける Pathways の性能は、JAX、TF、Ray と比べてどうなるか？

主な発見

• Pathways は 2048 TPUs を横断して SPMD 計算を実行する際、アクセラレータ利用率をほぼ 100% に達成する。
• Pathways は、データセンター網に接続された 2 つのアクセラレータ島の間でパイプライン化された 16 段階のトランスフォーマーモデル、またはシャードされた場合のSPMDベースラインと同等のスループットを提供する。
• 並列非同期ディスパッチはホスト側のボトルネックを低減し、パイプライン段数が増えるほど、逐次ディスパッチよりもより良いスループットを生み出します。
• Pathways は十分に大きな計算に対して JAX のスループットに一致し、数千のアクセラレータにスケールし、マイクロベンチマークでいくつかの単一コントローラ・フレームワークを上回る。
• このシステムはマルチテナンシーをサポートし、高い総スループットと効率的なリソース共有を実現する同時実行を可能にする。