TPU-Arbeitslasten in einem Docker-Container ausführen

Achtung: Diese Seite bezieht sich auf die Cloud TPU API, die nicht mehr aktiv weiterentwickelt wird. Die Cloud TPU API erhält nur Fehlerkorrekturen und Sicherheitsupdates. Für die neuesten Funktionen und die Unterstützung der neuesten TPU-Versionen empfehlen wir, Compute Engine oder Google Kubernetes Engine zum Verwalten von TPU-Ressourcen zu verwenden. Weitere Informationen finden Sie unter TPU-Ressourcen in Compute Engine und TPUs in GKE. Die Cloud TPU API wird auf TPU7x und höheren TPU-Versionen nicht unterstützt.

Docker-Container erleichtern die Konfiguration von Anwendungen, indem Ihr Code und alle erforderlichen Abhängigkeiten in einem verteilbaren Paket zusammengefasst werden. 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.

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

PyTorch-Modell in einem Docker-Container trainieren

Unter PyTorch auf Cloud TPU werden die PyTorch-Bibliotheken für TPU beschrieben und Links zur Dokumentation zum Abrufen der Python-Pakete und Docker-Images bereitgestellt. Sie können TorchTPU direkt verwenden (Informationen zum Container-Image für die öffentliche Version finden Sie unter Erste Schritte mit GKE) oder die Referenzbibliothek TorchTitan verwenden, um containerisierte Trainingsarbeitslasten zu erstellen und auszuführen.

JAX-Modell in einem Docker-Container trainieren

Erstellen Sie vor dem Ausführen der folgenden Befehle die folgenden Umgebungsvariablen:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-west4-a
export ACCELERATOR_TYPE=v5litepod-8
export RUNTIME_VERSION=v2-alpha-tpuv5-lite

Beschreibungen von Umgebungsvariablen

• PROJECT_ID: Ihre Google Cloud -Projekt-ID. Verwenden Sie ein vorhandenes Projekt oder erstellen Sie ein neues Projekt.
• TPU_NAME: Der Name der TPU.
• ZONE: Die Zone, in der die TPU-VM erstellt werden soll. Weitere Informationen zu unterstützten Zonen finden Sie unter TPU-Regionen und ‑Zonen.
• ACCELERATOR_TYPE: Der Beschleunigertyp gibt die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Weitere Informationen zu den unterstützten Beschleunigertypen für die einzelnen TPU-Versionen finden Sie unter TPU-Versionen.
• RUNTIME_VERSION: Die Cloud TPU-Softwareversion.

TPU-Gerät

1. Erstellen Sie eine TPU-VM.

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

2. Stellen Sie eine SSH-Verbindung zur TPU-VM her.

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE
   

3. Starten Sie den Docker-Daemon in der TPU-VM.

   sudo systemctl start docker
   

4. Starten Sie den Docker-Container.

   sudo docker run --net=host -ti --rm --name your-container-name \
   --privileged us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 \
   bash
   

5. Installieren Sie JAX.

   pip install jax[tpu]
   

6. Installieren Sie FLAX.

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

7. Installieren Sie das tensorflow- und das tensorflow-datasets-Paket.

   pip install tensorflow
   pip install tensorflow-datasets
   

8. Führen Sie das FLAX-MNIST-Trainings-Script aus.

   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 Trainings-Script abgeschlossen ist, bereinigen Sie 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. Löschen Sie die TPU-VM.

   gcloud compute tpus tpu-vm delete $TPU_NAME --zone=$ZONE
   

TPU-Slice

Wenn Sie JAX-Code auf einem TPU-Slice ausführen, müssen Sie den JAX-Code gleichzeitig auf allen TPU-Workern ausführen. Eine Möglichkeit hierfür ist die Verwendung des Befehls gcloud compute tpus tpu-vm ssh mit den Flags --worker=all und --command. Im folgenden Verfahren wird gezeigt, wie Sie ein Docker-Image erstellen, um die Einrichtung der einzelnen TPU-Worker zu vereinfachen.

1. Erstellen Sie in Ihrem aktuellen Verzeichnis eine Datei namens Dockerfile und fügen Sie den folgenden Text ein:

   FROM python:3.10
   RUN pip install jax[tpu]
   RUN pip install --upgrade clu
   RUN git clone https://github.com/google/flax.git
   RUN pip install --user -e flax
   RUN pip install tensorflow
   RUN pip install tensorflow-datasets
   WORKDIR ./flax/examples/mnist
   

2. Bereiten Sie Artifact Registry vor.

   gcloud artifacts repositories create your-repo \
   --repository-format=docker \
   --location=europe-west4 --description="Docker repository" \
   --project=$PROJECT_ID
   
   gcloud artifacts repositories list \
   --project=$PROJECT_ID
   
   gcloud auth configure-docker europe-west4-docker.pkg.dev
   

3. Erstellen Sie das Docker-Image.

   docker build -t your-image-name .
   

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

   docker tag your-image-name europe-west4-docker.pkg.dev/$PROJECT_ID/your-repo/your-image-name:your-tag
   

5. Übertragen Sie das Docker-Image per Push an Artifact Registry.

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

6. Erstellen Sie eine TPU-VM.

   gcloud compute tpus tpu-vm create $TPU_NAME \
   --zone=$ZONE \
   --accelerator-type=$ACCELERATOR_TYPE \
   --version=$RUNTIME_VERSION
   

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

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command='sudo usermod -a -G docker ${USER}'
   

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="gcloud auth configure-docker europe-west4-docker.pkg.dev --quiet"
   

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

8. Führen Sie den Container auf allen TPU-Workern aus.

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

9. Führen Sie das Trainings-Script auf allen TPU-Workern aus.

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --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 Trainings-Script abgeschlossen ist, bereinigen Sie die Ressourcen.

1. Fahren Sie den Container auf allen Workern herunter.

   gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
   --zone=$ZONE \
   --command="docker kill your-container-name"
   

2. Löschen Sie die TPU-VM.

   gcloud compute tpus tpu-vm delete $TPU_NAME \
   --zone=$ZONE
   

Nächste Schritte

• Cloud TPU-Anleitungen
• TPUs verwalten
• Cloud TPU-Architektur
• JAX-Code auf TPU-Slices ausführen