Cloud TPU-Anwendungen in einem Docker-Container ausführen

Docker-Container erleichtern die Konfiguration von Anwendungen, indem sie Ihren Code und alle benötigten Abhängigkeiten in einem verteilbaren Paket kombinieren. Sie können Docker-Container in TPU-VMs ausführen, um die Konfiguration und Freigabe Ihrer Cloud TPU-Anwendungen zu vereinfachen. In diesem Dokument wird beschrieben, wie Sie einen Docker-Container für jedes von Cloud TPU unterstützte ML-Framework einrichten.

TensorFlow-Modell in einem Docker-Container trainieren

TPU-Gerät

1. Erstellen Sie in Ihrem aktuellen Verzeichnis eine Datei mit dem Namen Dockerfile und fügen Sie den folgenden Text ein

   FROM python:3.8
   RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.10.0/tensorflow-2.10.0-cp38-cp38-linux_x86_64.whl
   RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.4.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-Bucket erstellen

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

3. TPU-VM erstellen

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

4. Dockerfile auf die TPU-VM kopieren

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

5. Stellen Sie eine SSH-Verbindung zur TPU-VM her:

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

6. Docker-Image erstellen

   sudo docker build -t your-image-name .
   

7. Docker-Container starten

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

8. Umgebungsvariablen festlegen

   export STORAGE_BUCKET=gs://your-bucket-name
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export MODEL_DIR=${STORAGE_BUCKET}/resnet-2x
   export ZONE=europe-west4-a
   export TPU_NAME=local
   

9. ResNet trainieren

   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"
   

TPU-Pod

1. Erstellen Sie in Ihrem aktuellen Verzeichnis eine Datei mit dem Namen Dockerfile und fügen Sie den folgenden Text ein

   FROM python:3.8
   RUN pip install https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/tensorflow/tf-2.10.0/tensorflow-2.10.0-cp38-cp38-linux_x86_64.whl
   RUN curl -L https://storage.googleapis.com/cloud-tpu-tpuvm-artifacts/libtpu/1.4.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 erstellen

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

3. Dockerfile auf die TPU-VM kopieren

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

4. Stellen Sie eine SSH-Verbindung zur TPU-VM her:

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

5. Docker-Image erstellen

   sudo docker build -t your-image-name .
   

6. Docker-Container starten

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

7. ResNet trainieren

   python3 official/vision/train.py \
   --tpu=your-tpu-name \
   --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"
   

Bereinigen Sie nach Abschluss des Trainingsskripts die Ressourcen.

1. Geben Sie exit ein, um den Docker-Container zu beenden
2. Geben Sie exit ein, um die TPU-VM zu beenden
3. TPU-VM löschen

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

PyTorch-Modell in einem Docker-Container trainieren

Hinweis: Derzeit wird das Ausführen von PyTorch in einem Docker-Container nur auf TPUs mit einem einzelnen Gerät (v2-8, v3-8) unterstützt.

1. Cloud TPU-VM erstellen

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

2. SSH-Verbindung zur TPU-VM herstellen

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

3. Container in der TPU-VM starten

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

4. libtpu installieren

   pip install torch_xla[tpuvm]
   export XRT_TPU_CONFIG="localservice;0;localhost:51011"
   

5. PyTorch XLA-Repository klonen

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

6. ImageNet trainieren

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

Bereinigen Sie nach Abschluss des Trainingsskripts die Ressourcen.

1. Geben Sie exit ein, um den Docker-Container zu beenden
2. Geben Sie exit ein, um die TPU-VM zu beenden
3. TPU-VM löschen

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

JAX-Modell in einem Docker-Container trainieren

TPU-Gerät

1. TPU-VM erstellen

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

2. SSH-Verbindung zur TPU-VM herstellen

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

3. Docker-Daemon in TPU-VM starten

   sudo systemctl start docker
   

4. Docker-Container starten

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

5. JAX installieren

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

6. FLAX installieren

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

7. FLAX MNIST-Trainingsskript ausführen

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

Bereinigen Sie nach Abschluss des Trainingsskripts die Ressourcen.

1. Geben Sie exit ein, um den Docker-Container zu beenden
2. Geben Sie exit ein, um die TPU-VM zu beenden

3. TPU-VM löschen

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

TPU-Pod

Wenn Sie JAX-Code auf einem TPU-Pod ausführen, müssen Sie ihn auf allen TPU-Workern gleichzeitig ausführen. Eine Möglichkeit dazu ist, den Befehl gcloud compute tpus tpu-vm ssh mit den Flags --worker=all und --command zu verwenden. Das folgende Verfahren zeigt, wie Sie ein Docker-Image erstellen, um die Einrichtung der einzelnen TPU-Worker zu vereinfachen.

1. Legen Sie in einer Terminalaufforderung die folgenden Umgebungsvariablen fest.

   export USER=your-user-id
   export PROJECT=your-project-name
   export REPO_NAME=your-repo-name
   export IMAGE_NAME=your-image-name
   export TAG=your-tag
   export TPU_NAME=your-tpu-name
   export ZONE=europe-west4-a
   export ACCEL_TYPE=v2-8
   

   Ersetzen Sie:

   • your-user-id durch Ihre Nutzer-ID
   • your-project-name durch den Namen Ihres Google Cloud-Projekts
   • your-repo-name durch den Namen Ihres Artifact Registry-Repositorys
   • your-image-name durch einen Namen für Ihr Docker-Image
   • your-tag mit einem Tag für das Docker-Image
   • your-tpu-name durch einen Namen für Ihre TPU

2. Erstellen Sie in Ihrem aktuellen Verzeichnis eine Datei mit dem Namen Dockerfile und fügen Sie den folgenden Text ein

   FROM google-cloud-cli:latest
   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
   

3. Docker-Image erstellen

   docker build -t $IMAGE .
   

4. Fügen Sie Ihrem Docker-Image ein Tag hinzu, bevor Sie es in die Artifact Registry übertragen. Weitere Informationen zum Arbeiten mit Artifact Registry finden Sie unter Mit Container-Images arbeiten.

   docker tag $IMAGE_NAME europe-west-docker.pkg.dev/$PROJECT_NAME/$REPO_NAME/$IMAGE_NAME:$TAG
   

5. Docker-Image per Push an die Artifact Registry übertragen

   docker push europe-west4-docker.pkg.dev/ml-writers/$REPO_NAME/$IMAGE_NAME:$TAG
   

6. TPU-VM erstellen

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone $ZONE \
   --accelerator-type $ACCEL_TYPE \
   --version tpu-vm-base
   

7. Rufen Sie das Docker-Image aus der Artifact Registry auf allen TPU-Workern ab.

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --command="sudo usermod -a -G docker $USER"
   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"
   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --command="docker pull europe-west4-docker.pkg.dev/your-project-name/your-repo-name/your-image-name:your-tag"
   

8. Container auf allen TPU-Workern ausführen

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

9. Führen Sie das Trainingsskript auf allen TPU-Workern aus:

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone $ZONE \
   --command="docker exec --privileged $CONTAINER_NAME python3 main.py --workdir=/tmp/mnist \
   --config=configs/default.py \
   --config.learning_rate=0.05 \
   --config.num_epochs=5"
   

Bereinigen Sie nach Abschluss des Trainingsskripts die Ressourcen. Löschen Sie die TPU-VM mit dem folgenden Befehl:

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

Nächste Schritte

• Cloud TPU-Anleitungen
• Cloud TPU-Nutzerhandbuch
• Cloud TPU-Systemarchitektur
• TensorFlow-Code auf TPU Pod-Slices ausführen
• JAX-Code auf TPU Pod-Slices ausführen