> [!abstract] 概要(arXiv abstract の日本語訳) > 我々はアクセラレータ向けの新しい大規模オーケストレーション層の設計を提示する。我々のシステム PATHWAYS は、現行モデルに対して最先端の性能を維持しながら、新しいシステムおよび ML 研究のアイデアの探索を可能にすることを明示的に目指して設計されている。PATHWAYS は、future を消費・生成する非同期演算子からなるシャーディングされたデータフローグラフを用い、数千のアクセラレータ上で異種混在の並列計算を効率的にギャングスケジュールしながら、専用インターコネクト上でのデータ転送を協調させる。PATHWAYS は、データプレーンに依存関係があってもコントロールプレーンを並行実行できるようにする、新規の非同期分散データフロー設計を利用する。この設計は、注意深いエンジニアリングとあわせて、PATHWAYS が単一コントローラモデルを採用することを可能にし、複雑な新しい並列化パターンを表現しやすくする。我々は、PATHWAYS が 2048 台の TPU 上で SPMD 計算を実行する際に最先端システムと性能パリティ(アクセラレータ利用率 約100%)を達成できること、また 16 段にパイプライン化された Transformer モデルや、データセンターネットワークで接続された2つのアクセラレータアイランドにシャーディングされたモデルにおいても、SPMD の場合に匹敵するスループットを提供できることを示す。 ## 論文情報 - タイトル: Pathways: Asynchronous Distributed Dataflow for ML - 著者: Paul Barham, Aakanksha Chowdhery, Jeff Dean, Sanjay Ghemawat, Steven Hand, Dan Hurt, Michael Isard, Hyeontaek Lim, Ruoming Pang, Sudip Roy, Brennan Saeta, Parker Schuh, Ryan Sepassi, Laurent El Shafey, Chandramohan A. Thekkath, Yonghui Wu(全員 Google 所属) - 媒体: Proceedings of the 5th MLSys Conference(MLSys 2022) - arXiv ID: [arXiv:2203.12533](https://arxiv.org/abs/2203.12533)(2022年3月23日投稿) ## 概要 [[Google]] の分散 ML アクセラレータ実行系である PATHWAYS を提案する論文。従来の多コントローラ(multi-controller)アーキテクチャ(JAX・PyTorch・TensorFlow の SPMD 実行)は低遅延だが排他的ハードウェア占有とパイプライン・疎な計算への不適合という課題を持ち、従来の単一コントローラ(single-controller)アーキテクチャ(TensorFlow v1)は柔軟だがディスパッチ遅延とスケジューリングの不整合という課題を持つ。PATHWAYS はシャーディングされた非同期データフローとギャングスケジューリングによって、この二者の「良いところ取り」を実現する。 ## 問題設定 入力は JAX(および TensorFlow)で記述された、既知の入出力型・形状を持つ「コンパイル済み関数(compiled functions)」の列であり、出力はそれらを thousands 規模のアクセラレータ上へ配置・実行した結果である。前提として、モデルは今後 SPMD(single program multiple data)だけでなく、パイプライン化・Mixture of Experts(MoE)のような計算スパース性・複数タスク間のリソース共有(マルチテナンシ)を必要とする MPMD(multiple program multiple data)的な性質を持つようになる、という将来予測がある。 ## 提案手法 ### アーキテクチャ PATHWAYS は「リソースマネージャ」「クライアント」「コーディネーション基盤」「per-island スケジューラ」「per-device エグゼキュータ」から構成される([[非同期分散データフロー]])。 **Figure 3: PATHWAYS システム全体構成** ![[_attachments/arxiv-2203.12533/fig03-system-overview.png]] (Figure 3. 左: 分散計算は各ノードが単一のコンパイル済み関数を表す DAG として表現される。中央: リソースマネージャがアイランドのアクセラレータの部分集合(「virtual slice」)を各コンパイル済み関数に割り当てる。右: 各アイランド専用の集中型スケジューラがギャングスケジューリングを行い、per-shard エグゼキュータがディスパッチする。赤矢印はコントロールメッセージ、青矢印はデータパス転送を示す。) - **リソースマネージャ**: すべてのアイランドに跨るデバイスを集中管理する。クライアントは特定の 2D/3D メッシュ形状を持つ「virtual slice」を要求し、リソースマネージャが物理デバイスへ動的に割り当てる。 - **クライアント**: トレース済みプログラムを実行する際、まず未実行のコンパイル済み関数に virtual device を割り当ててリソースマネージャへ登録し、サーバにバックグラウンドでコンパイルを指示する。次いでデバイス位置に依存しない PATHWAYS 中間表現(IR、カスタム MLIR ダイアレクト)を構築し、標準的なコンパイラパス群で段階的に「lowering」して物理デバイス位置を含む低レベル表現へ変換する。 - **コーディネーション実装**: PATHWAYS はクロスホスト協調のすべてを既存の(非公開)本番シャーディングデータフローシステム PLAQUE に依存する。低レベル PATHWAYS IR は PLAQUE プログラムへ直接変換され、シャーディングされた計算 1 個につきノード 1 個という compact な表現を保つ(N 分割された計算 A→B の連鎖でも常に 4 ノード)。 - **データ管理**: 各ホストは Ray のオブジェクトストアに類似した、アクセラレータ HBM 上のバッファも追跡できる sharded object store を管理する。プログラムやクライアントが失敗した場合に回収できるよう、オブジェクトには所有権ラベルが付く。 多コントローラと単一コントローラのディスパッチ遅延・通信パターンの違いは次のとおり。 **Figure 1: ディスパッチオーバーヘッドと通信パターンの比較** ![[_attachments/arxiv-2203.12533/fig01-dispatch-comparison.png]] (Figure 1. (a) JAX/PyTorch の SPMD は高速な PCIe 経由で非同期にアクセラレータ計算をエンキューする。(b) TensorFlow v1 の SPMD はより低速な DCN(データセンターネットワーク)経由のコントロールメッセージを要する。(c) TensorFlow v1 の非SPMDプログラムは、明示的な send(S)・recv(R)オペレータによるクロスホスト協調やデータ転送を要する。) ### プログラミングモデル PATHWAYS は JAX と TensorFlow から書かれたソースプログラムをターゲットにする実装を持つが、本論文の評価は JAX に集中する。JAX ユーザーは標準的な Python コードをデコレータで明示的にラップし、(潜在的に SPMD な)XLA 計算へコンパイルされる断片を示す。これらは既知の入出力型・形状、有界ループ、少ない(あるいは無い)条件分岐を特徴とするため事前にリソース要求を見積もれ、これを「コンパイル済み関数」と呼ぶ。 ```python def get_devices(n): """Allocates `n` virtual TPU devices on an island.""" device_set = pw.make_virtual_device_set() return device_set.add_slice(tpu_devices=n).tpus a = jax.pmap(lambda x: x * 2., devices=get_devices(2)) b = jax.pmap(lambda x: x + 1., devices=get_devices(2)) c = jax.pmap(lambda x: x / 2., devices=get_devices(2)) @pw.program # Program tracing (optional) def f(v): x = a(v) y = b(x) z = a(c(x)) return (y, z) print(f(numpy.array([1., 2.]))) # output: (array([3., 5.]), array([2., 4.])) ``` (Figure 2. PATHWAYS 上で複数の TPU アイランドに跨るシャーディング計算を実行する Python ユーザーコード例。) ユーザーは型・場所・インターコネクトトポロジに任意の制約を付けた「virtual devices」の集合を要求し、特定のコンパイル済み関数をそのデバイス上に配置できる。デフォルトでは各コンパイル済み関数が単一の(シャーディングされた)計算を含む独立した PATHWAYS プログラムに変換されるため、多数の関数を連続実行したいときは各関数ごとに個別の Python 呼び出しと RPC が必要になる。そこで PATHWAYS はプログラムトレーサを実装し、各コンパイル済み関数を1つの計算ノードとするデータフローグラフを持つ単一の PATHWAYS プログラムを生成できるようにした。JAX を PATHWAYS バックエンドへ差し替えると、JAX コードはほぼ無変更で動作するうえ、SPMD 計算がローカル接続の TPU コアだけでなくシステム上にプロビジョニングされたすべてのコアへアクセスできるようになり、複数の TPU pod をまたぐ初めてのスケーリングが可能になる。 ### アルゴリズム/手法の詳細: ギャングスケジューリングと並列非同期ディスパッチ TPU は単一スレッドで非プリエンプティブルなカーネルのみを実行するため、通信を伴う計算がホスト間で一貫した順序でエンキューされなければデッドロックする。このため PATHWAYS は各アイランドに専用の集中型スケジューラを持ち、プログラム間をまたいで一貫した実行順序を強制する「ギャングスケジューリング」を実装する。 **Figure 4: 逐次ディスパッチと並列ディスパッチの比較** ![[_attachments/arxiv-2203.12533/fig04-sequential-vs-parallel-dispatch.png]] (Figure 4. 3 ノードプログラムに対する逐次ディスパッチ(a)と並列ディスパッチ(b)。計算のデバイス上実行時間がスケジューリング・リソース割り当て・協調に費やす時間より短いとき、逐次ディスパッチではホスト側処理待ちで非同期パイプラインがストールする。並列非同期ディスパッチはホスト側処理を並行実行することでこのボトルネックを解消する。) 素朴な設計では、ホスト A がノード A をエンキューして出力の future を受け取り、それをホスト B へ送信し、ホスト B がノード B の入力を確保して起動準備を整える、という逐次的な流れになる。計算時間がホスト側処理(スケジューリング・リソース割り当て・協調)より短いと、非同期パイプラインがストールしてホスト側処理が全体のボトルネックになる(Figure 4a)。PATHWAYS は、コンパイル済み関数のリソース要求が静的に既知であることを利用し、後続ノードの入力形状を先行計算の完了前に計算できる点に着目して、新規の**並列非同期ディスパッチ**を導入した(Figure 4b)。静的にスケジュール可能なサブグラフに対しては、サブグラフ全体を記述する単一メッセージをスケジューラへ送ることでネットワークトラフィックを最小化する。データ依存の制御フローのようにノードのリソース要求が先行計算の完了まで未知の場合は、従来の逐次モデルへフォールバックする。 ### 実装上の工夫 - コーディネーション基盤には既存の PLAQUE(汎用データフローエンジン)を再利用し、DCN 通信の背景処理(設定情報の配布・プログラムのモニタリング・クリーンアップ・失敗時のエラー配信など)にも同じ基盤を用いる。 - Ray のような他の分散フレームワークで PATHWAYS を再実装することも可能と考えられるが、Ray は HBM オブジェクトストアや GPU インターコネクト経由のリモートオブジェクト転送プリミティブを欠くため、同等の性能を得るには追加実装が必要と述べる。 - PATHWAYS の単一コントローラ設計は、将来的にクライアントのリソースを協調なしに一時的に回収・再割り当てする透過的な suspend/resume・migration をサポートする余地を与える。 ## 新規性 既存の多コントローラ系(MPI・JAX・PyTorch・TensorFlow の最近の構成)は PCIe 経由の低遅延ディスパッチを持つが、排他的ハードウェア占有を前提とするため、リソース仮想化・マルチプレキシングの実現が難しい。既存の単一コントローラ系(TensorFlow v1)は柔軟な分散データフローとリソース仮想化を提供するが、DCN 経由のディスパッチ遅延が大きく、フルにマテリアライズされたシャーディング計算グラフがシャード数千規模で深刻なオーバーヘッドを生む。PATHWAYS は、[[非同期分散データフロー]]によってコントロールプレーンをデータプレーンの依存関係と並行実行させることで、単一コントローラの柔軟性と多コントローラの性能を両立させた点が新規性である。 ## 実験設定 - **ハードウェア構成**: Configuration (A) は 1 ホストあたり 4 TPU で最大 512 ホスト(2048 TPU、ICI 接続)。Configuration (B) は 1 ホストあたり 8 TPU で最大 64 ホスト(512 TPU)。Configuration (C) は 4 アイランド構成で各アイランド 4 ホスト・32 TPU。Ray/GPU の評価には Ray v1.3 + PyTorch 1.8.1 を p3.2xlarge VM(1×V100 GPU、8×CPU コア)上で実行し、Amazon placement group でスケジューリングした。 - **比較対象**: 主に多コントローラ JAX(業界標準ベンチマークで最先端性能を示すため)。マイクロベンチマークでは TensorFlow(TF)・Ray も比較。 - **評価指標**: 1 秒あたりの実行計算数(computations per second、スループット)、および実モデルの学習スループット(tokens/秒)。 ## 実験結果 ### 単一コントローラのディスパッチオーバーヘッド スカラーの AllReduce に続けてスカラー加算を行う自明なギャングスケジュール計算を、OpByOp(逐次呼び出し)・Chained(128 ノードの連鎖を単一呼び出しで実行)・Fused(単一ノードが 128 計算の連鎖を内包)の3方式で反復実行し、スループットを比較する。 **Figure 5: PATHWAYS のディスパッチオーバーヘッド比較** ![[_attachments/arxiv-2203.12533/fig05-dispatch-overhead.png]] (Figure 5. PATHWAYS は TF・Ray のような単一コントローラ系をすべての構成で上回り、Fused(-F)・Chained(-C)構成では最大 1000・256 TPU コアまで多コントローラ JAX の性能に一致する。各計算はスカラー AllReduce とスカラー加算の連鎖からなる。) OpByOp は最悪ケースの負荷実験であり、いずれのフレームワークにも自然な使い方ではないが、アクセラレータ数が増えるほど JAX 多コントローラが単一コントローラ系より優位になる。PATHWAYS のオーバーヘッドの大半は、クライアントがコーディネータからの出力ハンドル返却を待って次の計算をエンキューする点に起因する。十分な計算が Fused されれば PATHWAYS は最大 1000 TPU コアまで JAX と一致し、PATHWAYS Chained は最大 256 コアまで JAX OpByOp を上回る(C++ から直接連続してアクセラレータ計算をディスパッチできるため)。TensorFlow は集中バリア(control edge で実装)によるギャングスケジュール計算の直列化のため多コア環境で遅く、Ray は GPU→DRAM のコピーと Python アクター実行のオーバーヘッドを持つ。 **Figure 6: PATHWAYS と JAX のスループット一致点** ![[_attachments/arxiv-2203.12533/fig06-throughput-parity.png]] (Figure 6. Configuration (B) の 16 ホスト・128 TPU では計算サイズ 2.3 ms 以上で PATHWAYS が JAX のスループットに一致し、Configuration (A) の 512 ホスト・2048 TPU では計算サイズ 35 ms 以上で一致する。) この結果は、計算時間が十分に長ければ(2048 TPU 規模でも 35 ms 以上)PATHWAYS の単一コントローラオーバーヘッドが完全にマスクされることを示す。実モデルではスカラー計算より大きいため、§5.3 の結果ではこの一致が常に成立する。 次のマイクロベンチマークは§4.5 の並列非同期ディスパッチ機構の効果を評価する。各計算を異なるホスト上の 4 TPU コアの集合上で実行し、データはステージ間を ICI 経由で送る、より実際的なパイプラインベンチマークである。 **Figure 7: PATHWAYS における並列ディスパッチと逐次ディスパッチの比較** ![[_attachments/arxiv-2203.12533/fig07-parallel-vs-sequential-pipeline.png]] (Figure 7. パイプラインステージごとに異なる 4 TPU コア(別ホスト)を使用し、ICI 経由で次ステージへデータを転送する。並列非同期ディスパッチは、パイプラインステージ数が多い場合に固定クライアントオーバーヘッドとスケジューリングオーバーヘッドを償却する。) ホスト数増加に伴い固定クライアントオーバーヘッドが償却される段階、ステージ増加に伴う転送コストが優勢になる段階、固定スケジューリングオーバーヘッドが償却される段階という3つの「フェーズ」が観測される。逐次非同期ディスパッチに強制した場合との比較で、並列非同期ディスパッチの利得が確認できる。 ### マルチテナンシ Configuration (B) 上で、複数クライアントが同時に異なる PATHWAYS プログラムを投入する状況を評価する。 **Figure 8: 並行プログラムの集約スループット** ![[_attachments/arxiv-2203.12533/fig08-multitenancy-throughput.png]] (Figure 8. PATHWAYS は複数プログラム間でコンテキストスイッチのオーバーヘッドなしに、JAX と少なくとも同等の集約スループットを達成する(括弧内は計算時間 ms)。) PATHWAYS の最大スループットが小さな計算では JAX を上回るのは、リモートクライアントから受け付けられる計算数が、JAX がローカルに Python からディスパッチできる数より多いためである。 **Figure 9: 4 クライアント間のギャングスケジュール計算のインターリーブトレース** ![[_attachments/arxiv-2203.12533/fig09a-multitenancy-trace-proportional-1111.png]] ![[_attachments/arxiv-2203.12533/fig09b-multitenancy-trace-proportional-1248.png]] (Figure 9. 128 コアのサンプルにおける PATHWAYS 上のトレース。4 クライアント間で proportional-share 比率 1:1:1:1(上)・1:2:4:8(下)のギャングスケジュール並行プログラムのインターリーブを示す。) この実験は、PATHWAYS が 4 つの独立クライアントから投入されたプログラムをギャングスケジュールしつつ、公平性のためにアクセラレータ時間の割り当てを制御できることを示す。付録 D の Figure 11 は、この負荷でクライアント数を 1・4・8・16 と変えたときのコアトレースを示す。 **Figure 11: クライアント数を変えたマルチテナンシトレース(付録 D)** ![[_attachments/arxiv-2203.12533/fig11a-multitenancy-trace-1client.png]] ![[_attachments/arxiv-2203.12533/fig11b-multitenancy-trace-4clients.png]] ![[_attachments/arxiv-2203.12533/fig11c-multitenancy-trace-8clients.png]] ![[_attachments/arxiv-2203.12533/fig11d-multitenancy-trace-16clients.png]] (Figure 11. (a) 1 クライアント (b) 4 クライアント (c) 8 クライアント (d) 16 クライアント。単一クライアントでは 1 プログラムあたりの計算時間が 0.33 ms と小さくアクセラレータを飽和させられないが、PATHWAYS のマルチテナンシ支援により複数クライアントを使うことで利用率が約100%まで上がる。全クライアントプログラムはミリ秒スケール以下でインターリーブされ、コンテキストスイッチのオーバーヘッドはほとんど見られない。) ### 大規模モデル性能 JAX・TF のモデルをネイティブシステムと PATHWAYS 上で実行して、数値結果が一致することを確認した上で性能のみを比較する。 **表1: Text-to-Text Transformer(T5)の学習スループット(tokens/s)** | Model | Params | TPU cores | JAX | PATHWAYS | |---|---|---|---|---| | T5-Base | 270M | 32 | 618k | 618k | | T5-Large | 770M | 32 | 90.4k | 90.4k | | T5-3B | 3B | 512 | 282.8k | 282.8k | | T5-11B | 11B | 512 | 84.8k | 84.8k | (Table 1. [Raffel+, 2019] のモデル構成を用いた Text-to-text Transformer の学習スループット。モデルコードが同一のため、JAX と PATHWAYS は全モデルサイズで同一の学習ステップ数・同一のパープレキシティに到達し、性能も一致する。) 続いて、Decoder-only アーキテクチャの Transformer 言語モデル(62 層、モデル次元 2048、隠れ次元 8192、合計 30 億パラメータ)を Configuration (B)・(C) 上で学習し、SPMD 構成と GPipe 型スケジュールによるパイプライン構成を比較した。 **表2: 3B Transformer 言語モデルの学習スループット(tokens/s)。S はパイプラインステージ数、M はマイクロバッチ数** | Model configuration | TPU cores | PATHWAYS | |---|---|---| | Model-parallel(SPMD) | 128 | 125.7k | | Pipelining, S=4, M=16 | 128 | 133.7k | | Pipelining, S=8, M=32 | 128 | 132.7k | | Pipelining, S=16, M=64 | 128 | 131.4k | | Pipelining, S=16, M=64 | 512 | 507.8k | パイプラインステージ数を 4 から 16 に増やしてもスループットの低下は 133.7k から 131.4k tokens/秒とわずかであり、この事例では SPMD の集合通信オーバーヘッドがパイプラインバブルオーバーヘッドより大きいため、[[パイプライン並列化]]は SPMD に対して競争力のある性能を示す(Source: 本ソース Table 2)。 さらに PATHWAYS は DCN で接続された TPU アイランドをまたいでも効率的に学習できることを示した。S=16, M=64・128 コア構成では、Configuration (B) の単一アイランド 128 コアと、Configuration (C) の 4 アイランド × 32 コアとで同じスループット(131.4k tokens/秒)を達成した。 **Figure 10: 128 TPU にパイプライン化した 3B Transformer モデルのトレース** ![[_attachments/arxiv-2203.12533/fig10-pipeline-trace-128tpu.png]] (Figure 10. 3B Transformer モデルを 128 TPU にパイプライン化: PATHWAYS は DCN で接続された TPU アイランド群を跨いで、Configuration (C) の 4 アイランド×32 コアで Configuration (B) の単一アイランド 128 コアと同じスループット(131.4k tokens/秒)を効率的に学習できる。) DCN 転送はトレース中のコアグループ 8 行ごとに発生するが、通信時間が計算と実質的に重なるためトレースには視認されない。 最後に、64B・136B パラメータの Decoder-only Transformer を、2つのアクセラレータアイランドを使って学習した。PATHWAYS は 2 アイランド構成で、2倍のデバイス数を持つ単一アイランド構成と比べて約97%のスループットを達成した(136B モデルは 2 アイランド×1024 コア、64B モデルは 2 アイランド×512 コア、アイランド内は高速な ICI でリダクションし、アイランド間は DCN 転送)。 **Figure 12: 64B Transformer の2アイランド DCN トレース(付録 D)** ![[_attachments/arxiv-2203.12533/fig12-two-island-dcn-trace-64b.png]] (Figure 12. 512 TPU ずつの2アイランドでデータ並列学習する 64B Transformer モデル。トレースはアイランド間 DCN 転送の相対的な小さいオーバーヘッドを示す。上段8行(青)は第1アイランドのホスト上の TPU 計算、下段8行(緑)は第2アイランドのホスト上の TPU 計算に対応する。各アイランドは勾配を計算してから他方のアイランドへの勾配転送をエンキューし、転送完了後に受信した勾配を適用して次の学習ステップへ進む。付録 D では、128 ホストのペア規模でも DCN 転送のオーバーヘッドが最小限であり、ICI 通信を用いる等価な SPMD 構成に対して 97.2% の学習スループットを達成すると報告される。) ## 考察 - PATHWAYS は TPU の特性(長時間実行かつプリエンプション不能なカーネル、リッチな制御フローと通信プリミティブを XLA コンパイラで単一カーネルへ融合できる)を前提に低レベル設計判断を下しているが、著者らは高レベルのアーキテクチャ選択の大部分は大規模 GPU システムにも有効だろうと考察する。 - PATHWAYS の単一コントローラ設計は、利用可能なリソースを大規模に追跡・割り当てる能力をシステムへ与えるため、優先度・性能分離・アクセスコントロール・リソース会計といった多テナンシ要件を、既存研究より小さいタイムスケールかつ大きいリソースプールで探索する計画がある(§6.2)。 - 著者らは、Mixture of Experts やルーテッドカプセルネットワークのようなデータ依存のベクトル化制御フローを、良好な性能を保ったまま表現するクリーンなプログラミングモデルの提供を将来課題としている(§6.3)。 ## 強み / 弱点・課題 - **強み**: 単一コントローラの柔軟性(複雑な非 SPMD・MPMD 計算の表現、集中的なリソース管理・仮想化)と、多コントローラ並みの性能を、非同期分散データフローと並列非同期ディスパッチの組み合わせで両立させた。2048 TPU 規模の SPMD、16 段パイプライン、2 アイランド DCN 構成のいずれでも実測により性能パリティを裏付けている。 - **弱点・課題**: コーディネーション基盤 PLAQUE は Google 内部の非公開システムであり、外部での再現には Ray 等への再実装が必要になると著者ら自身が述べる。また、小さな計算(スカラー規模)では単一コントローラのオーバーヘッドがまだ露出しており(Figure 5・6)、著者らはこの領域の最適化には注力していないと明言する。データ依存のベクトル化制御流(MoE のルーティング等)への対応は本論文の範囲外で将来課題として残されている。