Pathways は、大規模でマルチタスクのスパース活性化 ML システムを構築できるように設計されたシステムです。数千から数万のアクセラレータを使用でき、処理要件に基づいてさまざまなタスクにさまざまな量のコンピューティングを動的に割り当てることができます。
Pathways は、単一の JAX クライアントで大規模な TPU スライス（数千の TPU チップにまたがる可能性がある）にまたがるワークロードをオーケストレートすることで、大規模な ML 計算を簡素化します。
Google は、Gemini などの大規模モデルのトレーニングに Pathways を社内で使用しています。Pathways on Cloud は、 Google Cloud のお客様に同じメリットをもたらします。
始める前に
-
このページのサンプルをどのように使うかに応じて、タブを選択してください。
gcloud
Google Cloud CLI をインストールし、 フェデレーション ID を使用して gcloud CLI にログインします。ログイン後、次のコマンドを実行して Google Cloud CLI を初期化します。
gcloud initREST
このページの REST API サンプルをローカル開発環境で使用するには、gcloud CLI に指定した認証情報を使用します。
Google Cloud CLI をインストールし、 フェデレーション ID を使用して gcloud CLI にログインします。
詳細については、 Google Cloud 認証ドキュメントの REST を使用して認証するをご覧ください。
始める前に
Pathways は、Google Kubernetes Engine（GKE）または Slurm を使用する Cluster Director クラスタで実行できます。プラットフォームのタブを選択します。
GKE
次の前提条件を満たしていることを確認してください。
このドキュメントでは、GKE で Pathways マネージド TPU を使用して、バッチ、リアルタイム、インタラクティブなワークロードを実行する方法の概要について説明します。このドキュメントは、GKE での TPU（シングルスライスと GKE でのマルチスライス TPU の両方）の使用に精通しており、マルチスライス TPU の使用経験があることを前提としています。
Cluster Director
次の前提条件を満たしていることを確認してください。
このドキュメントでは、バッチ ワークロードとインタラクティブ ワークロードに Cluster Director クラスタで TPU と Pathways を使用する方法の概要について説明します。Cluster Director、Cluster Director の TPU マシン、Slurm でのジョブの実行に精通しており、マルチスライス TPU の使用経験があることを前提としています。
クラスタを作成するには、Pathways を使用して Cluster Director クラスタを作成するをご覧ください。
単一コントローラとマルチコントローラ
複数のデバイスにわたる計算を管理、オーケストレートするには、主に次の 2 つの方法があります。
| 機能 | 単一コントローラ（Pathways） | マルチコントローラ（JAX のデフォルト） |
|---|---|---|
管理 |
単一の制御ポイント: 単一のクライアント プログラムが中央コントローラとして機能します。 |
分散制御: 複数のプロセスが参加し、それぞれに独自の Python インタープリタ インスタンスがあります。 |
表示 |
統合ビュー: クライアントは、すべてのデバイスを 1 つの統合システムとして認識します。 |
ローカライズされたビュー: 各 Python プロセスは、接続されているデバイスのみを認識します。 |
プログラミング |
簡素化されたプログラミング: 単一のクライアントを操作することで、システムが多数のローカル アクセラレータを備えた単一の大規模なマシンとして表示されます。 |
単一プログラム、複数データ（SPMD）: 主に SPMD パラダイムを使用し、すべてのデバイスで同じプログラムを実行する必要があります。 |
柔軟性 |
非対称パイプライン並列処理や計算のスパース性など、SPMD を超える複雑な計算パターンをサポートします。 |
リソース管理の柔軟性が低くなる可能性があります。特に、異なる TPU スライス間ではその傾向が強くなります。 |
Pathways のコンポーネント
以降のセクションでは、Pathways アーキテクチャの主なコンポーネントについて説明します。
Pathways リソース マネージャー
このコンポーネントは、Pathways システムの中央コントロール プレーンです。すべてのアクセラレータ リソースを管理し、ジョブのアクセラレータの割り当てを調整します。ワーカーの健全性をモニタリングし、ジョブのスケジューリング、一時停止、再開を処理します。エラーとシステム ステータスの単一の連絡先として機能します。このコンポーネントに必要なのは CPU リソースのみです。
Pathways クライアント
このコンポーネントは、Pathways システムへのエントリ ポイントとして機能する Interim Framework Runtime（IFRT）の実装です。プログラムから高レベル オペレーション（HLO）を受け取ります。Pathways クライアントは、Pathways リソース マネージャーと連携して、コードに基づいてコンパイルされたプログラムを実行する場所を決定します。特定の JAX クライアントにシステムの統合ビューを提供します。このコンポーネントに必要なのは CPU リソースのみです。
Pathways ワーカー
Pathways ワーカー プロセスは、アクセラレータ マシン（TPU VM）で実行されます。IFRT プロキシ サーバーからプログラムのコンパイル済み実行可能ファイルを受け取り、TPU で計算を実行します。Pathways ワーカーは、IFRT プロキシ サーバーを介してプログラムにデータを送り返します。このコンポーネントにはアクセラレータ リソースが必要です。
IFRT プロキシ クライアント
IFRT プロキシ クライアントは、Interim Framework Runtime（IFRT） API のオープンソース実装です。この API は、ユーザーコードを基盤となるランタイムから切り離し、コードのポータビリティと透明性を高めます。JAX は、デフォルトのマルチコントローラ ランタイムの代わりにこの実装を使用します。IFRT プロキシ クライアントは、プログラムと Pathways コンポーネント間の通信ブリッジとして機能します。IFRT プロキシ サーバーにリクエストを送信し、結果を受信します。このコンポーネントには CPU リソースのみが必要です。
IFRT プロキシ サーバー
IFRT プロキシ サーバーは gRPC サーバーで、IFRT プロキシ クライアントからリクエストを受信し、Pathways クライアントに転送します。Pathways クライアントは、作業の分散を処理します。このコンポーネントに必要なのは CPU リソースのみです。
サイドカー サーバー
サイドカー サーバーは、アクセラレータ VM 上の Pathways ワーカーと同じ場所に配置された gRPC サーバーです。ユーザー指定の Python コードをアクセラレータ VM で直接実行して、コントローラとアクセラレータ間のデータ転送レイテンシを短縮します。サイドカー サーバーは、gRPC トランスポート上のカスタム バージョン管理プロトコルを介して Pathways ワーカーとやり取りします。
GKE の Pathways コンポーネント
このセクションでは、Pathways コンポーネントをコンテナや Pod などの GKE コンポーネントにマッピングします。
Pathways コンテナ イメージは次の場所にあります。
| コンテナのタイプ | ロケーション |
|---|---|
IFRT プロキシ サーバー |
|
Pathways リソース マネージャー / ワーカー |
|
Pathways リソース マネージャー
GKE クラスタを作成したら、次の containerSpec を使用して Pathways リソース マネージャーをデプロイできます。
- name: pathways-rm image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:latest imagePullPolicy: Always env: - name: HOST_ADDRESS valueFrom: fieldRef: fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']" - name: TPU_SKIP_MDS_QUERY value: "true" args: - --server_port=29001 - --node_type=resource_manager - --instance_count=WORKLOAD_NODEPOOL_COUNT - --instance_type=SLICE_TOPOLOGY - --gcs_scratch_location=gs://BUCKET_NAME
コンテナは次の引数を使用します。
--server_port: Pathways リソース マネージャーが他の Pathways コンポーネントとの通信に使用するポート。--node_type: ノードタイプ。Pathways リソース マネージャーの場合は、この引数をresource_managerに設定します。他のコンテナでは必要ありません。--instance_count: TPU スライスの数。--instance_type: スライスの TPU タイプとトポロジ。tpuTPU_TYPE:TOPOLOGY形式（例:tpuv5e:4x4）。--gcs_scratch_location: 一時ファイルに使用される Cloud Storage バケット。
IFRT プロキシ サーバー
次の containerSpec を使用して、IFRT プロキシ サーバーをデプロイできます。
- name: pathways-proxy image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:latest imagePullPolicy: Always env: - name: PATHWAYS_HEAD valueFrom: fieldRef: fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']" args: - --resource_manager_address=$(PATHWAYS_HEAD):29001 - --server_port=29000 - --gcs_scratch_location=gs://BUCKET_NAME ports: - containerPort: 29000
コンテナは次の引数を使用します。
--resource_manager_address: プロキシ サーバーが Pathways リソース マネージャーとの通信に使用するホスト名とポート。ポートは、Pathways リソース マネージャー コンテナに使用される--server_port値と一致する必要があります。--server_port: IFRT プロキシ サーバーが IFRT プロキシ クライアントとの通信に使用するポート。--gcs_scratch_location: 一時ファイルに使用される Cloud Storage バケット。
Pathways ワーカー
次の containerSpec を使用して、Pathways ワーカーをデプロイできます。
- name: worker image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:latest imagePullPolicy: Always env: - name: PATHWAYS_HEAD valueFrom: fieldRef: fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']" - name: MEGASCALE_NUM_SLICES valueFrom: fieldRef: fieldPath: "metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas']" - name: MEGASCALE_SLICE_ID valueFrom: fieldRef: fieldPath: "metadata.labels['jobset.sigs.k8s.io/job-index']" - name: MEGASCALE_COORDINATOR_ADDRESS value: "$(PATHWAYS_HEAD)" args: - --server_port=29001 - --resource_manager_address=$(PATHWAYS_HEAD):29001 - --gcs_scratch_location=gs://BUCKET_NAME ports: - containerPort: 29001 resources: limits: google.com/tpu: "4"
コンテナは次の引数を使用します。
--resource_manager_address: TPU ワーカーが Pathways リソース マネージャーとの通信に使用するホスト名とポート。ポートは、Pathways リソース マネージャー コンテナに使用される--server_port値と一致する必要があります。--server_port: ワーカーがプロキシ サーバーと Pathways リソース マネージャーとの通信に使用するポート。--gcs_scratch_location: 一時ファイルに使用される Cloud Storage バケット。
Pathways リソース マネージャー、IFRT プロキシ サーバー、Pathways ワーカーはすべて異なるポートを使用できますが、この例では、Pathways リソース マネージャーと Pathways ワーカーが同じポートを共有しています。
Cluster Director の Pathways コンポーネント
このセクションでは、Pathways コンポーネントを Cluster Director クラスタの Slurm ジョブにマッピングします。
Pathways スクリプトは、Enroot を使用して Pathways コンポーネントを実行します。setup_env.sh スクリプトは、ビルドする Python 環境の JAX バージョン（0.11.1）に一致する次のコンテナ イメージを使用します。
| コンテナのタイプ | ロケーション |
|---|---|
IFRT プロキシ サーバー |
|
Pathways リソース マネージャー / ワーカー |
|
Pathways コンポーネントを起動するには、start_pathways.sh スクリプトを Pathways サービスジョブとして送信します。次のコマンドを自分で作成する必要はありません。スクリプトは、start_pathways.sh の送信時に設定した環境変数（SLICE_SHAPE、MIN_SLICES、MAX_SLICES、GCS_BUCKET など）から引数を取得します。
Pathways リソース マネージャー
Pathways サービスジョブは、1 つの CPU ノードで Pathways リソース マネージャーを実行します。Pathways サーバー コンテナで次のコマンドを実行します。
/pathways/server \ --server_port=29001 \ --node_type=resource_manager \ --instance_type=tpuTPU_TYPE:TOPOLOGY \ --instance_count=MAX_SLICES \ --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \ --enable_gcs_olm=false
このコマンドは、次の引数を使用します。
--server_port: Pathways リソース マネージャーがリッスンするポート。IFRT プロキシ サーバーと Pathways ワーカーは、このポートに接続します。--node_type: ノードタイプresource_manager。この引数を設定するのは、Pathways リソース マネージャーのみです。--instance_type: 各スライスの TPU タイプとトポロジ（SLICE_SHAPEから）。たとえば、SLICE_SHAPE=2x4はtpuv6e:2x4を設定します。--instance_count: TPU スライスの数（MAX_SLICES）。--gcs_scratch_location: コンパイル キャッシュなどの一時ファイルに使用される、GCS_BUCKETからの Cloud Storage バケット内のパス。--enable_gcs_olm: Pathways が Cloud Storage オブジェクトのライフサイクル ルールを作成して、古いコンパイル キャッシュ エントリをクリーンアップするかどうか。このスクリプトでは、バケットに対する追加の権限が必要になるため、この機能はオフになっています。
IFRT プロキシ サーバー
Pathways サービスジョブは、Pathways リソース マネージャーと同じ CPU ノードで IFRT プロキシ サーバーを実行します。IFRT プロキシ サーバー コンテナで次のコマンドを実行します。
/pathways/proxy_server \ --server_port=29000 \ --resource_manager_address=RM_IP:29001 \ --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \ --num_elastic_instances=MAX_SLICES
このコマンドは、次の引数を使用します。
--server_port: IFRT プロキシ サーバーがリッスンするポート。JAX プログラムは、このポートに接続します。--resource_manager_address: Pathways リソース マネージャーのアドレス。RM_IPは、CPU ノードの内部 IP アドレスです。--gcs_scratch_location: Pathways リソース マネージャーと同じ Cloud Storage パス。--num_elastic_instances: IFRT プロキシ セッションを停止せずに使用不可にできるスライスの数（MAX_SLICES）。ワークロードはMIN_SLICESを適用します。
少なくとも MIN_SLICES スライスが実行されている場合、サービスジョブは IFRT プロキシ サーバーのアドレスを ~/pathways/cluster.env に書き込みます。JAX プログラム（IFRT プロキシ クライアント）は、このファイルを読み取って Pathways に接続します。JAX プログラムは、ワークロード（run_maxtext.sh など）として、または Jupyter ノートブックで実行されます。
Pathways ワーカー
Pathways サービスジョブは、TPU スライスごとに 1 つの配列タスクを含む Slurm ジョブ配列（run_worker_slice.sh）を送信します。各配列タスクは、スライス内のすべての TPU VM の Pathways サーバー コンテナで次のコマンドを実行します。
/pathways/server \ --server_port=29001 \ --resource_manager_address=RM_IP:29001 \ --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \ --enable_gcs_olm=false
このコマンドは、次の引数を使用します。
--server_port: Pathways ワーカーがリッスンするポート。Pathways リソース マネージャーと IFRT プロキシ サーバーは、このポートに接続します。--resource_manager_address: Pathways リソース マネージャーのアドレス。IFRT プロキシ サーバーと同じです。--gcs_scratch_location、--enable_gcs_olm: Pathways リソース マネージャーと同じ。
このスクリプトは、スライスのノードと配列タスク インデックスから、各ワーカーの TPU トポロジと MEGASCALE_* 環境変数も設定します。
TPU スライスがプリエンプトされるか失敗すると、Slurm は配列タスクを再キューイングします。クラスタノードにパスワードなしの sudo がある場合、サービスジョブはスライスのノードをパワーダウンして、Slurm が再作成できるようにします。Pathways リソース マネージャーと IFRT プロキシ サーバーは引き続き実行されるため、JAX プログラムを新しいアドレスに接続する必要はありません。