Cloud TPU v5e トレーニング

Cloud TPU v5e は、Google Cloud の最新の生成 AI アクセラレータです。Pod あたり 256 チップのフットプリントが小さい v5e は、トランスフォーマー、text-to-image、および、畳み込みニューラルネットワーク（CNN）のトレーニング、微調整、サービス提供に適した最適なプロダクトとなるように最適化されています。Cloud TPU v5e をサービングに使用する方法については、v5e を使用した推論をご覧ください。

Cloud TPU v5e TPU のハードウェアと構成の詳細については、TPU v5e をご覧ください。

使ってみる

以降のセクションでは、TPU v5e の使用を開始する方法について説明します。

リクエストの割り当て

トレーニングに TPU v5e を使用するには、割り当てが必要です。オンデマンド TPU、予約 TPU、TPU Spot VM には、さまざまな割り当てタイプがあります。推論に TPU v5e を使用する場合は、個別の割り当てが必要です。割り当ての詳細については、割り当てをご覧ください。TPU v5e の割り当てをリクエストするには、Cloud セールスにお問い合わせください。

Google Cloud アカウントとプロジェクトを作成する

Cloud TPU を使用するには、Google Cloud アカウントとプロジェクトが必要です。詳細については、Cloud TPU 環境を設定するをご覧ください。

Cloud TPU の作成

queued-resource create コマンドを使用して、Cloud TPU v5es をキューに格納されたリソースとしてプロビジョニングすることをおすすめします。詳細については、キューに入れられたリソースを管理するをご覧ください。

Create Node API（gcloud compute tpus tpu-vm create）を使用して Cloud TPU v5es をプロビジョニングすることもできます。詳細については、TPU リソースの管理をご覧ください。

トレーニングに使用できる v5e 構成の詳細については、トレーニング用の Cloud TPU v5e タイプをご覧ください。

フレームワークの設定

このセクションでは、TPU v5e で JAX または PyTorch を使用したカスタムモデルのトレーニングの一般的な設定プロセスについて説明します。TensorFlow のサポートは、tpu-vm-tf-2.17.0-pjrt と tpu-vm-tf-2.17.0-pod-pjrt の TPU ランタイムバージョンで利用できます。

推論の設定手順については、v5e 推論の概要をご覧ください。

JAX を設定する

スライス形状が 8 チップを超える場合、1 つのスライスに複数の VM があります。この場合、SSH を使用して個別にログインすることなく、--worker=all フラグを使用してすべての TPU VM に 1 つのステップでインストールを実行する必要があります。

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

コマンドフラグの説明

変数	説明
TPU_NAME	キューに格納されたリソースリクエストの割り当て時に作成される TPU のユーザー割り当てテキスト ID。
PROJECT_ID	Google Cloud プロジェクト名。既存のプロジェクトを使用するか、Google Cloud プロジェクトを設定するで新しいプロジェクトを作成します。
ゾーン	サポートされているゾーンについては、TPU のリージョンとゾーンのドキュメントをご覧ください。
ワーカー	基盤となる TPU にアクセスできる TPU VM。

次のコマンドを実行して、デバイスの数を確認できます（ここに表示されている出力は、v5litepod-16 スライスで生成されたものです）。このコードは、JAX が Cloud TPU TensorCore を認識し、基本オペレーションを実行できることを確認することで、すべてが正しくインストールされていることをテストします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

出力は次のようになります。

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device_count() は、指定されたスライス内のチップの合計数を示します。jax.local_device_count() は、このスライス内の単一の VM からアクセス可能なチップの数を示します。

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

出力は次のようになります。

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

このドキュメントの JAX チュートリアルを試して、JAX を使用した v5e トレーニングを開始します。

PyTorch を設定する

v5e は PJRT ランタイムのみをサポートしているのでご注意ください。PyTorch 2.1 以降では、すべての TPU バージョンのデフォルトランタイムとして PJRT が使用されます。

このセクションでは、すべてのワーカー用のコマンドで PyTorch/XLA を使用して v5e 上で PJRT の使用を開始する方法について説明します。

依存関係をインストールする

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

torch、torch_xla、torchvision のホイール（pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222 など）を取り付ける際にエラーが発生した場合は、次のコマンドを使用してバージョンをダウングレードします。

pip3 install setuptools==62.1.0

PJRT でスクリプトを実行する

unset LD_PRELOAD

Python スクリプトを使用して v5e VM で計算を行う例を次に示します。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker all \
   --command='
      export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.local/lib/
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      unset LD_PRELOAD
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"'

これにより、次のような出力が生成されます。

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')

このドキュメントの PyTorch チュートリアルを試して、PyTorch を使用した v5e トレーニングを開始します。

セッションの終了時に TPU とキューに格納されたリソースを削除します。キューに格納されたリソースを削除するには、スライスを削除してから、キューに格納されたリソースを削除する、2 つのステップで行います。

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

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

これら 2 つのステップは、FAILED 状態にある、キューに格納されたリソースリクエストを削除するためにも使用できます。

JAX/FLAX の例

以降のセクションでは、TPU v5e で JAX モデルと FLAX モデルをトレーニングする方法の例について説明します。

v5e で ImageNet をトレーニングする

このチュートリアルでは、架空の入力データを使用して v5e で ImageNet をトレーニングする方法について説明します。実際のデータを使用する場合は、GitHub の README ファイルをご覧ください。

設定

環境変数を作成します。

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

TPU リソースを作成します。
```
gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}
```
注: QUOTA_TYPE フラグは、reserved または best-effort のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。
キューに格納されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。
```
gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}
```
QueuedResource が ACTIVE 状態の場合、出力は次のようになります。
```
 state: ACTIVE
```

最新バージョンの JAX と jaxlib をインストールします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

ImageNet モデルのクローンを作成し、対応する要件をインストールします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/google/flax.git && cd flax/examples/imagenet && pip install -r requirements.txt && pip install flax==0.7.4'

架空のデータを生成するには、データセットのサイズに関する情報が必要です。これは、ImageNet データセットのメタデータから収集できます。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='mkdir -p $HOME/flax/.tfds/metadata/imagenet2012/5.1.0 && curl https://raw.githubusercontent.com/tensorflow/datasets/v4.4.0/tensorflow_datasets/testing/metadata/imagenet2012/5.1.0/dataset_info.json --output $HOME/flax/.tfds/metadata/imagenet2012/5.1.0/dataset_info.json'

モデルをトレーニング

前の手順をすべて完了したら、モデルをトレーニングできます。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd flax/examples/imagenet && JAX_PLATFORMS=tpu python3 imagenet_fake_data_benchmark.py'

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

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

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

Hugging Face FLAX モデル

FLAX で実装された Hugging Face モデルは Cloud TPU v5e ですぐに使用できます。このセクションでは、一般的なモデルの実行手順について説明します。

Imagenette で ViT をトレーニングする

このチュートリアルでは、Cloud TPU v5e で Fast AI Imagenette データセットを使用して、HuggingFace から Vision Transformer（ViT）モデルをトレーニングする方法について説明します。

ViT モデルは、畳み込みネットワークと比較して優れた結果で ImageNet で Transformer エンコーダをトレーニングした最初のモデルです。詳しくは、次のリソースをご覧ください。

設定

環境変数を作成します。

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

TPU リソースを作成します。
```
gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}
```
注: QUOTA_TYPE フラグは、reserved または best-effort のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。
キューに格納されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。
```
gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}
```
キューに格納されたリソースが ACTIVE 状態の場合、出力は次のようになります。
```
 state: ACTIVE
```

JAX とそのライブラリをインストールします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Hugging　Face のリポジトリをダウンロードし、要件をインストールします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow==2.17.0 && pip install -r examples/flax/vision/requirements.txt'

Imagenette データセットをダウンロードします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'

モデルのトレーニング

事前にマッピングされたバッファ（4GB）を使用してモデルをトレーニングします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

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

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

ViT のベンチマーク結果

トレーニングスクリプトは、v5litepod-4、v5litepod-16、v5litepod-64 で実行されました。次の表に、異なるアクセラレータタイプでのスループットを示します。

アクセラレータタイプ	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
グローバルバッチサイズ	32	128	512
スループット（例/秒）	263.40	429.34	470.71

Pokémon データセットで Diffusion をトレーニングする

このチュートリアルでは、Cloud TPU v5e で Pokémon データセットを使用して、HuggingFace から Stable Diffusion モデルをトレーニングする方法について説明します。

Stable Diffusion モデルは、テキスト入力からフォトリアリスティックな画像を生成する、潜在的 text-to-image モデルです。詳しくは、次のリソースをご覧ください。

設定

モデル出力用のストレージバケットを設定します。

gcloud storage buckets create gs://your_bucket 

--project=your_project 

--location=us-west1 

export GCS_BUCKET_NAME=your_bucket

環境変数を作成する

export GCS_BUCKET_NAME=your_bucket
export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west1-c
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

コマンドフラグの説明

変数	説明
GCS_BUCKET_NAME	Google Cloud コンソール -> Cloud Storage -> バケットに表示されます。
PROJECT_ID	Google Cloud プロジェクト名。既存のプロジェクトを使用するか、Google Cloud プロジェクトを設定するで新しいプロジェクトを作成します。
ACCELERATOR_TYPE	TPU のバージョンについては、TPU のバージョンをご覧ください。
ZONE	サポートされているゾーンについては、TPU のリージョンとゾーンのドキュメントをご覧ください。
RUNTIME_VERSION	RUNTIME_VERSION に v2-alpha-tpuv5 を使用します。
SERVICE_ACCOUNT	これは、Google Cloud コンソール -> IAM] -> サービスアカウントで確認できるサービスアカウントのアドレスです。例: tpu-service-account@myprojectID。iam.gserviceaccount.com
TPU_NAME	キューに格納されたリソースリクエストの割り当て時に作成される TPU のユーザー割り当て ID。
QUEUED_RESOURCE_ID	キューに格納されたリソースリクエストのユーザー割り当て ID。キューに格納されたリソースについては、キューに格納されたリソースのドキュメントをご覧ください。
QUOTA_TYPE	`reserved`、`spot` のいずれかを設定できます。どちらも指定されていない場合、QUOTA_TYPE はデフォルトで `on-demand` になります。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。
VALID_UNTIL_DURATION	リクエストが有効である期間。さまざまな有効期間の詳細については、キューに入れられたリソースをご覧ください。

TPU リソースを作成します。
```
gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}
```
注: QUOTA_TYPE フラグは、予約がある場合は reserved、予約がない場合は spot にできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。
キューに格納されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。
```
gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}
```
キューに格納されたリソースが ACTIVE 状態の場合、出力は次のようになります。
```
 state: ACTIVE
```

JAX とそのライブラリをインストールします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

HuggingFace のリポジトリをダウンロードし、要件をインストールします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install tensorflow==2.17.0 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'

モデルのトレーニング

事前にマッピングされたバッファ（4GB）を使用してモデルをトレーニングします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE} --project ${PROJECT_ID} --worker=all --command="
   git clone https://github.com/google/maxdiffusion
   cd maxdiffusion
   git reset --hard 57629bcf4fa32fe5a57096b60b09f41f2fa5c35d # This identifies the GitHub commit to use.
   pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   pip3 install -r requirements.txt
   pip3 install .
   export LIBTPU_INIT_ARGS=""
   python -m src.maxdiffusion.models.train src/maxdiffusion/configs/base_2_base.yml run_name=your_run base_output_directory=gs://${GCS_BUCKET_NAME}/ enable_profiler=False"

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

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

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

拡散のベンチマーク結果

トレーニングスクリプトは、v5litepod-4、v5litepod-16、v5litepod-64 で実行されました。次の表では、スループットを示します。

アクセラレータタイプ	v5litepod-4	v5litepod-16	v5litepod-64
トレーニングステップ	1500	1500	1500
グローバルバッチサイズ	32	64	128
スループット（例/秒）	36.53	43.71	49.36

OSCAR データセットで GPT2 をトレーニングする

このチュートリアルでは、Cloud TPU v5e で OSCAR データセットを使用して、HuggingFace から GPT2 モデルをトレーニングする方法について説明します。

GPT2 は、ヒューマンラベリングのない未加工のテキストで事前にトレーニングされた Transformer モデルです。文の次の単語を予測するようにトレーニングされました。詳しくは、次のリソースをご覧ください。

設定

環境変数を作成します。

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

TPU リソースを作成します。
```
gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}
```
注: QUOTA_TYPE フラグは、reserved または best-effort のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。
キューに格納されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。
```
gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}
```
キューに格納されたリソースが ACTIVE 状態の場合、出力は次のようになります。
```
 state: ACTIVE
```

JAX とそのライブラリをインストールします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

HuggingFace のリポジトリをダウンロードし、要件をインストールします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install TensorFlow && pip install -r examples/flax/language-modeling/requirements.txt'

モデルをトレーニングするための構成ファイルをダウンロードします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt . --recursive'

モデルのトレーニング

事前にマッピングされたバッファ（4GB）を使用してモデルをトレーニングします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers/examples/flax/language-modeling && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 run_clm_flax.py --output_dir=./gpt --model_type=gpt2 --config_name=./gpt --tokenizer_name=./gpt --dataset_name=oscar --dataset_config_name=unshuffled_deduplicated_no --do_train --do_eval --block_size=512 --per_device_train_batch_size=4 --per_device_eval_batch_size=4 --learning_rate=5e-3 --warmup_steps=1000 --adam_beta1=0.9 --adam_beta2=0.98 --weight_decay=0.01 --overwrite_output_dir --num_train_epochs=3 --logging_steps=500 --eval_steps=2500'

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

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

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

GPT2 のベンチマーク結果

トレーニングスクリプトは、v5litepod-4、v5litepod-16、v5litepod-64 で実行されました。次の表では、スループットを示します。

	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
グローバルバッチサイズ	64	64	64
スループット（例/秒）	74.60	72.97	72.62

PyTorch/XLA

以下のセクションでは、TPU v5e で PyTorch/XLA モデルをトレーニングする方法の例について説明します。

PJRT ランタイムを使用して ResNet をトレーニングする

PyTorch/XLA は、PyTorch 2.0 以降、XRT から PjRt に移行しています。PyTorch/XLA トレーニングワークロード用に v5e を設定する手順が更新されています。

設定

環境変数を作成します。

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=tpu-name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

TPU リソースを作成します。
```
gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --{QUOTA_TYPE}
```
注: QUOTA_TYPE フラグは、reserved または best-effort のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。
QueuedResource が ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。
```
gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}
```
キューに格納されたリソースが ACTIVE 状態の場合、出力は次のようになります。
```
 state: ACTIVE
```

Torch/XLA 固有の依存関係をインストールする

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

ResNet モデルをトレーニングする

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      date
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export XLA_USE_BF16=1
      export LIBTPU_INIT_ARGS=--xla_jf_auto_cross_replica_sharding
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      git clone https://github.com/pytorch/xla.git
      cd xla/
      git reset --hard caf5168785c081cd7eb60b49fe4fffeb894c39d9
      python3 test/test_train_mp_imagenet.py --model=resnet50  --fake_data --num_epochs=1 —num_workers=16  --log_steps=300 --batch_size=64 --profile'

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

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

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

ベンチマークの結果

次の表に、ベンチマークのスループットを示します。

アクセラレータタイプ	スループット（例/秒）
v5litepod-4	4240 ex/s
v5litepod-16	10,810 ex/s
v5litepod-64	46,154 ex/s

v5e で GPT2 をトレーニングする

このチュートリアルでは、wikitext データセットを使用し、PyTorch/XLA で HuggingFace のリポジトリを使用して、v5e で GPT2 を実行する方法を説明します。

設定

環境変数を作成します。

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

TPU リソースを作成します。
```
gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}
```
注: QUOTA_TYPE フラグは、reserved または best-effort のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。
QueuedResource が ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。
```
gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}
```
キューに格納されたリソースが ACTIVE 状態の場合、出力は次のようになります。
```
state: ACTIVE
```

PyTorch/XLA の依存関係をインストールします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get -y update
      sudo apt install -y libopenblas-base
      pip3 install torchvision
      pip3 uninstall -y torch
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

HuggingFace のリポジトリをダウンロードし、要件をインストールします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      git clone https://github.com/pytorch/xla.git
      pip install --upgrade accelerate
      git clone https://github.com/huggingface/transformers.git
      cd transformers
      git checkout ebdb185befaa821304d461ed6aa20a17e4dc3aa2
      pip install .
      git log -1
      pip install datasets evaluate scikit-learn
      '

事前トレーニング済みモデルの構成をダウンロードします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/ --recursive
      gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'

モデルのトレーニング

バッチサイズ 16 を使用して 2B モデルをトレーニングします。

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      cd transformers/
      export LD_LIBRARY_PATH=/usr/local/lib/
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      python3 examples/pytorch/xla_spawn.py \
         --num_cores=4 \
         examples/pytorch/language-modeling/run_clm.py \
         --num_train_epochs=3 \
         --dataset_name=wikitext \
         --dataset_config_name=wikitext-2-raw-v1 \
         --per_device_train_batch_size=16 \
         --per_device_eval_batch_size=16 \
         --do_train \
         --do_eval \
         --logging_dir=./tensorboard-metrics \
         --cache_dir=./cache_dir \
         --output_dir=/tmp/test-clm \
         --overwrite_output_dir \
         --cache_dir=/tmp \
         --config_name=examples/pytorch/language-modeling/my_config_2.json \
         --tokenizer_name=gpt2 \
         --block_size=1024 \
         --optim=adafactor \
         --adafactor=true \
         --save_strategy=no \
         --logging_strategy=no \
         --fsdp=full_shard \
         --fsdp_config=examples/pytorch/language-modeling/fsdp_config.json'

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

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

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

ベンチマークの結果

トレーニングスクリプトは、v5litepod-4、v5litepod-16、v5litepod-64 で実行されました。次の表に、さまざまなアクセラレータタイプのベンチマークスループットを示します。

	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
config	6 億	20 億	16B
グローバルバッチサイズ	64	128	256
スループット（例/秒）	66	77	31

v5e で ViT をトレーニングする

このチュートリアルでは、cifar10 データセット上の PyTorch/XLA で HuggingFace のリポジトリを使用して、v5e で VIT を実行する方法について説明します。

設定

環境変数を作成します。

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export RUNTIME_VERSION=v2-alpha-tpuv5-lite
export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=tpu-name
export QUEUED_RESOURCE_ID=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

TPU リソースを作成します。
```
gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}
```
注: QUOTA_TYPE フラグは、reserved または best-effort のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。
QueuedResource が ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。
```
 gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}
```
キューに格納されたリソースが ACTIVE 状態の場合、出力は次のようになります。
```
 state: ACTIVE
```

PyTorch/XLA の依存関係をインストールする

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all
   --command='
      sudo apt-get update -y
      sudo apt-get install libomp5 -y
      pip3 install mkl mkl-include
      pip3 install tf-nightly tb-nightly tbp-nightly
      pip3 install numpy
      sudo apt-get install libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'

HuggingFace のリポジトリをダウンロードし、要件をインストールします。

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command="
      git clone https://github.com/suexu1025/transformers.git vittransformers; \
      cd vittransformers; \
      pip3 install .; \
      pip3 install datasets; \
      wget https://github.com/pytorch/xla/blob/master/scripts/capture_profile.py"

モデルのトレーニング

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export TF_CPP_MIN_LOG_LEVEL=0
      export XLA_USE_BF16=1
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      cd vittransformers
      python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/image-pretraining/run_mae.py --dataset_name=cifar10 \
      --remove_unused_columns=False \
      --label_names=pixel_values \
      --mask_ratio=0.75 \
      --norm_pix_loss=True \
      --do_train=true \
      --do_eval=true \
      --base_learning_rate=1.5e-4 \
      --lr_scheduler_type=cosine \
      --weight_decay=0.05 \
      --num_train_epochs=3 \
      --warmup_ratio=0.05 \
      --per_device_train_batch_size=8 \
      --per_device_eval_batch_size=8 \
      --logging_strategy=steps \
      --logging_steps=30 \
      --evaluation_strategy=epoch \
      --save_strategy=epoch \
      --load_best_model_at_end=True \
      --save_total_limit=3 \
      --seed=1337 \
      --output_dir=MAE \
      --overwrite_output_dir=true \
      --logging_dir=./tensorboard-metrics \
      --tpu_metrics_debug=true'

TPU とキューに格納されたリソースを削除する

セッションの終了時に TPU とキューに格納されたリソースを削除します。

gcloud compute tpus tpu-vm delete ${TPU_NAME}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
   --project=${PROJECT_ID}
   --zone=${ZONE}
   --quiet

ベンチマークの結果

次の表に、さまざまなアクセラレータタイプのベンチマークスループットを示します。

	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	3	3	3
グローバルバッチサイズ	32	128	512
スループット（例/秒）	201	657	2,844

TensorFlow 2.x

以降のセクションでは、TPU v5e で TensorFlow 2.x モデルをトレーニングする方法の例について説明します。

単一ホスト v5e で Resnet をトレーニングする

このチュートリアルでは、架空のデータセットを使用して v5litepod-4 または v5litepod-8 で ImageNet をトレーニングする方法について説明します。別のデータセットを使用する場合は、データセットの準備をご覧ください。

設定

環境変数を作成します。

export PROJECT_ID=your-project-ID
export ACCELERATOR_TYPE=v5litepod-4
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pjrt
export TPU_NAME=your-tpu-name
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

ACCELERATOR_TYPE は、v5litepod-4 または v5litepod-8 のいずれかになります。

TPU リソースを作成します。
```
gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --${QUOTA_TYPE}
```
注: QUOTA_TYPE フラグは、reserved または best-effort のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。
キューに格納されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。キューに格納されたリソースの状態を確認するには、次のコマンドを使用します。
```
gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}
```

SSH を使用して TPU に接続します

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

いくつかの環境変数を設定します

export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

モデルリポジトリのディレクトリに移動し、要件をインストールします。
```
cd ${MODELS_REPO} && git checkout r2.15.0
pip install -r official/requirements.txt
```

モデルのトレーニング

トレーニングスクリプトを実行します。

python3 official/vision/train.py \
   --tpu=local \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

TPU とキューに格納されたリソースを削除する

TPU を削除します

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

キューに格納されたリソースリクエストを削除します

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

マルチホスト v5e で Resnet をトレーニングする

このチュートリアルでは、架空のデータセットを使用して v5litepod-16 以上での ImageNet をトレーニングする方法について説明します。別のデータセットを使用する場合は、データセットの準備をご覧ください。

環境変数を作成します。

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pod-pjrt
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

ACCELERATOR_TYPE は v5litepod-16 かそれ以上にすることができます。

TPU リソースを作成します。
```
gcloud compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
   --node-id=${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --accelerator-type=${ACCELERATOR_TYPE} \
   --runtime-version=${RUNTIME_VERSION} \
   --${QUOTA_TYPE}
```
注: QUOTA_TYPE フラグは、reserved または best-effort のいずれかにできます。どちらも指定されていない場合、デフォルトは on-demand です。Cloud TPU でサポートされている割り当てのさまざまなタイプについては、割り当てをご覧ください。
キューに格納されたリソースが ACTIVE 状態になると、TPU VM に SSH 接続できるようになります。キューに格納されたリソースの状態を確認するには、次のコマンドを使用します。
```
gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}
```

SSH を使用して TPU（ワーカーゼロ）に接続します

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

いくつかの環境変数を設定します

export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export TPU_LOAD_LIBRARY=0
export TPU_NAME=your_tpu_name

モデルリポジトリのディレクトリに移動し、要件をインストールします。
```
 cd $MODELS_REPO && git checkout r2.15.0
 pip install -r official/requirements.txt
```

モデルのトレーニング

トレーニングスクリプトを実行します。

python3 official/vision/train.py \
   --tpu=${TPU_NAME} \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*"

TPU とキューに格納されたリソースを削除する

TPU を削除します

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

キューに格納されたリソースリクエストを削除します

gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet