Cloud TPU-Anwendungen in einem Docker-Container ausführen
Docker-Container erleichtern die Anwendungskonfiguration, da der Code und alle erforderlichen Abhängigkeiten in einem verteilbaren Paket zusammengefasst werden. 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 für jedes von Cloud TPU unterstützte ML-Framework einen Docker-Container einrichten.
TensorFlow-Modell in einem Docker-Container trainieren
TPU-Gerät
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
Cloud Storage-Bucket erstellen
gsutil mb -c standard -l europe-west4 gs://your-bucket-name
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.15.0-pjrt
Dockerfile in die TPU-VM kopieren
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
Stellen Sie eine SSH-Verbindung zur TPU-VM her:
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=europe-west4-a
Docker-Image erstellen
sudo docker build -t your-image-name .
Docker-Container starten
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
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
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 nach Abschluss des Trainingsskripts die Ressourcen.
- Geben Sie
exit
ein, um den Docker-Container zu beenden - Geben Sie
exit
ein, um die TPU-VM zu beenden - TPU-VM löschen
$ gcloud compute tpus tpu-vm delete your-tpu-name --zone=europe-west4-a
TPU-Pod
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
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.15.0-pod-pjrt
Dockerfile in die TPU-VM kopieren
gcloud compute tpus tpu-vm scp ./Dockerfile your-tpu-name:
Stellen Sie eine SSH-Verbindung zur TPU-VM her:
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=europe-west4-a
Docker-Image erstellen
sudo docker build -t your-image-name .
Docker-Container starten
sudo docker run -ti --rm --net=host --name your-container-name --privileged your-image-name bash
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 nach Abschluss des Trainingsskripts die Ressourcen.
- Geben Sie
exit
ein, um den Docker-Container zu beenden - Geben Sie
exit
ein, um die TPU-VM zu beenden - 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
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
SSH-Verbindung zu TPU-VM herstellen
gcloud compute tpus tpu-vm ssh your-tpu-name \ --zone=europe-west4-a
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
TPU-Laufzeit konfigurieren
Es gibt zwei PyTorch/XLA-Laufzeitoptionen: PJRT und XRT. Wir empfehlen die Verwendung von PJRT, sofern Sie keinen Grund für die Verwendung von XRT haben. Wenn Sie mehr über die verschiedenen Laufzeitkonfigurationen erfahren möchten, haben Sie einen Grund für die Verwendung von XRT. Weitere Informationen zu den verschiedenen Laufzeitkonfigurationen finden Sie in der PJRT-Laufzeitdokumentation.
PJRT
export PJRT_DEVICE=TPU
Logo: XRT
export XRT_TPU_CONFIG="localservice;0;localhost:51011"
PyTorch XLA-Repository klonen
git clone --recursive https://github.com/pytorch/xla.git
ResNet50 trainieren
python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
Bereinigen Sie nach Abschluss des Trainingsskripts die Ressourcen.
- Geben Sie
exit
ein, um den Docker-Container zu beenden - Geben Sie
exit
ein, um die TPU-VM zu beenden - 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, muss er auf allen TPU-Workern gleichzeitig ausgeführt werden. Eine Möglichkeit dazu ist die Verwendung des Befehls gcloud compute tpus tpu-vm ssh
mit den Flags --worker=all
und --command
. Das folgende Verfahren zeigt, wie Sie ein Docker-Image erstellen, um die Einrichtung der TPU-Worker zu vereinfachen.
TPU-VM erstellen
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=us-central2-b \ --accelerator-type=v4-32 \ --version=tpu-ubuntu2204-base
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"
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
macht das TPU-Gerät für den Container verfügbar.--net=host
bindet alle Ports des Containers an die TPU-VM, um die Kommunikation zwischen den Hosts im Pod zu ermöglichen.-e
hat Umgebungsvariablen festgelegt.
Bereinigen Sie nach Abschluss des Trainingsskripts die Ressourcen.
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
TPU-VM erstellen
gcloud compute tpus tpu-vm create your-tpu-name \ --zone=europe-west4-a \ --accelerator-type=v2-8 \ --version=tpu-ubuntu2204-base
SSH-Verbindung zu TPU-VM herstellen
gcloud compute tpus tpu-vm ssh your-tpu-name --zone=europe-west4-a
Docker-Daemon in TPU-VM starten
sudo systemctl start docker
Docker-Container starten
sudo docker run -ti --rm --name your-container-name --privileged --network=host python:3.8 bash
JAX installieren
pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
FLAX installieren
pip install --upgrade clu git clone https://github.com/google/flax.git pip install --user -e flax
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 nach Abschluss des Trainingsskripts die Ressourcen.
- Geben Sie
exit
ein, um den Docker-Container zu beenden - Geben Sie
exit
ein, um die TPU-VM zu beenden 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 auf allen TPU-Workern gleichzeitig ausgeführt werden. Eine Möglichkeit dazu ist die Verwendung des Befehls gcloud compute tpus tpu-vm ssh
mit den Flags --worker=all
und --command
. Das folgende Verfahren zeigt, wie Sie ein Docker-Image erstellen, um die Einrichtung der einzelnen TPU-Worker zu vereinfachen.
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
Docker-Image erstellen
docker build -t your-image-name .
Fügen Sie Ihrem Docker-Image ein Tag hinzu, bevor Sie es per Push an Artifact Registry übertragen. 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
Docker-Image per Push in Artifact Registry hochladen
docker push europe-west4-docker.pkg.dev/your-project/your-repo/your-image-name:your-tag
TPU-VM erstellen
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:your-tag bash"
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 nach Abschluss des Trainingsskripts die Ressourcen.
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 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-Segmenten ausführen
- JAX-Code auf TPU-Pod-Segmenten ausführen