TPU-Arbeitslasten in einem Docker-Container ausführen

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.

PyTorch-Modell in einem Docker-Container trainieren

Bevor Sie die folgenden Befehle ausführen, erstellen Sie 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 Cloud 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. Achten Sie darauf, dass Ihrem Google Cloud Nutzer die Rolle „Artifact Registry-Leser“ zugewiesen wurde. Weitere Informationen finden Sie unter Gewähren von Artifact Registry-Rollen.

  4. Starten Sie einen Container in der TPU-VM mit dem Nightly-PyTorch/XLA-Image.

    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. Konfigurieren Sie die TPU-Laufzeit.

    Es gibt zwei PyTorch/XLA-Laufzeitoptionen: PJRT und XRT. Wir empfehlen, PJRT zu verwenden, wenn kein Grund für die Verwendung von XRT vorliegt. 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"

  6. Klonen Sie das PyTorch XLA-Repository.

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

  7. Trainieren Sie ResNet50.

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

Wenn das Trainings-Script 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. Löschen Sie die TPU-VM.

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

TPU-Slice

Wenn Sie PyTorch-Code auf einem TPU-Slice ausführen, müssen Sie den 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 eine TPU-VM.

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

  2. Fügen Sie den aktuellen Nutzer der Docker-Gruppe hinzu.

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

  3. Klonen Sie das PyTorch XLA-Repository.

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
--zone=$ZONE \
--command="git clone --recursive https://github.com/pytorch/xla.git"

  4. Führen Sie das Trainings-Script in einem Container auf allen TPU-Workern aus.

    gcloud compute tpus tpu-vm ssh $TPU_NAME --worker=all \
--zone=$ZONE \
--command="docker run --rm --privileged --net=host  -v ~/xla:/xla -e PJRT_DEVICE=TPU us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/xla:r2.6.0_3.10_tpuvm_cxx11 python /xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1"

    Docker-Befehls-Flags:

    • Mit --rm wird der Container entfernt, nachdem sein Prozess beendet wurde.
    • Mit --privileged wird das TPU-Gerät für den Container verfügbar gemacht.
    • Mit --net=host werden alle Ports des Containers an die TPU-VM gebunden, um die Kommunikation zwischen den Hosts im Pod zu ermöglichen.
    • Mit -e werden Umgebungsvariablen festgelegt.

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 $TPU_NAME \
--zone=$ZONE

JAX-Modell in einem Docker-Container trainieren

Bevor Sie die folgenden Befehle ausführen, erstellen Sie 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 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. 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, sollten Sie die Ressourcen bereinigen.

  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