マルチホスト TPU スライスを作成する

インスタンス テンプレートの作成

TPU v6e VM のインスタンス テンプレートを作成するには、gcloud compute instance-templates create コマンドを使用します。

gcloud compute instance-templates create quickstart-tpu-instance-template \
    --machine-type=ct6e-standard-4t \
    --maintenance-policy=TERMINATE \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --image-project=ubuntu-os-accelerator-images \
    --region=us-east5

ワークロード ポリシーの作成

ワークロード ポリシーは、コンピューティング インスタンスの物理プロパティを定義します。TPU スライスでは、アクセラレータ トポロジによって TPU チップの物理的な配置が定義されます。相互接続されたマルチホスト TPU スライスでは、アクセラレータ トポロジを指定する必要があります。

マルチホスト TPU スライスのワークロード ポリシーを作成するには、--accelerator-topology フラグを指定して gcloud compute resource-policies create workload-policy コマンドを使用します。次のコマンドは、2x4 トポロジのワークロード ポリシーを作成します。

gcloud compute resource-policies create workload-policy quickstart-tpu-workload-policy \
    --type=high-throughput \
    --accelerator-topology=2x4 \
    --region=us-east5

MIG を作成する

次のコマンドを実行して、マルチホスト TPU スライスを形成する MIG を作成します。

1. マルチホスト TPU スライスを形成する MIG を作成するには、gcloud compute instance-groups managed create コマンドを使用します。

   gcloud compute instance-groups managed create quickstart-tpu-mig \
       --size=2 \
       --target-size-policy-mode=bulk \
       --template=quickstart-tpu-instance-template \
       --region=us-east5 \
       --target-distribution-shape=any-single-zone \
       --instance-redistribution-type=none \
       --default-action-on-vm-failure=do-nothing \
       --workload-policy=projects/PROJECT_ID/regions/us-east5/resourcePolicies/quickstart-tpu-workload-policy
   

   PROJECT_ID は、実際の Cloud de Confiance by S3NS プロジェクト ID に置き換えます。

2. 必要に応じて、次のコマンドを使用して、マネージド インスタンスが実行されていることを確認します。

   • MIG の全体的なステータスを表示するには、gcloud compute instance-groups managed describe コマンドを使用します。

     gcloud compute instance-groups managed describe quickstart-tpu-mig \
         --region=us-east5
     

   • 各インスタンスのステータスを表示するには、gcloud compute instance-groups managed list-instances コマンドを使用します。

     gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
         --region=us-east5
     

JAX をインストールする

MIG 内のすべての TPU VM インスタンスの仮想環境に、依存関係と JAX フレームワークをインストールします。TPU VM に 3.11 より前のバージョンの Python がインストールされている場合は、最新バージョンの JAX を実行するために Python 3.11 をインストールする必要があります。

1. TPU VM で実行されている Python のバージョンを確認します。

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='python3 --version'
   

   バージョンが Python 3.11 より前の場合は、Python 3.11 をインストールします。

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt update && \
       sudo apt install -y software-properties-common && \
       sudo add-apt-repository -y ppa:deadsnakes/ppa && \
       sudo apt update && \
       sudo apt install -y python3.11 python3.11-dev'
   

2. 仮想環境を作成します。

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='sudo apt install -y python3.11-venv && \
       python3.11 -m venv ~/jax_venv'
   

3. 仮想環境に JAX をインストールします。

   gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
       --region=us-east5 \
       --uri \
   | xargs -I {} -P 0 gcloud compute ssh {} \
       --command='source ~/jax_venv/bin/activate && \
       pip install --upgrade pip -q && \
       pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html -q'
   

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

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

コードを準備する

各インスタンスに example.py というファイルを作成します。

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command="cat << 'EOF' > ~/example.py
import jax

# Initialize the slice
jax.distributed.initialize()

# The total number of TPU cores in the slice
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 slice
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)
EOF"

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

スライスの各 TPU VM で example.py プログラムを実行します。

gcloud compute instance-groups managed list-instances quickstart-tpu-mig \
    --region=us-east5 \
    --uri \
| xargs -I {} -P 0 gcloud compute ssh {} \
    --command='source ~/jax_venv/bin/activate && python3 ~/example.py'

出力例を以下に示します。

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