JAX ワークロードを Pathways に移植する

Pathways を使用した JAX の分散型の性質により、通信オーバーヘッドが原因で一部のオペレーションが適切にスケーリングされない場合があります。Pathways は、非同期ディスパッチなどの機能を使用してこれらのオーバーヘッドを最小限に抑えますが、JAX ワークロードを Pathways に移植する場合や、Pathways ワークロードで JAX を多数のアクセラレータにスケーリングする場合は、注意が必要な点がいくつかあります。

始める前に

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

• Cluster Toolkit を設定する
• インストール済みの Kubernetes ツール
• gcloud CLI がインストールされていること
• Google Kubernetes Engine API を有効にする

プロセス インデックス

Pathways を使用する JAX は、Pathways クラスタ内のすべてのデバイスをローカルとして扱います。これにより、デバイス管理が簡素化され、JAX は利用可能なすべてのリソースを使用できます。実際には、これは次のことを意味します。

• jax.process_index() は、すべてのデバイスで常に 0 です。
• jax.devices() と jax.local_devices() は、ジョブ全体のすべての TPU デバイスを返します。

ハードウェアの種類とコロケーション

パフォーマンスを最大限に高めるには、すべての Pathways コンポーネントとユーザー ジョブを同じ Google Cloud クラウドゾーンに配置します。IFRT プロキシやリソース マネージャーなどの大規模な CPU を使用します。64 個の vCPU と 256 GB のメモリを備えた専用の n2-standard-64 をおすすめします。

PathwaysUtils

Pathways-utils は、Python ベースの GitHub リポジトリです。Pathways on Cloud アーキテクチャで JAX ワークロードのデプロイと実行を効率化するために必要なユーティリティとツールを提供します。このパッケージは、クラウド環境に必要な適応を処理するため、JAX デベロッパーはプラットフォーム固有の構成を最小限に抑え、コアの ML ワークフローに集中できます。具体的には、次のようなメリットがあります。

• 「プロキシ」JAX バックエンド: このカスタム バックエンドを使用すると、JAX_PLATFORMS=proxy 環境変数を設定することで、JAX アプリケーションで Pathways インフラストラクチャを使用できます。
• 統合プロファイリング ユーティリティ: アプリケーションのパフォーマンスを把握できるプロファイリング機能。jax.profiler.start_trace や jax.profiler.start_server などの標準の JAX プロファイリング API を使用すると、JAX コードだけでなく、基盤となる Pathways コンポーネントもプロファイリングできます。これにより、クラウド環境内の実行を全体的に把握できます。
• Orbax を使用した分散チェックポインティング: Pathways 環境内で Orbax ライブラリを使用するときに、分散チェックポイントを使用し、チェックポイントを復元できるカスタム Orbax チェックポイント ハンドラ。この統合は、pathwaysutils をインポートする限り、既存の Orbax チェックポイント コードを変更することなく動作することを目的としています。
• Elastic Training Primitives: Pathways を使用して堅牢でスケーラブルなトレーニング ワークフローを構築するために使用できる、基本的な Elastic Training Primitives を提供します。これらのプリミティブを使用すると、トレーニング ジョブが利用可能なリソースの変更に動的に適応し、クラウド環境の効率とレジリエンスを向上させることができます。

チェックポイント処理

Orbax は、Cloud Storage を使用した分散チェックポイントと復元のために Pathways で徹底的にテストされています。環境変数 ENABLE_PATHWAYS_PERSISTENCE=1 を設定し、train.py で import pathwaysutils; pathwaysutils.initialize() を呼び出すと、IFRT プロキシを介してチェックポイント オペレーションを効率的に処理するカスタム ArrayHandler が登録されます。これにより、アクセラレータ上の Pathways ワーカーはデータを直接保存して復元できます。

コロケーション Python

コロケーション Python は、ユーザー指定の Python コードを TPU または GPU ホストで直接実行できるオープンソースの JAX API です。これは、マルチコントローラ JAX でより簡単です。このアプローチを使用すると、データの読み込みやチェックポイントなどのコンピューティング負荷の高いタスクで、クライアントと TPU マシン間のデータ転送を回避できます。コロケーションされた Python JAX API を実行するように Pathways クラスタを構成するには、サンプルのサイドカー Dockerfile を使用してサイドカー コンテナ イメージをビルドして push します。

データ読み込み

トレーニング中は、データセットからバッチを繰り返し読み込み、モデルにフィードします。アクセラレータ使用率の低下を回避するには、バッチをホスト間でシャーディングする効率的な非同期データローダーが必要です。Pathways でトレーニングを実行すると、データローダーは CPU VM（マルチコントローラ設定で使用される TPU VM とは異なります）で実行され、TPU VM にデータをディスパッチします。これにより、データの読み取りレイテンシが高くなりますが、CPU ホストで X 個のバッチを事前に読み取り、読み取ったデータを TPU に非同期でディスパッチすることで、レイテンシを部分的に軽減できます。このソリューションは、小規模から中規模のスケールで実行する場合に十分です。

GKE のヘッド ワークロード コンテナでマルチプロセス データローダー（Grain など）を実行すると、コンテナは pathways-proxy サイドカーと pathways-rm サイドカーとヘッドノードを共有します。この共有アーキテクチャのため、次の構成を行う必要があります。

• ワーカー数を明示的に設定する: コンテナ内では、os.cpu_count() はコンテナの CPU 上限ではなく、ホスト VM の合計 vCPU を報告します。os.cpu_count() から取得するのではなく、明示的なワーカー数（コンテナの CPU 上限以下）を渡します。
• メモリバック /dev/shm ボリュームをマウントする: マルチプロセス データローダーは、プロセス間通信に POSIX 共有メモリ（/dev/shm）を使用します。プリフェッチされたバッチがデフォルトのコンテナ共有メモリの上限を超えないように、メモリバックの emptyDir（medium: Memory）ボリュームが /dev/shm にマウントされていることを確認します。

大規模なパフォーマンスを最適化するには、コロケーション Python を使用して、アクセラレータでデータ パイプラインを直接実行し、入力データ パイプラインをコロケーションすることをおすすめします。これにより、CPU のボトルネックが解消され、TPU の高速相互接続をデータ転送に活用できます。

TFDS ベースの入力パイプラインを移行するリファレンス実装は、multihost_dataloading.py の RemoteIterator 実装にあります。この実装は、コロケーションされた Python JAX API を使用して、マルチコントローラ JAX と Pathways の両方で分散方式で動作します。

Jax のバージョニング

Pathways のリリースは、互換性と安定性を確保するために JAX バージョンと密結合されています。潜在的な問題を回避するには、Pathways アーティファクトと JAX バージョンが一致していることを確認します。Pathways の各リリースでは、jax-<version> 形式のタグを使用して、互換性のある JAX バージョンを明示的に指定します。

コンパイル キャッシュ

Pathways 永続コンパイル キャッシュは、Pathways サーバーがコンパイル済みの XLA 実行可能ファイルを Cloud Storage などの永続的な場所に保存して、冗長なコンパイルを回避できるようにする機能です。この機能はデフォルトで有効になっています。キャッシュの場所は、--gcs_scratch_location フラグとしてリソース マネージャーと Pathways ワーカー コンテナに渡されます。関連するストレージ費用を最小限に抑えるため、キャッシュは Cloud Storage のロケーションにライフサイクル ポリシーを適用します。Cloud Storage バケットあたりのポリシー数の上限は 50 個です。そのため、すべてのワークロードで共通の Cloud Storage ロケーションを使用することをおすすめします。

--gcs_scratch_location フラグを省略するか、空のままにすると、Pathways は自動的にメモリ内コンパイル キャッシュにフォールバックします。このキャッシュは、Pathways デプロイ間でコンパイル済みの XLA 実行可能ファイルを保持しません。

このキャッシュは、Pathways ワークロードで pathwaysutils.initialize() によって無効にされる JAX コンパイル キャッシュに似ています。

コンパイル キャッシュには、次の Cloud Storage 権限が必要です。

• storage.buckets.get: バケット メタデータを取得します。
• storage.buckets.update: Pathways がオブジェクト ライフサイクル ポリシーを設定して、キャッシュの強制排除に TTL を適用するために必要です。
• storage.objects.list: バケット内の既存のキャッシュ オブジェクトを一覧表示します。
• storage.objects.create: 新しいコンパイル済み実行可能ファイルをキャッシュに書き込みます。
• storage.objects.get: バケットからキャッシュに保存された実行可能ファイルを読み取ります。

プロファイリング

JAX プロファイラを使用すると、JAX プログラムのトレースを生成できます。Pathways でサポートされている一般的な方法は 2 つあります。

• プログラマティック
  • JAX コードからプログラムでプロファイルをキャプチャする
• 手動
  • JAX コードからプロファイラ サーバーを起動した後にオンデマンドでプロファイルをキャプチャする

どちらの場合も、プロファイルは Cloud Storage バケットに書き込まれます。Cloud Storage バケットに複数のトレース ファイルが作成されます。これらのファイルは、異なるタイムスタンプ フォルダに作成される可能性があります。次に例を示します。

• トレースを呼び出したメインの Python プロセス（通常はノートブック VM）: <jax-client-vm-name>.xplane.pb
• Pathways IFRT プロキシ: client.<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
• Pathways Resource Manager: server.*<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
• Pathways ワーカー: server.*<tpu-node-name>.xplane.pb

これらのトレース ファイルは、次のコマンドを実行して TensorBoard で分析できます。TensorBoard とそのすべてのプロファイリング ツールの詳細については、プロファイラを使用して TensorFlow のパフォーマンスを最適化するをご覧ください。

# verify trace files are present
gsutil ls -l -r gs://BUCKET/PREFIX

# View on tensorboard
tensorboard --logdir=gs://BUCKET/PREFIX

次のように置き換えます。

• BUCKET : トレース ファイルを保存する Cloud Storage バケット
• PREFIX: トレース ファイルを保存する Cloud Storage バケット内のパス

プログラムによるプロファイルのキャプチャ

コード内からプロファイルをキャプチャします。プロファイルは、gs://<bucket>/<prefix> 内のタイムスタンプ ディレクトリに保存されます。

import jax
import pathwaysutils

pathwaysutils.initialize()

jax.profiler.start_trace("gs://BUCKET/PREFIX")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

手動プロファイル キャプチャ

プロファイル情報を手動でキャプチャするには、Python コードからプロファイラ サーバーを起動する必要があります。

import jax
import pathwaysutils

pathwaysutils.initialize()
jax.profiler.start_server(jax_profiler_port)

# Your JAX code
jax.profiler.stop_server() # this is functionally a no-op

プロファイラ サーバーの実行中に、プロファイルをキャプチャして、ターゲットの Cloud Storage の場所にデータをエクスポートできます。

export DURATION_IN_SECS=6
export OUTPUT_DIR=gs://<var>BUCKET</var>/<var>PREFIX</var>

curl -d "{\"duration_ms\":\"${DURATION_IN_SECS} * 1000 }}\", \"repository_path\":\"${OUTPUT_DIR}\"}" -H "Content-Type: application/json" -X POST http://localhost:<jax_profiler_port>

Compile や Execute などの IFRT プロキシ クライアント メソッドのタイミング情報は、プログラムのトレース内で確認できます。これらのイベントは、コンパイルと実行中に IFRT プロキシ gRPC サーバーとのやり取りを詳細に記録したもので、GrpcClientSessionUserFuturesWorkQueue という名前のスレッドに表示されます。トレースでこのスレッドを調べることで、これらのオペレーションのパフォーマンスに関する分析情報を得ることができます。

XLA フラグ

Pathways を使用する場合は、pathways-proxy コンテナに XLA フラグを設定する必要があります。

Cluster Toolkit（gcluster job submit コマンド）を使用する場合は、--pathways-proxy-args を使用して XLA フラグを設定します。

--pathways-proxy-args="--xla_flag_1=value1 --xla_flag_2=value2"

HLO ダンプ

XLA コンパイラに渡される High Level Optimizer（HLO）入力を検査するには、--pathways-proxy-args を使用して、指定された Cloud Storage の場所に HLO をダンプするように Pathways を構成します。

--pathways-proxy-args="--xla_dump_to=gs://BUCKET_NAME/HLO_PREFIX/"

次のステップ

• Pathways を使用して GKE クラスタを作成する
• Pathways を使用したマルチホスト推論
• Pathways を使用してバッチ ワークロードを実行する
• Pathways のインタラクティブ モード
• Pathways を使用した復元力のあるトレーニング
• Pathways on Cloud のトラブルシューティング