PyTorch-Code auf TPU-Pod-Slices ausführen

PyTorch/XLA erfordert, dass alle TPU-VMs auf den Modellcode und die -Daten zugreifen können. Sie können ein Startskript verwenden, um die Software herunterzuladen, die zum Verteilen der Modelldaten auf alle TPU-VMs erforderlich ist.

Wenn Sie Ihre TPU-VMs mit einer Virtual Private Cloud (VPC) verbinden, müssen Sie in Ihrem Projekt eine Firewallregel hinzufügen, um den Eingang für die Ports 8470 bis 8479 zuzulassen. Weitere Informationen zum Hinzufügen von Firewallregeln finden Sie unter Firewallregeln verwenden.

Umgebung einrichten

1. Führen Sie in der Cloud Shell den folgenden Befehl aus, um sicherzustellen, dass Sie die aktuelle Version von gcloud ausführen:

   $ gcloud components update

   Wenn Sie gcloud installieren möchten, verwenden Sie den folgenden Befehl:

   $ sudo apt install -y google-cloud-sdk

2. Erstellen Sie einige Umgebungsvariablen:

   $ export PROJECT_ID=project-id
   $ export TPU_NAME=tpu-name
   $ export ZONE=us-central2-b
   $ export RUNTIME_VERSION=tpu-ubuntu2204-base
   $ export ACCELERATOR_TYPE=v4-32

TPU-VM erstellen

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

Trainingsscript konfigurieren und ausführen

1. Fügen Sie Ihrem Projekt Ihr SSH-Zertifikat hinzu:

   ssh-add ~/.ssh/google_compute_engine

2. PyTorch/XLA auf allen TPU-VM-Arbeitsstationen installieren

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
     --zone=${ZONE} \
     --project=${PROJECT_ID} \
     --worker=all --command="
     pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html"

3. XLA auf allen TPU VM-Arbeitsstationen klonen

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
     --zone=${ZONE} \
     --project=${PROJECT_ID} \
     --worker=all --command="git clone -b r2.5 https://github.com/pytorch/xla.git"

4. Trainingsskript auf allen Workern ausführen

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
     --zone=${ZONE} \
     --project=${PROJECT_ID} \
     --worker=all \
     --command="PJRT_DEVICE=TPU python3 ~/xla/test/test_train_mp_imagenet.py  \
     --fake_data \
     --model=resnet50  \
     --num_epochs=1 2>&1 | tee ~/logs.txt"
     

   Das Training dauert etwa 5 Minuten. Nach Abschluss des Vorgangs sollte eine Meldung wie die folgende angezeigt werden:

   Epoch 1 test end 23:49:15, Accuracy=100.00
   10.164.0.11 [0] Max Accuracy: 100.00%
   

Bereinigen

Wenn Sie mit Ihrer TPU-VM fertig sind, führen Sie die folgenden Schritte aus, um Ihre Ressourcen zu bereinigen.

1. Trennen Sie die Verbindung zur Compute Engine-Instanz, sofern noch nicht geschehen:

   (vm)$ exit

   Die Eingabeaufforderung sollte nun username@projectname lauten und angeben, dass Sie sich in Cloud Shell befinden.

2. Löschen Sie Ihre Cloud TPU- und Compute Engine-Ressourcen.

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

3. Prüfen Sie, ob die Ressourcen gelöscht wurden. Führen Sie dazu gcloud compute tpus execution-groups list aus. Der Löschvorgang kann einige Minuten dauern. Die Ausgabe des folgenden Befehls sollte keine der in dieser Anleitung erstellten Ressourcen enthalten:

   $ gcloud compute tpus tpu-vm list --zone=${ZONE}