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

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

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 バケットを作成する

    gsutil mb -c standard -l europe-west4 gs://your-bucket-name

  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.16.1-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.16.1-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 ランタイムを構成する

    PyTorch/XLA のランタイム オプションは PJRT と XRT の 2 つです。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

TPU Pod で PyTorch コードを実行する場合は、すべての TPU ワーカーでコードを同時に実行する必要があります。これを行う方法の 1 つに、--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 コードを同時に実行する必要があります。これを行う方法の 1 つに、--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. Artifact Registry に push する前に、Docker イメージにタグを追加します。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

次のステップ