半年以上前から推進されてきたGoogleの次世代AIアーキテクチャとジェフ・ディーンのPathwaysがついに論文化

現在の AI システムが直面している問題について議論する際、非効率性はよく言及されるものの 1 つです。

Google の人工知能担当ディレクターのジェフ・ディーンは、かつてブログ記事にこう書いています。「今日の人工知能システムは常に新しい問題をゼロから学習します。数学モデルのパラメータはランダムな数字から始まります。新しいスキル (縄跳びなど) を習得するたびに、バランスの取り方、ジャンプの仕方、手の動きの調整方法など、以前に学習したことはすべて忘れてしまい、ゼロから再学習するのと同じです。これは、今日のほとんどの機械学習モデルをトレーニングする方法とほぼ同じです。既存のモデルを拡張して新しいタスクを学習させるのではなく、1 つのことを行う新しいモデルをゼロからトレーニングします (または、一般的なモデルを特定のタスクに特化させることもあります)。その結果、何千もの個別のタスクに対して何千ものモデルを開発することになります。この方法で新しいタスクを 1 つずつ学習すると、時間がかかるだけでなく、より多くのデータが必要になります。」

この状況を変えるために、ジェフ・ディーン氏らは昨年、「Pathways」と呼ばれる汎用AIアーキテクチャを提案しました。彼は、Pathways は 1 つのアーキテクチャで複数のタスクを同時に処理するように設計されており、新しいタスクを迅速に学習して世界をよりよく理解する能力を備えていると紹介しました。

このアーキテクチャの特徴は次のようにまとめられます。

• 何千ものことを実行するようにモデルをトレーニングする機能。
• 現在のモデルは 1 つの感覚のみに焦点を当てていますが、Pathways は複数の感覚に対応できます。
• 現在のモデルは密で非効率的ですが、Pathways によりモデルが疎で効率的になります。

アイデアを発表してから半年以上経って、ジェフ・ディーンはついに多くの技術的な詳細を含む Pathways 論文を発表しました。

論文リンク: https://arxiv.org/pdf/2203.12533.pdf

この論文では、PATHWAYS は、未来を消費および生成する非同期演算子のシャード化されたデータフロー グラフを使用し、専用の相互接続を介してデータ転送を調整しながら、数千のアクセラレータ上で異種の並列計算を効率的にギャング スケジュールすると述べられています。 PATHWAYS は、データ プレーンの依存関係にかかわらず、コントロール プレーンを並列実行できる新しい非同期分散データフロー設計を使用します。この設計により、PATHWAYS は単一のコントローラー モデルを採用できるようになり、複雑な新しい並列パターンをより簡単に表現できるようになります。

実験結果によると、2048 TPU で SPMD (シングル プログラム マルチ データ) 計算を実行すると、PATHWAYS のパフォーマンス (アクセラレータの使用率は 100% に近い) は SOTA システムに匹敵し、スループットは 16 ステージにまたがる、またはデータ センター ネットワークで接続された 2 つのアクセラレータ アイランドに分割された Transformer モデルの SPMD ケースに匹敵します。

論文の詳細は次のとおりです。

背景

過去 10 年間で、ディープラーニングは画像理解や自然言語処理など多くの分野で目覚ましい進歩を遂げてきました。これは、ML モデル、アクセラレータ ハードウェア、そしてこの 2 つを結び付けるソフトウェア システムの共進化の結果です。この共進化の隠れた危険性は、ディープラーニング システムが現在のワークロードに過度に集中し、将来のニーズを予測できなくなる可能性があることです。

PATHWAYS は、将来の ML ワークロードに必要な特定の機能をターゲットにして、分散 ML 向けに構築された新しいシステムです。現在、これらのワークロードは SOTA システムからサポートされていません。

たとえば、今日の SOTA ML ワークロードのほとんどは、MPI にヒントを得た Single Program Multiple Data (SPMD) モデルを使用しています。このモデルでは、すべてのアクセラレータが同じ計算を同期的に実行し、アクセラレータ間の通信は AllReduce などの集合によって記述されます。

しかし近年、研究者は ML 計算において SPMD による制約を受け始めています。大規模な言語モデルは、純粋なデータ並列処理ではなく、パイプライン並列処理を使用して拡張されてきました。Mixture of Experts (MoE) などのモデルは、きめ細かい制御フローとアクセラレータ間の異種計算を使用して最も自然に表現される計算スパース性を検討し始めています。システム設計者は、MPI スタイルのシステムでパイプライン化された同種の MoE モデルを実行するための巧妙な手法を採用し始めていますが、MPI プログラミング モデルは、ユーザーと基盤となるシステムの両方にとって制限が厳しすぎます。

一方、新しいアクセラレータ世代の出現により、ML における異種環境はますます一般的になりつつあります。高帯域幅の相互接続によって接続された同種のアクセラレータの大規模な「島」への排他的アクセスを提供するには、すべてのアクセラレータを常にビジー状態に保つために単一のユーザー プログラムが動作する必要があるため、コストがかかり、無駄になることがよくあります。この制限により、研究者はさらに「複数プログラム複数データ」(MPMD) コンピューティングへと進みました。 MPMD は、全体的な計算の一部を、より小さく、より容易に利用できるアクセラレータ アイランドのセットにマッピングすることで、柔軟性を高めます。使用率を向上させるために、一部の ML ハードウェア リソース管理研究者は、ワークロード間でハードウェアをきめ細かく再利用し、ワークロードの弾力性を実現し、フォールト トレランスを向上させています。

最終的に、研究者たちは、ビッグデータで大規模にトレーニングされ、さまざまな下流のタスクに適応できる一連の基礎モデルを標準化し始めました。このようなモデルのトレーニングと推論により、多くのタスク間でリソースを再利用し、それらの間で状態を効率的に共有することで、クラスターの使用率を向上させる機会が提供されます。たとえば、複数の研究者が、同じアクセラレータを使用してベースモデルレイヤーを固定したまま、異なるタスクに合わせてベースモデルを同時に微調整する場合があります。共有サブモデルのトレーニングや推論では、異なるタスクの例を単一のベクトル化されたバッチに組み合わせて、アクセラレータの使用率を高める技術が役立ちます。

提案された PATHWAYS は、機能とパフォーマンスの点で SOTA ML システムに匹敵し、将来の ML ワークロードをサポートするために必要な機能を提供します。これは、PATHWAYS ランタイムがシステム管理コンピューティング アイランド上の多数のクライアントに代わってプログラムを実行できるようにするクライアント サーバー アーキテクチャを使用します。

PATHWAYS は、複数の TPU ポッドにわたってプログラムを透過的かつ効率的に実行するように設計された最初のシステムです。新しいデータフロー実行モデルを採用することで、数千のアクセラレータに拡張できます。 PATHWAYS プログラミング モデルを使用すると、SPMD 以外の計算を簡単に表現でき、集中型リソース管理と仮想化をサポートしてアクセラレータの使用率が向上します。

PATHWAYS システムアーキテクチャ

PATHWAYS は、TPU 計算を表現および実行するための XLA (TensorFlow、2019)、分散 CPU 計算を表現および実行するための TensorFlow グラフとエグゼキューター (Abadi 他、2016)、JAX (Bradbury 他、2016) や TensorFlow API などの Python プログラミング フレームワークなどの以前のシステムを基盤としています。これらのビルディング ブロックを使用することで、PATHWAYS は互換性を維持しながら、最小限のコード変更で既存の ML モデルを実行できます。

リソース マネージャー

PATHWAYS のバックエンドは、上の図 3 に示すように、DCN を介して相互接続された密結合アイランドにグループ化されたアクセラレータのセットで構成されています。 PATHWAYS には、島内のすべてのデバイスを集中管理する役割を担う「リソース マネージャー」があります。クライアントは、島の「仮想スライス」に、その通信パターンに適した特定の 2D または 3D メッシュ形状を持たせることを要求する場合があります。各仮想スライスには、クライアントがグリッド上で計算がどのようにレイアウトされるかを表現できる「仮想デバイス」が含まれています。リソース マネージャーは、必要な相互接続トポロジ、メモリ容量などを満たす仮想デバイスに物理デバイスを動的に割り当てます。

元のリソース マネージャーは、利用可能なすべてのデバイスに計算を分散し、仮想デバイスと物理デバイス間の 1 対 1 のマッピングを維持することで、負荷を静的に分散しようとする単純なヒューリスティックを使用して実装されました。将来のワークロードで必要になった場合は、より高度な割り当てアルゴリズムを採用できます。たとえば、すべてのクライアント計算のリソース要件とシステムの現在の状態を考慮して、物理デバイスの最適な割り当てを概算します。

PATHWAYS を使用すると、リソース マネージャーが利用可能なデバイスを追跡しながら、バックエンド コンピューティング リソースを動的に追加および削除できます。単一コントローラ設計によって実現される仮想デバイスと物理デバイス間の間接レイヤーにより、透過的なサスペンド/再開や移行などの機能を将来的にサポートできるようになります。これにより、クライアントの仮想デバイスは、ユーザー プログラムの支援なしに、リソースを一時的に再利用したり、再割り当てしたりできるようになります。

クライアント

ユーザーがトレースされたプログラムを実行する場合、PATHWAYS クライアント ライブラリを呼び出します。このライブラリは、最初に、これまで実行されていない計算に仮想デバイスを割り当て、計算をリソース マネージャーに登録して、サーバーがバックグラウンドで計算をコンパイルするようにトリガーします。

次に、クライアントは、カスタム MLIR (Lattner et al.、2021) 方言として表現される、デバイスの局所性に依存しないプログラムの PATHWAYS 中間表現 (IR) を構築します。 IR は一連の標準コンパイラ パスを通じて徐々にレベルが下げられ、最終出力には物理デバイスの場所の低レベル表現が含まれます。この低レベル プログラムは、物理デバイス間のネットワーク接続を考慮し、データ交換が必要な場合の分散操作と収集操作など、ソース コンピューティング シャードからの出力を宛先シャードの場所に転送する操作が含まれています。仮想デバイスの位置が変わらない通常の場合には、低レベルプログラムを再実行することが効果的です。リソースマネージャが仮想デバイスと物理デバイスのマッピング関係を変更した場合は、プログラムを再実行することができます。

古いシングル コントローラー システムのクライアントは、数千のアクセラレーターに分散された計算スライスの数千の個別の計算と対応するデータ バッファーを調整する役割を担っているため、すぐにパフォーマンスのボトルネックになる可能性があります。 PATHWAYS クライアントは、シャード バッファ抽象化を使用して、複数のデバイスに分散される可能性のある論理バッファを表します。この抽象化により、個々のシャードではなく論理バッファの粒度で簿記タスク (参照カウントを含む) のコストを償却することで、クライアントのスケーリングが可能になります。

協調的な実施

PATHWAYS は、DCN を使用してすべてのホスト間調整を実行するために PLAQUE に依存しています。 PLAQUE は、高ファンアウトまたは高ファンインの通信を必要とし、スケーラビリティとレイテンシが重要となる多くの顧客向けサービスで Google が使用する、既存の（クローズド ソースの）本番環境用シャード データフロー システムです。低レベルの PATHWAYS IR は PLAQUE プログラムに直接変換され、データフロー グラフとして表現されます。 PATHWAYS には調整基質に対する厳しい要件があり、PLAQUE はすべての要件を満たしています。

まず、PATHWAYS IR を記述するために使用される表現には、複数のシャードにわたる計算をコンパクトに表現できるように、各シャード計算に対して 1 つのノードが含まれている必要があります。つまり、N 個の計算シャードを持つ 2 つの計算 A と B の連鎖実行です。各計算シャードには、N がいくつあっても、データフロー表現で 4 つのノードがあります。Arg → Compute (A) → Compute (B) → Result。 PLAQUE ランタイム実装では、各ノードはターゲット シャードでタグ付けされた出力データ タプルを生成するため、データ並列実行を実行すると、隣接する IR ノードの各ペア間で N 個のデータ タプルが流れます。

コーディネーション ランタイムは、シャード エッジに沿ったスパース データ交換もサポートする必要があります。このデータ交換では、標準の進行状況追跡メカニズム (Akidau 他、2013 年、Murray 他、2013 年) を使用して、シャードのすべてのメッセージが受信されたことを検出し、動的に選択されたシャードのサブセット間でメッセージを送信できます。効率的なスパース通信により、DCN がアクセラレータ上のデータ依存制御フローのボトルネックになることが回避されます。これは、PATHWAYS によって実現される重要な機能の 1 つです。

下の図 4 に示すように、調整サブストレートは、送信スケジュール メッセージとデータ ハンドルのクリティカル パスで DCN メッセージを送信するために使用されるため、クリティカル メッセージを低遅延で送信し、高スループットが必要な場合は同じホストにメッセージをバッチで送信する必要があります。

スケーラブルな汎用データ フロー エンジンを使用して DCN 通信を処理すると、PATHWAYS が構成情報の配布、プログラムの監視、プログラムのクリーンアップ、障害発生時のエラーの報告などのバックグラウンド管理タスクにも使用できるため便利です。

Google は、低レベルの調整フレームワークを実装するために、PLAQUE の代わりに Ray (Moritz et al.、2018) などの別の分散フレームワークを使用して完全な PATHWAYS 設計を再実装することが可能であると考えています。この実装では、PATHWAYS エグゼキュータとスケジューラは、基盤となる Ray クラスター スケジューリングの上に PATHWAYS スケジューリングを実装する長時間実行 Ray アクターに置き換えられ、エグゼキュータは GPU 計算とコレクションに PyTorch を使用できます。

ギャングスケジュールのダイナミックスケジューリング

前述したように、共有アクセラレータのセットで SPMD 計算をサポートするための要件の 1 つは、効率的なギャング スケジューリングをサポートすることです。

PATHWAYS ランタイムには、各アイランドの集中型スケジューラが含まれており、アイランド上のすべての計算の一貫した順序付けを実行します。 PATHWAYS がプログラムを実行キューに入れると、PLAQUE データ フロー プログラムは次のアクションを実行します。

• バッファリングされた future を入力として、各アクセラレータでネイティブ コンパイルされた関数の実行をキューに入れます。
• 関数の実行出力のバッファリングされた未来を取得するには、リモート アクセラレータのキューにネットワーク送信を追加します。
• スケジューラと通信して、アイランドで実行されているすべてのプログラムにわたって関数の実行の一貫した順序を決定します。

スケジューラは、アクセラレータをミリ秒単位で割り当てる戦略を実装する必要があります。ただし、現在の実装では単純に FIFO 順に作業をキューに入れますが、より洗練されたスケジューラでは、推定実行時間に基づいて計算の順序を変更する可能性があります。

並列非同期スケジューリング

アクセラレータ上で計算を実行する場合、システムは非同期 API を活用して計算と調整を重ねることができます。下の図 4a の 3 つのノード図に示されているように、四角形は 3 つのノード A、B、C に対応しており、それぞれホスト A、B、C に接続されたアクセラレータ上で実行されます。すべてのノード計算は、通常のコンパイルされた関数です。ホスト A はノード A をキューに追加し、A からの将来の出力を受信して​​ホスト B に送信します。ホスト B はノード B の入力をソートし、入力バッファ アドレスをノード A に送信し、ノード B の機能を開始するための準備の大部分を実行します。ノード A が完了すると、その出力はアクセラレータ インターコネクトを介してノード B の入力バッファに直接送信され、その後ホスト B がノード B を起動します。 1 つのノードが完了してから別のノードが起動するまでの遅延は、データ転送時間よりも長くなります。

上記の設計は、先行ノードの計算時間が、スケジューリング、リソース分類、およびホスト間の調整にかかる時間を超える場合に適しています。ただし、計算時間が短すぎると、非同期パイプラインが停止し、ホスト側の作業が計算シーケンス全体の実行における主要なボトルネックになります。コンパイルされた関数が規則的である場合、後続のノードの入力形状は、先行する計算がキューに入れられる前に実際に計算できます。

そのため、Google は、下の図 4b に示すように、新しい並列非同期スケジューリング設計を導入しました。この方式では、先行処理がすでにキューに参加した後にノード作業をシリアル化するのではなく、通常のコンパイル済み関数の静的に既知のリソースを利用して、コンピューティング ノード上でホスト側の作業を並列に実行します。従来の機能では作業を並列にスケジュールすることしかできないことを考慮して、PATHWAYS は最適化手段として並列スケジューリングを使用し、先行計算が完了するまでノード リソース要件が不明な場合は従来のモデルに戻ります。

計算サブグラフが静的にスケジュール可能な場合、プログラムはサブグラフ全体を記述した単一のメッセージをスケジューラに送信し、スケジューラはサブグラフ内のすべてのアクティブなシャードの実行を連続して命令することができます。単一メッセージ設計はネットワーク トラフィックを最小限に抑えることを目的としていますが、スケジューラがすべてのサブグラフ シャードをバッチとしてキューに入れる必要はありません。計算は、他の同時実行プログラムによって送信された計算とインターリーブされる可能性があります。

3 ノード プログラムの順次スケジューリング (a) と並列スケジューリング (b) の比較。

実験結果

Google は、実際の機械学習モデル (SPMD プログラムとして表すことができる) のトレーニングにおける PATHWAYS のパフォーマンスを実証しました。まず、エンコーダー/デコーダー アーキテクチャを使用して Transformer モデルを実行する JAX マルチ コントローラーと比較します。

以下の表 1 は、異なる数のアクセラレータでトレーニングした場合の、さまざまなサイズのテキストからテキストへの Transformer モデルのトレーニング スループット (トークン/秒) を示しています。予想どおり、モデル コードが同一であるため、JAX と PATHWAYS でトレーニングされたモデルは同じステップ数で同じパープレキシティを達成します。

次に、Google は、デコーダー アーキテクチャのみを使用して Transformer 言語モデルをトレーニングする場合の、さまざまな構成での PATHWAYS のパフォーマンスを比較しました。表 2 に示すように、PATHWAYS のトレーニング スループットは、各パイプライン ステージの TPU コアの数に比例して増加しており、これは他のシステムと一致しています。

上記の結果は、以下の図 5 と一致しており、PATHWAYS のスループットはホストの数に応じて直線的に増加することを示しています。パイプラインのステージ数を増やすと最小オーバーヘッドが改善され、ステージ数が 4 から 16 に増えるとスループットは 133.7k トークン/秒から 131.4k トークン/秒に減少します。パイプライン モデルのパフォーマンスを SPMD を使用した同等のモデルと比較したところ、少なくともこのケースでは、SPMD 計算内の集約通信のオーバーヘッドがパイプライン バブルのオーバーヘッドよりも高いため、パイプラインのパフォーマンスは SPMD と同等であることが確認されました。

​​