TPU Pod スライスでの JAX コードの実行

JAX コードを単一の TPU ボードで実行したら、TPU Pod スライスで実行してコードをスケールアップできます。TPU Pod スライスは、専用の高速ネットワーク接続で互いに接続された複数の TPU ボードです。このドキュメントでは、TPU Pod スライスでの JAX コードの実行についての概要を説明します。詳しくは、マルチホスト環境とマルチプロセス環境での JAX の使用をご覧ください。

データストレージにマウントされた NFS を使用する場合は、Pod スライス内のすべての TPU VM に OS Login を設定する必要があります。詳細については、データストレージに NFS を使用するをご覧ください。

TPU Pod スライスの作成

このドキュメントのコマンドを実行する前に、アカウントと Cloud TPU プロジェクトを設定するの手順に従ってください。ローカルマシンで次のコマンドを実行します。

gcloud コマンドを使用して TPU Pod スライスを作成します。たとえば、v4-32 Pod スライスを作成するには、次のコマンドを使用します。

$ gcloud compute tpus tpu-vm create tpu-name  \
  --zone=us-central2-b \
  --accelerator-type=v4-32  \
  --version=tpu-ubuntu2204-base

Pod スライスに JAX をインストールする

TPU Pod スライスを作成したら、TPU Pod スライスのすべてのホストに JAX をインストールする必要があります。--worker=all オプションを使用すると、1 つのコマンドですべてのホストに JAX をインストールできます。

  gcloud compute tpus tpu-vm ssh tpu-name \
  --zone=us-central2-b --worker=all --command="pip install \
  --upgrade 'jax[tpu]>0.3.0' \
  -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"

Pod スライスで JAX コードを実行する

TPU Pod スライスで JAX コードを実行するには、TPU Pod スライスの各ホストでコードを実行する必要があります。jax.device_count() 呼び出しは、Pod スライスの各ホストで呼び出されるまで応答しなくなります。次の例は、TPU Pod スライスで単純な JAX 計算を実行する方法を示しています。

コードの準備

scp コマンドでは、gcloud バージョン 344.0.0 以降が必要です。 gcloud --version を使用して gcloud のバージョンを確認し、必要に応じて gcloud components upgrade を実行します。

次のコードを使用して、example.py という名前のファイルを作成します。

# The following code snippet will be run on all TPU hosts
import jax

# The total number of TPU cores in the Pod
device_count = jax.device_count()

# The number of TPU cores attached to this host
local_device_count = jax.local_device_count()

# The psum is performed over all mapped devices across the Pod
xs = jax.numpy.ones(jax.local_device_count())
r = jax.pmap(lambda x: jax.lax.psum(x, 'i'), axis_name='i')(xs)

# Print from a single host to avoid duplicated output
if jax.process_index() == 0:
    print('global device count:', jax.device_count())
    print('local device count:', jax.local_device_count())
    print('pmap result:', r)

`example.py` を Pod スライス内のすべての TPU ワーカー VM にコピーする

$ gcloud compute tpus tpu-vm scp example.py tpu-name: \
  --worker=all \
  --zone=us-central2-b

以前に scp コマンドを使用したことがない場合、次のようなエラーが表示されることがあります。

ERROR: (gcloud.alpha.compute.tpus.tpu-vm.scp) SSH Key is not present in the SSH
agent. Please run `ssh-add /.../.ssh/google_compute_engine` to add it, and try
again.

このエラーを解決するには、エラーメッセージに示されている ssh-add コマンドを実行してから、コマンドを再実行します。

Pod スライスでコードを実行する

すべての VM で example.py プログラムを起動します。

$ gcloud compute tpus tpu-vm ssh tpu-name \
  --zone=us-central2-b \
  --worker=all \
  --command="python3 example.py"

出力（v4-32 Pod スライスで生成）

global device count: 16
local device count: 4
pmap result: [16. 16. 16. 16.]

クリーンアップ

終了した後は、gcloud コマンドを使用して TPU VM リソースを解放できます。

$ gcloud compute tpus tpu-vm delete tpu-name \
  --zone=us-central2-b