PyTorch-Code auf TPU-Pod-Slices ausführen
Bevor Sie die Befehle in diesem Dokument ausführen, müssen Sie den Anleitungen in Konto und Cloud TPU-Projekt einrichten gefolgt sein.
Nachdem Sie den PyTorch-Code auf einer einzelnen TPU-VM ausgeführt haben, können Sie den Code skalieren, indem Sie ihn auf einem TPU-Pod-Slice ausführen. TPU Pod-Slices sind mehrere TPU-Boards, die über dedizierte Hochgeschwindigkeits-Netzwerkverbindungen miteinander verbunden sind. Dieses Dokument bietet eine Einführung zum Ausführen von PyTorch-Code auf TPU-Pod-Slices.
Cloud TPU-Pod-Slice erstellen
Definieren Sie einige Umgebungsvariablen, um die Befehle nutzerfreundlicher zu gestalten.
export PROJECT_ID=your-project export ACCELERATOR_TYPE=v5p-32 export ZONE=europe-west4-b export RUNTIME_VERSION=v2-alpha-tpuv5 export TPU_NAME=your-tpu-name
Beschreibungen von Umgebungsvariablen
PROJECT_ID
- Ihre Google Cloud Projekt-ID.
ACCELERATOR_TYPE
- Mit dem Beschleunigertyp geben Sie die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Weitere Informationen zu den unterstützten Beschleunigertypen für jede TPU-Version finden Sie unter TPU-Versionen.
ZONE
- Die Zone, in der Sie die Cloud TPU erstellen möchten.
RUNTIME_VERSION
- Die Version der Cloud TPU-Software
TPU_NAME
- Der vom Nutzer zugewiesene Name für Ihre Cloud TPU.
Führen Sie den folgenden Befehl aus, um die TPU-VM zu erstellen:
$ gcloud compute tpus tpu-vm create $TPU_NAME \ --zone=$ZONE \ --project=$PROJECT_ID \ --accelerator-type=$ACCELERATOR_TYPE \ --version=$RUNTIME_VERSION
PyTorch/XLA auf Ihrem Pod-Slice installieren
Nachdem Sie das TPU Pod-Slice erstellt haben, müssen Sie PyTorch auf allen Hosts im TPU Pod-Slice installieren. Verwenden Sie dazu den Befehl gcloud compute tpus tpu-vm ssh
mit den Parametern --worker=all
und --commamnd
.
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"
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 https://github.com/pytorch/xla.git"
Trainingsskript auf Ihrem TPU-Pod-Speicherplatz ausführen
Führen Sie das Trainingsskript auf allen Workern aus. Das Trainingsskript verwendet die Sharding-Strategie „Single Program Multiple Data“ (SPMD). Weitere Informationen zu SPMD finden Sie im PyTorch/XLA SPMD-Nutzerhandbuch.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --zone=${ZONE} \ --project=${PROJECT_ID} \ --worker=all \ --command="PJRT_DEVICE=TPU python3 ~/xla/test/spmd/test_train_spmd_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.
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.Löschen Sie Ihre Cloud TPU- und Compute Engine-Ressourcen.
$ gcloud compute tpus tpu-vm delete \ --zone=${ZONE}
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}