Pathways は、大規模でマルチタスクのスパース活性化 ML システムを構築できるように設計されたシステムです。数千から数万のアクセラレータを使用でき、処理要件に基づいてさまざまなタスクにさまざまな量のコンピューティングを動的に割り当てることができます。
Pathways は、単一の JAX クライアントで複数の大規模な TPU スライスにまたがるワークロードをオーケストレートすることで、大規模な ML 計算を簡素化します。このワークロードは、数千個の TPU チップにまたがる可能性があります。
Pathways は、Gemini などの大規模モデルのトレーニング用に Google 社内で使用されています。Pathways on Cloud は、 Google Cloud のお客様にも同様のメリットをもたらします。
始める前に
インストールに必要なもの:
- インストールされている Kubernetes ツール
- gcloud CLI がインストールされている
- TPU API を有効にした
- Google Kubernetes Engine API を有効にした
このドキュメントでは、Google Kubernetes Engine（GKE）で Pathways マネージド TPU を使用してバッチ、リアルタイム、インタラクティブ ワークロードを実行する方法の概要について説明します。GKE での TPU の使用（Google Kubernetes Engine でのシングルスライス TPU とマルチスライス TPU の両方を含む）に精通していることと、マルチスライス TPU の使用経験があることを前提としています。
単一コントローラとマルチコントローラ
複数のデバイスにわたって計算を管理およびオーケストレートする方法は、主に次の 2 つがあります。
機能 |
シングル コントローラ（Pathways） |
マルチ コントローラ（JAX のデフォルト） |
管理 |
単一の制御ポイント: 単一のクライアント プログラムが中央コントローラとして機能します。 |
分散制御: 複数のプロセスが参加し、それぞれに独自の Python インタープリタ インスタンスがあります。 |
表示 |
統合ビュー: クライアントはすべてのデバイスを単一の統合システムとして認識します。 |
ローカライズされたビュー: 各 Python プロセスには、接続されているデバイスのみが表示されます。 |
プログラミング |
プログラミングの簡素化: ユーザーは単一のクライアントを操作するため、システムは多くのローカル アクセラレータを備えた単一の大型マシンとして表示されます。 |
SPMD: 主に SPMD パラダイムを使用します。すべてのデバイスで同じプログラムを実行する必要があります。 |
柔軟性 |
非対称パイプライン並列処理や計算のスパース性など、SPMD を超えるより複雑な計算パターンをサポートします。 |
リソース管理の柔軟性が低下する可能性があります。特に、異なる TPU スライス間では柔軟性が低下します。 |
Pathways コンポーネント
次のセクションでは、Pathways アーキテクチャの主なコンポーネントの概要について説明します。
Pathways リソース マネージャー
これは、Pathways システムの中央コントロール プレーンです。すべてのアクセラレータ リソースを管理し、ユーザージョブのアクセラレータの割り当てを調整します。ワーカーの健全性をモニタリングし、ジョブのスケジューリング、一時停止、再開を処理します。エラーとシステム ステータスの単一の連絡窓口として機能します。このコンポーネントに必要なのは CPU リソースのみです。
Pathways クライアント
これは、Pathways システムへのエントリ ポイントとして機能する Interim Framework Runtime（IFRT）の実装です。プログラムから High-Level Operations（HLO）を受け取ります。Pathways クライアントは、Pathways リソース マネージャーと連携して、ユーザーコードに基づいてコンパイルされたプログラムを実行する場所を決定します。特定の JAX クライアントにシステムの統合ビューを提供します。このコンポーネントに必要なのは CPU リソースのみです。
Pathways ワーカー
これらは、アクセラレータ マシン（TPU VM）で実行されるプロセスです。IFRT プロキシ サーバーからプログラムのコンパイル済み実行可能ファイルを受け取り、TPU で計算を実行します。Pathways ワーカーは、IFRT プロキシ サーバーを介してプログラムにデータを送り返します。このコンポーネントにはアクセラレータ リソースが必要です。
IFRT プロキシ クライアント
これは、ユーザーコードを基盤となるランタイムから切り離し、コードの移植性と透明性を高める Interim Framework Runtime（IFRT）API の OSS 実装です。JAX は、デフォルトのマルチコントローラ ランタイムの代替としてこの実装を使用します。IFRT プロキシ クライアントは、プログラムと Pathways コンポーネント間の通信ブリッジとして機能します。IFRT プロキシ サーバーにリクエストを送信し、結果を受け取ります。これは IFRT API の OSS 実装です。このコンポーネントに必要なのは CPU リソースのみです。
IFRT プロキシ サーバー
この gRPC サーバーは、IFRT プロキシ クライアントからリクエストを受け取り、Pathways クライアントに転送します。Pathways クライアントは、実際の作業の分散を処理します。このコンポーネントに必要なのは CPU リソースのみです。
サイドカー サーバー
この gRPC サーバーは、アクセラレータ VM 上の Pathways ワーカーと同じ場所に配置され、アクセラレータ VM 上でユーザー指定の Python コードを直接実行して、コントローラからアクセラレータへのデータ転送レイテンシを短縮します。サイドカー サーバーは、gRPC トランスポート上のカスタム バージョン付きプロトコルを介して Pathways ワーカーとやり取りします。
PathwaysJob API
PathwaysJob API は、ML トレーニングとバッチ推論ワークロードのデプロイに使用する OSS Kubernetes ネイティブ API です。PathwaysJob のコントローラは、JobSet API を活用して、すべての Pathways コンポーネントのライフサイクルと調整を管理します。このカスタム リソース定義（CRD）は、Pathways ワークロードを定義するためのハイレベル インターフェースを提供します。これにより、一般的なシナリオで個々の Pod 仕様を直接管理する必要がなくなります。すべてのパラメータとその意味の一覧については、GitHub の PathwaysJob API ドキュメントをご覧ください。
apiVersion: pathways-job.pathways.domain/v1 kind: PathwaysJob metadata: name: pathways-USER spec: maxRestarts: MAX_RESTARTS pathwaysVersion: jax-JAX_VERSION workers: - type: $(TPU_MACHINE_TYPE) topology: $(TOPOLOGY) numSlices: $(WORKLOAD_NODEPOOL_COUNT) maxSliceRestarts: # Optional customComponents: # This section is completely optional - componentType: proxy_server image: CUSTOM_PROXY_SERVER customFlags: - --flag_name_1=value_1 customEnv: - name: key_1 value: value_1 - componentType: pathways_server image: CUSTOM_PATHWAYS_SERVER customFlags: - --flag_name_1=value_1 customEnv: - name: key_1 value: value_1 - componentType: worker image: CUSTOM_WORKER customFlags: - --flag_name_1=value_1 customEnv: - name: key_1 value: value_1 - componentType: colocated_python_sidecar image: CUSTOM_SIDECAR_IMAGE customFlags: - --flag_name_1=value_1 customEnv: - name: key_1 value: value_1 pathwaysDir: "gs://BUCKET_NAME" # Pre-create this bucket. controller: deploymentMode: default # Default mode deploys pathways cpu resources (resource # manager and proxy server) on a dedicated CPU node, recommended for training elasticSlices: ELASTIC_SLICES template: spec: containers: - name: main image: python:3.11 command: - bash - -c - | pip install pathwaysutils python3 -c 'import pathwaysutils; import jax; pathwaysutils.initialize(); print(jax.devices())'
次の表に、PathwaysJob API の設定を示します。
| 属性 | 説明 |
|---|---|
apiVersion |
PathwaysJob カスタム リソース定義（CRD）の API バージョン（pathways-job.pathways.domain/v1）を指定します。 |
kind |
Kubernetes オブジェクトを PathwaysJob として識別します。 |
metadata.name |
Kubernetes の PathwaysJob オブジェクトの名前。通常は pathways- |
spec |
PathwaysJob の望ましい状態と構成を定義します。 |
spec.maxRestarts |
障害が発生した場合にシステムが PathwaysJob を自動的に再起動できる最大回数。 |
spec.pathwaysVersion |
（省略可）このジョブの Pathways 環境で使用する JAX フレームワークの目的のバージョンを指定します（例: jax-0.5.3）。 |
spec.workers |
PathwaysJob のワーカープールの構成を定義する配列。通常は TPU リソースを使用します。 |
spec.workers[].type |
ワーカーノードに使用する TPU マシンのタイプ（たとえば、$TPU_MACHINE_TYPE は ct6e-standard-4t になります）。 |
spec.workers[].topology |
ワーカーに割り当てられた TPU スライスのトポロジ（たとえば、$TOPOLOGY は 2x2、4x4、2x2x2 など）。 |
spec.workers[].numSlices |
ワーカープールにプロビジョニングする TPU スライスの数（たとえば、$WORKLOAD_NODEPOOL_COUNT は 2 になります）。 |
spec.workers[].maxSliceRestarts |
（省略可）スライス内の個々のワーカーが失敗した場合に再起動できる最大回数。 |
spec.customComponents |
（省略可）メインジョブとともにカスタム コンポーネント（プロキシ サーバー、Pathways サーバー、追加のワーカーなど）を定義してデプロイできる配列。 |
spec.customComponents[].componentType |
定義するカスタム コンポーネントのタイプ（proxy_server、pathways_server、worker、colocated_python_sidecar など）を指定します。 |
spec.customComponents[].image |
このカスタム コンポーネントのコンテナに使用する Docker イメージ。 |
spec.customComponents[].customFlags |
コンテナの起動時に渡されるカスタム コマンドライン フラグの配列。 |
spec.customComponents[].customEnv |
コンテナ内で設定するカスタム環境変数の配列。各要素には名前と値があります。 |
spec.pathwaysDir |
PathwaysJob がコンパイル アーティファクトやその他の一時データの保存に使用する Cloud Storage バケット。このバケットは、ワークロードを実行する前に作成する必要があります。 |
spec.controller |
ジョブの実行全体を管理する Pathways コントローラ構成設定。 |
spec.controller.deploymentMode |
Pathways コントローラの CPU リソース（Pathways リソース マネージャーとプロキシ サーバー）のデプロイ方法を指定します。デフォルト モードでは、専用の CPU ノードにデプロイされます。colocate_head_with_workers では、TPU ワーカーとともにデプロイされます。 |
spec.controller.elasticSlices |
（省略可）ジョブの実行中に使用できなくなる TPU スライスの最大数。この数を超えると、異常と見なされます。 |
spec.controller.template |
（省略可）ユーザー ジョブの Pod テンプレートを定義します。これはバッチ ワークロードには必要ですが、インタラクティブ ワークロードには必要ありません。 |
spec.controller.template.spec |
ユーザー ジョブの Pod の仕様。 |
spec.controller.template.spec.containers |
ユーザー ジョブ内で実行されるコンテナを定義する配列 |
spec.controller.template.spec.containers[].name |
ユーザー ジョブ内のコンテナの名前（このサンプルでは main）。 |
spec.controller.template.spec.containers[].image |
メイン コンテナのコンテナに使用する Docker イメージ（このサンプルでは python:3.11）。 |
spec.controller.template.spec.containers[].command |
メインコンテナの起動時に実行されるコマンド。このサンプルでは、`pathwaysutils` をインストールし、Pathways を初期化して、JAX デバイスを出力します。 |
GKE の Pathways コンポーネント
このセクションでは、Pathways コンポーネントをコンテナや Pod などの Google Kubernetes Engine コンポーネントにマッピングします。
Pathways コンテナ イメージは次の場所にあります。
コンテナタイプ |
ロケーション |
IFRT プロキシ サーバー |
|
Pathways リソース マネージャー/ワーカー |
|
Pathways リソース マネージャー
GKE クラスタを作成したら、次の containerSpec を使用してパスウェイ リソース マネージャーをデプロイできます。
- name: pathways-rm image: us-docker.pkg.dev/cloud-tpu-v2-images/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 タイプとトポロジ。tpu{TPU type}:{TPU topology}の形式（例:tpuv5e:4x4）。--gcs_scratch_location: 一時ファイルに使用される Cloud Storage バケット。
IFRT プロキシ サーバー
次の containerSpec を使用して、IFRT プロキシ サーバーをデプロイできます。
- name: pathways-proxy image: us-docker.pkg.dev/cloud-tpu-v2-images/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: 一時ファイルに使用される GCS バケット。
Pathways ワーカー
次の containerSpec を使用して、Pathways ワーカーをデプロイできます。
- name: worker image: us-docker.pkg.dev/cloud-tpu-v2-images/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 ワーカーが同じポートを共有しています。
次のステップ
- Pathways を使用して GKE クラスタを作成する
- Pathways でバッチ ワークロードを実行する
- Pathways を使用してマルチホスト推論を実行する
- Pathways でインタラクティブ ワークロードを実行する
- Pathways を使用した復元力トレーニング
- JAX ワークロードを Pathways に移植する
- Cloud 上の Pathways のトラブルシューティング