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

Docker-Container erleichtern die Konfiguration von Anwendungen, da sie den Code und alle erforderlichen Abhängigkeiten in einem verteilbaren Paket kombinieren. Sie können Docker-Container auf 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.

Wichtig: Sie können maximal 256 Docker-Pods auf einer TPU-VM ausführen.

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.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-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.16.1-pjrt
   

4. Dockerfile in 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
   

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"
   

Wenn das Trainingsskript abgeschlossen ist, sollten Sie die Ressourcen bereinigen.

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

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.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 erstellen

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

Wenn das Trainingsskript abgeschlossen ist, sollten Sie die Ressourcen bereinigen.

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

TPU-Gerät

1. Cloud TPU-VM erstellen

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

2. SSH-Verbindung zu TPU-VM herstellen

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

3. Starten Sie auf der TPU-VM einen Container mit dem nächtlichen PyTorch/XLA-Image.

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

   Hinweis: Nach dem Ausführen dieses Befehls ändert sich die Eingabeaufforderung, sobald das Terminal mit dem ausgeführten Container verbunden ist.

4. TPU-Laufzeit konfigurieren

   Es gibt zwei PyTorch/XLA-Laufzeitoptionen: PJRT und XRT. Wir empfehlen die Verwendung von PJRT, sofern es keinen Grund dafür gibt. Wenn Sie mehr über die verschiedenen Laufzeitkonfigurationen erfahren möchten, sehen Sie, dass es einen Grund gibt, XRT zu verwenden. Weitere Informationen zu den verschiedenen Laufzeitkonfigurationen finden Sie in der PJRT-Laufzeitdokumentation.

   PJRT

   export PJRT_DEVICE=TPU
   

   XRT

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

5. PyTorch XLA-Repository klonen

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

6. ResNet50 trainieren

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

Wenn das Trainingsskript abgeschlossen ist, sollten Sie die Ressourcen bereinigen.

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 PyTorch-Code auf einem TPU-Pod ausführen, müssen Sie den Code gleichzeitig auf allen TPU-Workern ausführen. Eine Möglichkeit dafür ist die Verwendung des Befehls gcloud compute tpus tpu-vm ssh mit den Flags --worker=all und --command. Das folgende Verfahren zeigt Ihnen, wie Sie ein Docker-Image erstellen, um die Einrichtung der einzelnen TPU-Worker zu vereinfachen.

1. TPU-VM erstellen

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

2. Aktuellen Nutzer der Docker-Gruppe hinzufügen

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

3. Führen Sie das Trainingsskript in einem Container auf allen TPU-Workern aus.

   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-Befehls-Flags:

   • --rm entfernt den Container, nachdem sein Prozess beendet wurde.
   • --privileged stellt das TPU-Gerät für den Container bereit.
   • --net=host bindet alle Ports des Containers an die TPU-VM, um die Kommunikation zwischen den Hosts im Pod zu ermöglichen.
   • -e legt Umgebungsvariablen fest.

Wenn das Trainingsskript abgeschlossen ist, sollten Sie die Ressourcen bereinigen.

Löschen Sie die TPU-VM mit dem folgenden Befehl:

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

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-ubuntu2204-base
   

2. SSH-Verbindung zu 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] -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
   

Wenn das Trainingsskript abgeschlossen ist, sollten Sie die Ressourcen bereinigen.

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, muss er gleichzeitig auf allen TPU-Workern ausgeführt werden. Eine Möglichkeit dazu ist der Befehl gcloud compute tpus tpu-vm ssh mit den Flags --worker=all und --command. Im folgenden Verfahren erfahren Sie, wie Sie ein Docker-Image erstellen, um die Einrichtung der TPU-Worker zu vereinfachen.

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 "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-Image erstellen

   docker build -t your-image-name .
   

3. Fügen Sie dem Docker-Image ein Tag hinzu, bevor Sie es per Push an Artifact Registry übertragen. Weitere Informationen zum Arbeiten mit Artifact Registry finden Sie unter Mit Container-Images arbeiten.

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

4. Docker-Image per Push an Artifact Registry übertragen

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

5. TPU-VM erstellen

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

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

   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. Führen Sie den Container auf allen TPU-Workern aus.

   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. Führen Sie das Trainingsskript auf allen TPU-Workern aus:

   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"
   

Wenn das Trainingsskript abgeschlossen ist, sollten Sie die Ressourcen bereinigen.

1. Fahren Sie den Container auf allen Workern herunter:

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

2. 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
• TPUs verwalten
• Cloud TPU-Systemarchitektur
• TensorFlow-Code auf TPU-Pod-Slices ausführen
• JAX-Code auf TPU-Pod-Slices ausführen