Pathways on Cloud の概要

Pathways は、大規模でマルチタスクのスパース活性化 ML システムを構築できるように設計されたシステムです。これにより、数千から数万のアクセラレータを使用でき、処理要件に基づいてさまざまなタスクにさまざまな量のコンピューティングを動的に割り当てることができます。

Pathways は、単一の JAX クライアントで複数の大規模な TPU スライスにまたがるワークロードをオーケストレートすることで、大規模な ML 計算を簡素化します。このワークロードは、数千個の TPU チップにまたがる可能性があります。

Pathways は、Gemini などの大規模モデルのトレーニング用に Google 社内で使用されています。Pathways on Cloud は、 Google Cloud のお客様にも同様のメリットをもたらします。

始める前に

インストールに必要なもの:

このドキュメントでは、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 クライアントに転送します。このコンポーネントに必要なのは CPU リソースのみです。

サイドカー サーバー

この gRPC サーバーは、アクセラレータ VM 上の Pathways ワーカーと同じ場所に配置され、アクセラレータ VM 上でユーザー指定の Python コードを直接実行して、コントローラからアクセラレータへのデータ転送レイテンシを短縮します。サイドカー サーバーは、gRPC トランスポート上のカスタム バージョン管理プロトコルを介して Pathways ワーカーとやり取りします。

Pathways コンポーネントの関係を示します。
Pathways コンポーネント

GKE の Pathways コンポーネント

このセクションでは、Pathways コンポーネントをコンテナや Pod などの Google Kubernetes Engine コンポーネントにマッピングします。

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 を使用してパスウェイ リソース マネージャーをデプロイできます。

  - 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 ワーカーが同じポートを共有しています。

次のステップ