PyTorch-Code auf TPU-Slices ausführen

Bevor Sie die Befehle in diesem Dokument ausführen können, müssen Sie der Anleitung in Konto und Cloud TPU-Projekt einrichten gefolgt sein.

Sobald der PyTorch-Code auf einer einzelnen TPU-VM ausgeführt wird, können Sie den Code hochskalieren, indem Sie ihn auf einem TPU-Slice ausführen. TPU-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-Slices.

Cloud TPU-Slice erstellen

Definieren Sie einige Umgebungsvariablen, um die Verwendung der Befehle zu erleichtern.

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=europe-west4-b
export ACCELERATOR_TYPE=v5p-32
export RUNTIME_VERSION=v2-alpha-tpuv5

Beschreibungen von Umgebungsvariablen

Variable	Beschreibung
`PROJECT_ID`	Ihre Projekt-ID von Google Cloud . Verwenden Sie ein vorhandenes oder erstellen Sie ein neues Projekt.
`TPU_NAME`	Der Name der TPU.
`ZONE`	Die Zone, in der die TPU-VM erstellt werden soll. Weitere Informationen zu unterstützten Zonen finden Sie unter TPU-Regionen und -Zonen.
`ACCELERATOR_TYPE`	Der Beschleunigertyp gibt die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Weitere Informationen zu den unterstützten Beschleunigertypen für die einzelnen TPU-Versionen finden Sie unter TPU-Versionen.
`RUNTIME_VERSION`	Die Softwareversion der Cloud TPU.

Erstellen Sie Ihre TPU-VM mit dem folgenden Befehl:

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

PyTorch/XLA auf Slice installieren

Nachdem Sie den TPU-Slice erstellt haben, müssen Sie PyTorch auf allen Hosts im TPU-Slice installieren. Dazu können Sie den Befehl gcloud compute tpus tpu-vm ssh mit den Parametern --worker=all und --commamnd verwenden.

Wenn die folgenden Befehle aufgrund eines SSH-Verbindungsfehlers fehlschlagen, liegt das möglicherweise daran, dass die TPU-VMs keine externen IP-Adressen haben. Wenn Sie ohne externe IP-Adresse auf eine TPU-VM zugreifen möchten, folgen Sie der Anleitung unter Verbindung zu einer TPU-VM ohne öffentliche IP-Adresse herstellen.

Installieren Sie PyTorch/XLA auf allen TPU-VM-Workern:

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"

Klonen Sie XLA auf allen TPU-VM-Workern:

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

Trainings-Script auf TPU-Slice ausführen

Führen Sie das Trainingsskript auf allen Workern aus. Im Trainingsskript wird eine SPMD-Fragmentierungsstrategie (Single Program Multiple Data) verwendet. 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 15 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 Cloud TPU-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-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 tpu-vm 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}
```