PyTorch-Code auf TPU-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-Speicherblock 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-Speichereinheit 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 Slice installieren

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

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

  1. Installieren Sie PyTorch/XLA auf allen TPU-VM-Arbeitsstationen:

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

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

  2. Löschen Sie Ihre Cloud TPU-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 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}