Docker コンテナで Cloud TPU アプリケーションを実行する

Docker コンテナでは、コードとすべての必要な依存関係が 1 つの配布可能なパッケージにまとまることで、アプリケーションの構成が簡単になります。TPU VM 内で Docker コンテナを実行すると、Cloud TPU アプリケーションの構成と共有を簡素化できます。このドキュメントでは、Cloud TPU でサポートされている各 ML フレームワークに Docker コンテナを設定する方法について説明します。

重要: 1 つの TPU VM で最大 256 個の Docker Pod を実行できます。

Docker コンテナで TensorFlow モデルをトレーニングする

TPU デバイス

1. 現在のディレクトリに Dockerfile という名前のファイルを作成し、次のテキストを貼り付ける

   FROM python:3.8
   RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl
   RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so
   RUN git clone https://github.com/tensorflow/models.git
   WORKDIR ./models
   RUN pip install -r official/requirements.txt
   ENV PYTHONPATH=/models
   

2. Cloud Storage バケットを作成する

   gcloud storage buckets create gs://your-bucket-name --location=europe-west4
   

3. TPU VM を作成する

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-vm-tf-2.17.0-pjrt
   

4. Dockerfile を TPU VM にコピーする

   gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
   

5. TPU VM に SSH 接続する

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=europe-west4-a
   

6. Docker イメージをビルドする

   sudo docker build -t your-image-name . 
   

7. Docker コンテナを起動する

   sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
   

8. 環境変数を設定する

   export STORAGE_BUCKET=gs://your-bucket-name
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x
   

9. ResNet50 をトレーニングする

   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="task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*,trainer.train_steps=100"
   

トレーニング スクリプトが完了したら、必ずリソースをクリーンアップします。

1. 「exit」と入力して Docker コンテナを終了する
2. 「exit」と入力して TPU VM を終了する
3. TPU VM を削除する

    $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
   

TPU Pod

1. 現在のディレクトリに Dockerfile という名前のファイルを作成し、次のテキストを貼り付ける

   FROM python:3.8
   RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.12.0/tensorflow-2.12.0-cp38-cp38-linux_x86_64.whl
   RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.6.0/libtpu.so -o /lib/libtpu.so
   RUN git clone https://github.com/tensorflow/models.git
   WORKDIR ./models
   RUN pip install -r official/requirements.txt
   ENV PYTHONPATH=/models
   

2. TPU VM を作成する

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v3-32 \
   --version=tpu-vm-tf-2.17.0-pod-pjrt
   

3. Dockerfile を TPU VM にコピーする

   gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
   

4. TPU VM に SSH 接続する

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=europe-west4-a
   

5. Docker イメージをビルドする

   sudo docker build -t your-image-name . 
   

6. Docker コンテナを起動する

   sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
   

7. ResNet50 をトレーニングする

   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="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"
   

トレーニング スクリプトが完了したら、必ずリソースをクリーンアップします。

1. 「exit」と入力して Docker コンテナを終了する
2. 「exit」と入力して TPU VM を終了する
3. TPU VM を削除する

     $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
   

Docker コンテナで PyTorch モデルをトレーニングする

TPU デバイス

1. Cloud TPU VM を作成する

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-ubuntu2204-base
   

2. TPU VM に SSH 接続する

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=europe-west4-a
   

3. 夜間の PyTorch/XLA イメージを使用して TPU VM でコンテナを起動します。

   sudo docker run -ti --rm --name your-container-name --privileged gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm bash
   

   注: このコマンドを実行すると、コマンド プロンプトが変更され、ターミナルが実行中のコンテナに接続されていることが示されます。

4. TPU ランタイムを構成する

   PJRT と XRT の 2 つの PyTorch/XLA ランタイム オプションがあります。XRT を使用する理由がない限り、PJRT を使用することをおすすめします。さまざまなランタイム構成の詳細については、XRT を使用する理由をご覧ください。さまざまなランタイム構成の詳細については、PJRT ランタイム ドキュメントをご覧ください。

   PJRT

   export PJRT_DEVICE=TPU
   

   XRT

   export XRT_TPU_CONFIG="localservice;0;localhost:51011"
   

5. PyTorch XLA リポジトリのクローンを作成する

   git clone --recursive https://github.com/pytorch/xla.git
   

6. ResNet50 をトレーニングする

   python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
   

トレーニング スクリプトが完了したら、必ずリソースをクリーンアップします。

1. 「exit」と入力して Docker コンテナを終了する
2. 「exit」と入力して TPU VM を終了する
3. TPU VM を削除する

    $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
   

TPU Pod

PyTorch コードを TPU Pod で実行する場合は、すべての TPU ワーカーで同時にコードを実行する必要があります。これを行う方法の一つに、--worker=all フラグと --command フラグを指定して gcloud compute tpus tpu-vm ssh コマンドを使用する方法があります。次の手順では、各 TPU ワーカーを簡単に設定できるように Docker イメージを作成する方法を示します。

1. TPU VM を作成する

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

2. 現在のユーザーを Docker グループに追加する

   gcloud compute tpus tpu-vm ssh your-tpu-name \
   --zone=us-central2-b \
   --worker=all \
   --command="sudo usermod -a -G docker $USER"
   

3. すべての TPU ワーカーのコンテナでトレーニング スクリプトを実行します。

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=us-central2-b \
   --command="docker run --rm --privileged --net=host  -e PJRT_DEVICE=TPU gcr.io/tpu-pytorch/xla:r2.0_3.8_tpuvm python /pytorch/xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1"
   

   Docker コマンドフラグ:

   • --rm は、プロセスが終了した後にコンテナを削除します。
   • --privileged は、TPU デバイスをコンテナに公開します。
   • --net=host は、コンテナのすべてのポートを TPU VM にバインドして、Pod 内のホスト間の通信を許可します。
   • -e は、環境変数を設定します。

トレーニング スクリプトが完了したら、必ずリソースをクリーンアップします。

次のコマンドを使用して TPU VM を削除します。

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

Docker コンテナで JAX モデルをトレーニングする

TPU デバイス

1. TPU VM を作成します。

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type=v2-8 \
   --version=tpu-ubuntu2204-base
   

2. TPU VM に SSH 接続する

   gcloud compute tpus tpu-vm ssh your-tpu-name  --zone=europe-west4-a
   

3. TPU VM で Docker デーモンを起動する

   sudo systemctl start docker
   

4. Docker コンテナを起動する

   sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash
   

5. JAX をインストールする

   pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   

6. FLAX のインストール

   pip install --upgrade clu
   git clone https://github.com/google/flax.git
   pip install --user -e flax
   

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

   cd flax/examples/mnist
   python3 main.py --workdir=/tmp/mnist \
   --config=configs/default.py \
   --config.learning_rate=0.05 \
   --config.num_epochs=5
   

トレーニング スクリプトが完了したら、必ずリソースをクリーンアップします。

1. 「exit」と入力して Docker コンテナを終了する
2. 「exit」と入力して TPU VM を終了する

3. TPU VM を削除する

   $ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
   

TPU Pod

TPU Pod で JAX コードを実行する場合は、すべての TPU ワーカーで同時に JAX コードを実行する必要があります。これを行う方法の一つに、--worker=all フラグと --command フラグを指定して gcloud compute tpus tpu-vm ssh コマンドを使用する方法があります。次の手順では、各 TPU ワーカーを簡単に設定できるように Docker イメージを作成する方法を示します。

1. 現在のディレクトリに Dockerfile という名前のファイルを作成し、次のテキストを貼り付ける

   FROM python:3.8
   RUN pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
   RUN pip install --upgrade clu
   RUN git clone https://github.com/google/flax.git
   RUN pip install --user -e flax
   WORKDIR ./flax/examples/mnist
   

2. Docker イメージをビルドする

   docker build -t your-image-name .
   

3. Docker イメージにタグを追加してから、Artifact Registry に push します。Artifact Registry の操作の詳細については、コンテナ イメージの使用をご覧ください。

   docker tag your-image-name europe-west-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
   

4. Docker イメージを Artifact Registry に push する

   docker push europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
   

5. TPU VM を作成する

   gcloud compute tpus tpu-vm create your-tpu-name \
   --zone=europe-west4-a \
   --accelerator-type==v2-8 \
   --version=tpu-ubuntu2204-base
   

6. すべての TPU ワーカーで Artifact Registry から Docker イメージを pull します。

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="sudo usermod -a -G docker ${USER}"
   

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"
   

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker pull europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag"
   

7. すべての TPU ワーカーでコンテナを実行します。

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   zone=europe-west4-a \
   --command="docker run -ti -d --privileged --net=host --name your-container-name europe-west4-docker.pkg.dev/your-project/your-repo/your-image:your-tag bash"
   

8. すべてのワーカーでトレーニング スクリプトを実行する

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker exec --privileged your-container-name python3 main.py --workdir=/tmp/mnist \
   --config=configs/default.py \
   --config.learning_rate=0.05 \
   --config.num_epochs=5"
   

トレーニング スクリプトが完了したら、必ずリソースをクリーンアップします。

1. すべてのワーカーでコンテナをシャットダウンします。

   gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \
   --zone=europe-west4-a \
   --command="docker kill your-container-name"
   

2. 次のコマンドを使用して TPU VM を削除します。

   $ gcloud compute tpus tpu-vm delete your-tpu-name \
   --zone=europe-west4-a
   

次のステップ

• Cloud TPU のチュートリアル
• TPU を管理する
• Cloud TPU システム アーキテクチャ
• TPU Pod スライスでの TensorFlow コードの実行
• TPU Pod スライスでの JAX コードの実行