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

Docker-Container erleichtern die Konfiguration von Anwendungen, da Code und alle erforderlichen Abhängigkeiten in einem distribuierbaren Paket kombiniert 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.

TensorFlow-Modell in einem Docker-Container trainieren

TPU-Gerät

  1. Erstellen Sie im 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

    gcloud storage buckets create gs://your-bucket-name --location=europe-west4
  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.18.0-pjrt
  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
  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"

Bereinigen Sie die Ressourcen nach Abschluss des Trainingsskripts.

  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 im 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.18.0-pod-pjrt
  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=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"

Bereinigen Sie die Ressourcen nach Abschluss des Trainingsskripts.

  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 zur TPU-VM herstellen

    gcloud compute tpus tpu-vm ssh your-tpu-name \
    --zone=europe-west4-a
  3. Starten Sie einen Container in der TPU-VM 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
  4. TPU-Laufzeit konfigurieren

    Es gibt zwei PyTorch/XLA-Laufzeitoptionen: PJRT und XRT. Wir empfehlen, PJRT zu verwenden, es sei denn, Sie haben einen Grund, XRT zu verwenden. Weitere Informationen zu den verschiedenen Laufzeitkonfigurationen finden Sie unter „Gründe für die Verwendung von XRT“. 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

Bereinigen Sie die Ressourcen nach Abschluss des Trainingsskripts.

  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 ihn gleichzeitig auf allen TPU-Workern ausführen. Dazu können Sie den Befehl gcloud compute tpus tpu-vm ssh mit den Flags --worker=all und --command verwenden. Im Folgenden wird beschrieben, 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. Den 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 entfernen Sie den Container, nachdem der Prozess beendet wurde.
    • --privileged stellt das TPU-Gerät dem Container zur Verfügung.
    • --net=host bindet alle Ports des Containers an die TPU-VM, um die Kommunikation zwischen den Hosts im Pod zu ermöglichen.
    • -e Umgebungsvariablen festlegen.

Bereinigen Sie die Ressourcen nach Abschluss des Trainingsskripts.

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 zur TPU-VM herstellen

    gcloud compute tpus tpu-vm ssh your-tpu-name  --zone=europe-west4-a
  3. Docker-Daemon in der 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

Bereinigen Sie die Ressourcen nach Abschluss des Trainingsskripts.

  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 gleichzeitig auf allen TPU-Workern ausführen. Dazu können Sie den Befehl gcloud compute tpus tpu-vm ssh mit den Flags --worker=all und --command verwenden. Im Folgenden wird beschrieben, wie Sie ein Docker-Image erstellen, um die Einrichtung der einzelnen TPU-Arbeitsstationen zu vereinfachen.

  1. Erstellen Sie im 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 Ihrem Docker-Image ein Tag hinzu, bevor Sie es in Artifact Registry veröffentlichen. Weitere Informationen zur Arbeit 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 in Artifact Registry veröffentlichen

    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-Arbeitsstationen aus der 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"

Bereinigen Sie die Ressourcen nach Abschluss des Trainingsskripts.

  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

JAX-Modell in einem Docker-Container mit JAX Stable Stack trainieren

Sie können die Docker-Images MaxText und MaxDiffusion mit dem Basis-Image JAX Stable Stack erstellen.

JAX Stable Stack bietet eine einheitliche Umgebung für MaxText und MaxDiffusion, da JAX mit Kernpaketen wie orbax, flax, optax und libtpu.so gebündelt wird. Diese Bibliotheken werden auf Kompatibilität getestet und bieten eine stabile Grundlage für die Erstellung und Ausführung von MaxText und MaxDiffusion. So werden potenzielle Konflikte aufgrund von inkompatiblen Paketversionen vermieden.

Der JAX Stable Stack enthält eine vollständig veröffentlichte und qualifizierte libtpu.so, die Kernbibliothek, die die TPU-Programmkompilierung, ‑ausführung und ‑Netzwerkkonfiguration steuert. Die libtpu-Version ersetzt den bisher von JAX verwendeten Nightly-Build und sorgt mit Qualifikationstests auf PJRT-Ebene in HLO/StableHLO-IRs für eine konsistente Funktion von XLA-Berechnungen auf TPUs.

Wenn Sie das Docker-Image für MaxText und MaxDiffusion mit dem JAX Stable Stack erstellen möchten, legen Sie beim Ausführen des docker_build_dependency_image.sh-Scripts die Variable MODE auf stable_stack und die Variable BASEIMAGE auf das gewünschte Basis-Image fest.

Im folgenden Beispiel wird us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.4.35-rev1 als Basisbild angegeben:

   bash docker_build_dependency_image.sh MODE=stable_stack BASEIMAGE=us-docker.pkg.dev/cloud-tpu-images/jax-stable-stack/tpu:jax0.4.35-rev1
   

Eine Liste der verfügbaren JAX Stable Stack-Basis-Images finden Sie unter JAX Stable Stack-Images in Artifact Registry.

Nächste Schritte