Pathways on Cloud の概要

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 を使用して、バッチ、リアルタイム、インタラクティブなワークロードを実行する方法の概要について説明します。Google Kubernetes Engine での単一スライス TPU と Google Kubernetes Engine でのマルチスライス TPU の両方を含む、GKE での 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 ワーカーと通信します。

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 コンポーネントにマッピングします。

注: このセクションで説明するコンテナ仕様は例であり、デモのみを目的として使用されます。基盤となるコンテナは抽象化され、PathwaysJob カスタム リソース定義（CRD）API を介して管理されます。これらのコンテナの構成と 管理については、 PathwaysJob API をご覧ください。個々のコンテナ フラグ、環境変数、イメージの場所の操作は、customComponents 仕様でサポートされています。

Pathways コンテナ イメージは次の場所にあります。

コンテナタイプ	ロケーション
IFRT プロキシ サーバー	us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-<jax-version>
Pathways リソース マネージャー/ワーカー	us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-<jax-version>

Pathways リソース マネージャー

GKE クラスタを作成したら、次の containerSpec を使用して Pathways リソース マネージャーをデプロイできます。

  - 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: 一時ファイルに使用される Cloud Storage バケット。

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 に移植する
• クラウド上の Pathways のトラブルシューティング