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

  1. 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.

  2. 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.

  1. 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"

  2. 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.

  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}