PyTorch-Modelle auf Cloud TPU Pods trainieren

---

In dieser Anleitung wird gezeigt, wie Sie das Training Ihres Modells von einer einzelnen Cloud TPU (v2-8 oder v3-8) auf einen Cloud TPU Pod mit der TPU-Knotenkonfiguration hochskalieren. Cloud TPU-Beschleuniger in einem TPU-Pod sind über Interconnect-Verbindungen mit hoher Bandbreite verbunden, sodass sie Trainingsjobs effizient hochskalieren können.

Wichtig: Sie können TPUs mit zwei verschiedenen Architekturen verwenden: TPU-Knoten und TPU-VMs. In dieser Anleitung wird davon ausgegangen, dass Sie TPU-Knoten verwenden. Die beiden VM-Architekturen werden unter Systemarchitektur beschrieben. Wenn Sie nicht wissen, dass Sie TPU-Knoten benötigen, empfehlen wir TPU-VMs. Weitere Informationen zur Verwendung von PyTorch auf TPU-Pods mit TPU-VMs finden Sie im Nutzerhandbuch für PyTorch XLA.

Weitere Informationen zu den Cloud TPU Pods finden Sie auf der Cloud TPU-Produktseite oder in dieser Cloud TPU-Präsentation.

Das folgende Diagramm bietet eine Übersicht über die Einrichtung des verteilten Clusters. Eine Instanzgruppe von VMs ist mit einem TPU-Pod verbunden. Für jede Gruppe von 8 TPU-Kernen ist eine VM erforderlich. Die VMs speisen Daten zu den TPU-Kernen und alle Trainings findet auf dem TPU Pod statt.

Hinweis: Wenn Sie bereitgestellte NFS für die Datenspeicherung verwenden möchten, müssen Sie OS Login für alle TPU-VMs im Pod-Slice festlegen. Weitere Informationen finden Sie unter NFS als Datenspeicher verwenden.

Lernziele

• Eine Compute Engine-Instanzgruppe und einen Cloud TPU Pod für das Training mit PyTorch/XLA einrichten
• PyTorch/XLA-Training auf einem Cloud TPU Pod ausführen

Warnung: In diesem Modell wird ein Dataset eines Drittanbieters verwendet. Google gibt keine Zusicherungen, Gewährleistungen oder Garantien hinsichtlich der Gültigkeit oder anderer Aspekte dieses Datasets.

Hinweis

Bevor Sie das verteilte Training auf Cloud TPU Pods starten, achten Sie darauf, dass Ihr Modell auf einem einzelnen Cloud TPU-Gerät (v2-8 oder v3-8) trainiert wird. Wenn Ihr Modell erhebliche Leistungsprobleme auf einem einzelnen Gerät hat, lesen Sie die Leitfäden Best Practices und Fehlerbehebung.

Sobald Ihr einzelnes TPU-Gerät erfolgreich trainiert wurde, führen Sie die folgenden Schritte aus, um einen Cloud TPU Pod einzurichten und zu trainieren:

1. Konfigurieren Sie den Befehl gcloud.

2. [Optional] VM-Laufwerk-Image in ein VM-Image aufnehmen

3. Instanzvorlage aus einem VM-Image erstellen

4. Instanzgruppe aus Ihrer Instanzvorlage erstellen

5. SSH-Verbindung zu Ihrer Compute Engine-VM herstellen

6. Anhand der Firewallregeln prüfen, ob die Kommunikation zwischen VMs zugelassen wird

7. Cloud TPU Pod erstellen

8. Verteiltes Training auf dem Pod ausführen

9. Bereinigen

gcloud-Befehl konfigurieren

Konfigurieren Sie Ihr Google Cloud-Projekt mit gcloud:

Erstellen Sie eine Variable für Ihre Projekt-ID.

export PROJECT_ID=project-id

Legen Sie Ihre Projekt-ID in gcloud als Standardprojekt fest.

gcloud config set project ${PROJECT_ID}

Wenn Sie diesen Befehl zum ersten Mal in einer neuen Cloud Shell-VM ausführen, wird die Seite Authorize Cloud Shell angezeigt. Klicken Sie unten auf der Seite auf Authorize, damit gcloud API-Aufrufe mit Ihren Anmeldedaten ausführen kann.

Konfigurieren Sie die Standardzone mit gcloud:

gcloud config set compute/zone europe-west4-a

[Optional] Erstellen Sie ein VM-Laufwerk-Image

Hinweis: Sie können diesen Schritt überspringen, wenn Sie keine zusätzlichen Pakete oder Datasets für das Training benötigen und ein standardmäßiges PyTorch-/XLA-Image verwenden können.

Sie können das Laufwerk-Image von der VM verwenden, die Sie zum Trainieren der einzelnen TPU verwendet haben (auf der das Dataset, die Pakete usw. bereits installiert sind). Bevor Sie ein Image erstellen, beenden Sie die VM mit dem Befehl gcloud:

gcloud compute instances stop vm-name

Als Nächstes erstellen Sie ein VM-Image mit dem Befehl gcloud:

gcloud compute images create image-name  \
    --source-disk instance-name \
    --source-disk-zone europe-west4-a \
    --family=torch-xla \
    --storage-location europe-west4

Instanzvorlage aus einem VM-Image erstellen

Erstellen Sie eine Standardinstanzvorlage. Wenn Sie eine Instanzvorlage erstellen, können Sie das im obigen Schritt erstellte VM-Image verwenden. ODER Sie können das von Google bereitgestellte öffentliche PyTorch/XLA-Image verwenden. Verwenden Sie zum Erstellen einer Instanzvorlage den Befehl gcloud:

gcloud compute instance-templates create instance-template-name \
    --machine-type n1-standard-16 \
    --image-project=${PROJECT_ID} \
    --image=image-name \
    --scopes=https://www.googleapis.com/auth/cloud-platform

Instanzgruppe aus Ihrer Instanzvorlage erstellen

gcloud compute instance-groups managed create instance-group-name \
    --size 4 \
    --template template-name \
    --zone europe-west4-a

SSH-Verbindung zu Ihrer Compute Engine-VM herstellen

Nachdem Sie die Instanzgruppe erstellt haben, stellen Sie eine SSH-Verbindung zu einer der Instanzen (VMs) in Ihrer Instanzgruppe her. Verwenden Sie den folgenden Befehl, um alle Instanzen in der Instanz aufzulisten, die den Befehl gcloud gruppieren:

gcloud compute instance-groups list-instances instance-group-name

Stellen Sie eine SSH-Verbindung zu einer der Instanzen her, die über den Befehl list-instances aufgelistet werden.

gcloud compute ssh instance-name --zone=europe-west4-a

Prüfen, ob die VMs in der Instanzgruppe miteinander kommunizieren können

Prüfen Sie mit dem Befehl nmap, ob die VMs in Ihrer Instanzgruppe miteinander kommunizieren können. Führen Sie den Befehl nmap von der VM aus, mit der Sie verbunden sind, und ersetzen Sie instance-name durch den Instanznamen einer anderen VM in der Instanzgruppe.

(vm)$ nmap -Pn -p 8477 instance-name

Starting Nmap 7.40 ( https://nmap.org ) at 2019-10-02 21:35 UTC
Nmap scan report for pytorch-20190923-n4tx.c.jysohntpu.internal (10.164.0.3)
Host is up (0.00034s latency).
PORT     STATE  SERVICE
8477/tcp closed unknown

Wenn das Feld STATE (Status) nicht den Wert Filtered (Gefiltert) enthält, sind die Firewallregeln richtig eingerichtet.

Cloud TPU Pod erstellen

gcloud compute tpus create tpu-name \
    --zone=europe-west4-a \
    --network=default \
    --accelerator-type=v2-32 \
    --version=pytorch-1.13

Verteiltes Training auf dem Pod ausführen

Hinweis: In diesem Beispiel wird davon ausgegangen, dass Sie eine Conda-Umgebung für verteiltes Training verwenden. Außerdem aktiviert dieses Beispiel XLA_USE_BF16=1 zum Trainingszeitpunkt. Sie können auch die Variable --env verwenden, um alle anderen Umgebungsvariablen aufzulisten, die verteilt werden sollen.

1. Exportieren Sie in Ihrem VM-Sitzungsfenster den Cloud TPU-Namen und aktivieren Sie die Conda-Umgebung.

   (vm)$ export TPU_NAME=tpu-name

   (vm)$ conda activate torch-xla-1.13

2. Führen Sie das Trainingsskript aus:

   Hinweis: Die Conda-Umgebung muss der TPU-Softwareversion entsprechen. TPU-Software pytorch-1.13 entspricht der Conda-Umgebung torch-xla-1.13 und TPU-Software pytorch-nightly entspricht der Conda-Umgebung torch-xla-nightly.

   (torch-xla-1.13)$ python -m torch_xla.distributed.xla_dist \
         --tpu=$TPU_NAME \
         --conda-env=torch-xla-1.13 \
         --env XLA_USE_BF16=1 \
         --env ANY_OTHER=ENV_VAR \
         -- python /usr/share/torch-xla-1.13/pytorch/xla/test/test_train_mp_imagenet.py \
         --fake_data

Wenn Sie den obigen Befehl ausführen, wird in etwa die folgende Ausgabe angezeigt. Beachten Sie, dass hier --fake_data verwendet wird. Das Training dauert etwa eine halbe Stunde auf einem TPU Pod vom Typ v3-32.

2020-08-06 02:38:29  [] Command to distribute: "python" "/usr/share/torch-xla-nightly/pytorch/xla/test/test_train_mp_imagenet.py" "--fake_data"
2020-08-06 02:38:29  [] Cluster configuration: {client_workers: [{10.164.0.43, n1-standard-96, europe-west4-a, my-instance-group-hm88}, {10.164.0.109, n1-standard-96, europe-west4-a, my-instance-group-n3q2}, {10.164.0.46, n1-standard-96, europe-west4-a, my-instance-group-s0xl}, {10.164.0.49, n1-standard-96, europe-west4-a, my-instance-group-zp14}], service_workers: [{10.131.144.61, 8470, v3-32, europe-west4-a, pytorch-nightly, my-tpu-slice}, {10.131.144.59, 8470, v3-32, europe-west4-a, pytorch-nightly, my-tpu-slice}, {10.131.144.58, 8470, v3-32, europe-west4-a, pytorch-nightly, my-tpu-slice}, {10.131.144.60, 8470, v3-32, europe-west4-a, pytorch-nightly, my-tpu-slice}]}
2020-08-06 02:38:31 10.164.0.43 [0]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:31 10.164.0.43 [0]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2757      0 --:--:-- --:--:-- --:--:--  3166
2020-08-06 02:38:34 10.164.0.43 [0]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:34 10.164.0.43 [0]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2623      0 --:--:-- --:--:-- --:--:--  2714
2020-08-06 02:38:37 10.164.0.46 [2]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:37 10.164.0.46 [2]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2583      0 --:--:-- --:--:-- --:--:--  2714
2020-08-06 02:38:37 10.164.0.49 [3]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:37 10.164.0.49 [3]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2530      0 --:--:-- --:--:-- --:--:--  2714
2020-08-06 02:38:37 10.164.0.109 [1]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:37 10.164.0.109 [1]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2317      0 --:--:-- --:--:-- --:--:--  2375
2020-08-06 02:38:40 10.164.0.46 [2]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:40 10.164.0.49 [3]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:40 10.164.0.46 [2]                                  Dload  Upload   Total   Spent    Left  Speed
2020-08-06 02:38:40 10.164.0.49 [3]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2748      0 --:--:-- --:--:-- --:--:--  3166
100    19  100    19    0     0   2584      0 --:--:-- --:--:-- --:--:--  2714
2020-08-06 02:38:40 10.164.0.109 [1]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:40 10.164.0.109 [1]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2495      0 --:--:-- --:--:-- --:--:--  2714
2020-08-06 02:38:43 10.164.0.49 [3]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:43 10.164.0.49 [3]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2654      0 --:--:-- --:--:-- --:--:--  2714
2020-08-06 02:38:43 10.164.0.43 [0]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:43 10.164.0.43 [0]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2784      0 --:--:-- --:--:-- --:--:--  3166
2020-08-06 02:38:43 10.164.0.46 [2]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:43 10.164.0.46 [2]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2691      0 --:--:-- --:--:-- --:--:--  3166
2020-08-06 02:38:43 10.164.0.109 [1]   % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current
2020-08-06 02:38:43 10.164.0.109 [1]                                  Dload  Upload   Total   Spent    Left  Speed
100    19  100    19    0     0   2589      0 --:--:-- --:--:-- --:--:--  2714
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/14 Epoch=1 Step=0 Loss=6.87500 Rate=258.47 GlobalRate=258.47 Time=02:38:57
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/15 Epoch=1 Step=0 Loss=6.87500 Rate=149.45 GlobalRate=149.45 Time=02:38:57
2020-08-06 02:38:57 10.164.0.43 [0] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.43 [0] Epoch 1 train begin 02:38:52
2020-08-06 02:38:57 10.164.0.43 [0] | Training Device=xla:1/0 Epoch=1 Step=0 Loss=6.87500 Rate=25.72 GlobalRate=25.72 Time=02:38:57
2020-08-06 02:38:57 10.164.0.43 [0] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.43 [0] | Training Device=xla:0/6 Epoch=1 Step=0 Loss=6.87500 Rate=89.01 GlobalRate=89.01 Time=02:38:57
2020-08-06 02:38:57 10.164.0.43 [0] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.43 [0] | Training Device=xla:0/1 Epoch=1 Step=0 Loss=6.87500 Rate=64.15 GlobalRate=64.15 Time=02:38:57
2020-08-06 02:38:57 10.164.0.43 [0] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.43 [0] | Training Device=xla:0/2 Epoch=1 Step=0 Loss=6.87500 Rate=93.19 GlobalRate=93.19 Time=02:38:57
2020-08-06 02:38:57 10.164.0.43 [0] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.43 [0] | Training Device=xla:0/7 Epoch=1 Step=0 Loss=6.87500 Rate=58.78 GlobalRate=58.78 Time=02:38:57
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] Epoch 1 train begin 02:38:56
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:1/8 Epoch=1 Step=0 Loss=6.87500 Rate=100.43 GlobalRate=100.43 Time=02:38:57
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/13 Epoch=1 Step=0 Loss=6.87500 Rate=66.83 GlobalRate=66.83 Time=02:38:57
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/11 Epoch=1 Step=0 Loss=6.87500 Rate=64.28 GlobalRate=64.28 Time=02:38:57
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/10 Epoch=1 Step=0 Loss=6.87500 Rate=73.17 GlobalRate=73.17 Time=02:38:57
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/9 Epoch=1 Step=0 Loss=6.87500 Rate=27.29 GlobalRate=27.29 Time=02:38:57
2020-08-06 02:38:57 10.164.0.109 [1] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.109 [1] | Training Device=xla:0/12 Epoch=1 Step=0 Loss=6.87500 Rate=110.29 GlobalRate=110.29 Time=02:38:57
2020-08-06 02:38:57 10.164.0.46 [2] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.46 [2] | Training Device=xla:0/20 Epoch=1 Step=0 Loss=6.87500 Rate=100.85 GlobalRate=100.85 Time=02:38:57
2020-08-06 02:38:57 10.164.0.46 [2] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.46 [2] | Training Device=xla:0/22 Epoch=1 Step=0 Loss=6.87500 Rate=93.52 GlobalRate=93.52 Time=02:38:57
2020-08-06 02:38:57 10.164.0.46 [2] ==> Preparing data..
2020-08-06 02:38:57 10.164.0.46 [2] | Training Device=xla:0/23 Epoch=1 Step=0 Loss=6.87500 Rate=165.86 GlobalRate=165.86 Time=02:38:57

Bereinigen

Damit Ihrem Google Cloud-Konto die in dieser Anleitung verwendeten Ressourcen nicht in Rechnung gestellt werden, löschen Sie entweder das Projekt, das die Ressourcen enthält, oder Sie behalten das Projekt und löschen die einzelnen Ressourcen.

1. Trennen Sie die Verbindung zur Compute Engine-VM:

   (vm)$ exit

2. Löschen Sie die Instanzgruppe:

   gcloud compute instance-groups managed delete instance-group-name

3. Löschen Sie Ihren TPU-Pod:

   gcloud compute tpus delete ${TPU_NAME} --zone=europe-west4-a

4. Löschen Sie die Instanzgruppenvorlage:

   gcloud compute instance-templates delete instance-template-name

5. [Optional] Löschen Sie das VM-Laufwerk-Image:

   gcloud compute images delete image-name

Nächste Schritte

Testen Sie die PyTorch Colabs:

• Erste Schritte mit PyTorch auf Cloud TPUs
• MNIST auf TPUs trainieren
• ResNet18 auf TPUs mit Cifar10-Dataset trainieren
• Inferenz mit dem vortrainierten ResNet50-Modell
• Schnelle neuronale Übertragung
• MultiCore Training AlexNet on Fashion MNIST
• Single Core Training AlexNet on Fashion MNIST