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
TPU-Gerät
Erstellen Sie eine Cloud TPU-VM.
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-ubuntu2204-base
Stellen Sie eine SSH-Verbindung zur TPU-VM her.
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=europe-west4-a
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.
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
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"
Klonen Sie das PyTorch XLA-Repository.
git clone --recursive https://github.com/pytorch/xla.git
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.
- Geben Sie
exit
ein, um den Docker-Container zu beenden. - Geben Sie
exit
ein, um die TPU-VM zu beenden. Löschen Sie die TPU-VM.
gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
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.
Erstellen Sie eine TPU-VM.
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=us-central2-b \ --accelerator-type=v4-32 \ --version=tpu-ubuntu2204-base
Fügen Sie den aktuellen Nutzer der Docker-Gruppe hinzu.
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=us-central2-b \ --worker=all \ --command='sudo usermod -a -G docker $USER'
Klonen Sie das PyTorch XLA-Repository.
gcloud compute tpus tpu-vm ssh your-tpu-name --worker=all \ --zone=us-central2-b \ --command="git clone --recursive https://github.com/pytorch/xla.git"
Führen Sie das Trainings-Script 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 -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.
- Mit
Wenn das Trainings-Script 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
Erstellen Sie eine TPU-VM.
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-ubuntu2204-base
Stellen Sie eine SSH-Verbindung zur TPU-VM her.
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Starten Sie den Docker-Daemon in der TPU-VM.
sudo systemctl start docker
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
Installieren Sie JAX.
pip install jax[tpu]
Installieren Sie FLAX.
pip install --upgrade clu git clone https://github.com/google/flax.git pip install --user -e flax
Installieren Sie das
tensorflow
- und dastensorflow-dataset
-Paket.pip install tensorflow pip install tensorflow-datasets
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, sollten Sie die Ressourcen bereinigen.
- Geben Sie
exit
ein, um den Docker-Container zu beenden. - Geben Sie
exit
ein, um die TPU-VM zu beenden. Löschen Sie die TPU-VM.
gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
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.
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
Bereiten Sie Artifact Registry vor.
gcloud artifacts repositories create your-repo \ --repository-format=docker \ --location=europe-west4 --description="Docker repository" \ --project=your-project gcloud artifacts repositories list \ --project=your-project gcloud auth configure-docker europe-west4-docker.pkg.dev
Erstellen Sie das Docker-Image.
docker build -t your-image-name .
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/your-project/your-repo/your-image-name:your-tag
Übertragen Sie das Docker-Image per Push an Artifact Registry.
docker push europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
Erstellen Sie eine TPU-VM.
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-ubuntu2204-base
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"
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-name:your-tag bash"
Führen Sie das Trainings-Script 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 Trainings-Script abgeschlossen ist, sollten Sie die Ressourcen bereinigen.
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"
Löschen Sie die TPU-VM.
gcloud compute tpus tpu-vm delete your-tpu-name \ --zone=europe-west4-a