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

このドキュメントでは、TPU Pod で TensorFlow を使用して計算を行う方法を説明します。次の手順を行います。

TensorFlow ソフトウェアを使用して TPU Pod スライスを作成する
SSH を使用して TPU VM に接続する
サンプルスクリプトを作成して実行する

TPU VM は、Cloud TPU API を呼び出す権限をサービスアカウントに依存します。デフォルトでは、TPU VM はデフォルトの Compute Engine サービスアカウントを使用します。これには、必要なすべての Cloud TPU 権限が含まれています。独自のサービスアカウントを使用する場合は、TPU 閲覧者のロールをサービスアカウントに追加する必要があります。Google Cloud のロールの詳細については、ロールについてをご覧ください。TPU VM を作成するときに、--service-account フラグを使用して独自のサービスアカウントを指定できます。

環境の設定

Cloud Shell で次のコマンドを実行して、gcloud の最新バージョンを実行していることを確認します。
```
$ gcloud components update
```
gcloud をインストールする必要がある場合は、次のコマンドを使用します。
```
$ sudo apt install -y google-cloud-sdk
```

いくつかの環境変数を作成します。

$ export PROJECT_ID=project-id
$ export TPU_NAME=tpu-name
$ export ZONE=europe-west4-a
$ export RUNTIME_VERSION=tpu-vm-tf-2.18.0-pod-pjrt
$ export ACCELERATOR_TYPE=v3-32

TensorFlow ランタイムを使用して v3-32 TPU Pod スライスを作成する

$ gcloud compute tpus tpu-vm create ${TPU_NAME}} \
  --zone=${ZONE} \
  --accelerator-type=${ACCELERATOR_TYPE} \
  --version=${RUNTIME_VERSION}

コマンドフラグの説明

zone: Cloud TPU を作成するゾーン。
accelerator-type: アクセラレータタイプでは、作成する Cloud TPU のバージョンとサイズを指定します。TPU のバージョンごとにサポートされているアクセラレータタイプの詳細については、TPU のバージョンをご覧ください。
version: Cloud TPU ソフトウェアのバージョン。

SSH を使用して Cloud TPU VM に接続する

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
      --zone=${ZONE}

サンプルスクリプトを作成して実行する

次の環境変数を設定します。

(vm)$ export TPU_NAME=tpu-name
(vm)$ export TPU_LOAD_LIBRARY=0

現在のディレクトリに tpu-test.py という名前のファイルを作成し、次のスクリプトをコピーして貼り付けます。

import tensorflow as tf
print("Tensorflow version " + tf.__version__)

cluster_resolver = tf.distribute.cluster_resolver.TPUClusterResolver()
print('Running on TPU ', cluster_resolver.cluster_spec().as_dict()['worker'])

tf.config.experimental_connect_to_cluster(cluster_resolver)
tf.tpu.experimental.initialize_tpu_system(cluster_resolver)
strategy = tf.distribute.experimental.TPUStrategy(cluster_resolver)

@tf.function
def add_fn(x,y):
  z = x + y
  return z

x = tf.constant(1.)
y = tf.constant(1.)
z = strategy.run(add_fn, args=(x,y))
print(z)

次のコマンドでこのスクリプトを実行します。

(vm)$ python3 tpu-test.py

このスクリプトは、TPU Pod スライスの各 TensorCore で計算を実行します。出力は次のようになります。

PerReplica:{
  0: tf.Tensor(2.0, shape=(), dtype=float32),
  1: tf.Tensor(2.0, shape=(), dtype=float32),
  2: tf.Tensor(2.0, shape=(), dtype=float32),
  3: tf.Tensor(2.0, shape=(), dtype=float32),
  4: tf.Tensor(2.0, shape=(), dtype=float32),
  5: tf.Tensor(2.0, shape=(), dtype=float32),
  6: tf.Tensor(2.0, shape=(), dtype=float32),
  7: tf.Tensor(2.0, shape=(), dtype=float32),
  8: tf.Tensor(2.0, shape=(), dtype=float32),
  9: tf.Tensor(2.0, shape=(), dtype=float32),
  10: tf.Tensor(2.0, shape=(), dtype=float32),
  11: tf.Tensor(2.0, shape=(), dtype=float32),
  12: tf.Tensor(2.0, shape=(), dtype=float32),
  13: tf.Tensor(2.0, shape=(), dtype=float32),
  14: tf.Tensor(2.0, shape=(), dtype=float32),
  15: tf.Tensor(2.0, shape=(), dtype=float32),
  16: tf.Tensor(2.0, shape=(), dtype=float32),
  17: tf.Tensor(2.0, shape=(), dtype=float32),
  18: tf.Tensor(2.0, shape=(), dtype=float32),
  19: tf.Tensor(2.0, shape=(), dtype=float32),
  20: tf.Tensor(2.0, shape=(), dtype=float32),
  21: tf.Tensor(2.0, shape=(), dtype=float32),
  22: tf.Tensor(2.0, shape=(), dtype=float32),
  23: tf.Tensor(2.0, shape=(), dtype=float32),
  24: tf.Tensor(2.0, shape=(), dtype=float32),
  25: tf.Tensor(2.0, shape=(), dtype=float32),
  26: tf.Tensor(2.0, shape=(), dtype=float32),
  27: tf.Tensor(2.0, shape=(), dtype=float32),
  28: tf.Tensor(2.0, shape=(), dtype=float32),
  29: tf.Tensor(2.0, shape=(), dtype=float32),
  30: tf.Tensor(2.0, shape=(), dtype=float32),
  31: tf.Tensor(2.0, shape=(), dtype=float32)
}

クリーンアップ

TPU VM の使用を終了したら、次の手順に沿ってリソースをクリーンアップします。

Compute Engine から切断します。
```
(vm)$ exit
```

Cloud TPU を削除します。

$ gcloud compute tpus tpu-vm delete ${TPU_NAME} \
  --zone=${ZONE}

次のコマンドを実行して、リソースが削除されたことを確認します。TPU がリストに表示されないことを確認します。削除には数分かかることがあります。
```
$ gcloud compute tpus tpu-vm list \
  --zone=${ZONE}
```