Nutzerhandbuch für Cloud TPU PyTorch/XLA

Wichtig: Sie können TPUs mit zwei verschiedenen Architekturen verwenden: TPU-Knoten und TPU-VMs. In dieser Anleitung wird davon ausgegangen, dass Sie TPU-VMs 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-Knoten finden Sie unter PyTorch-Modelle auf Cloud TPU-Pods trainieren.

ML-Arbeitslasten mit PyTorch/XLA ausführen

In diesem Leitfaden wird die einfache Berechnung auf einer v4-TPU mit PyTorch beschrieben.

Grundlegende Einstellungen

1. Erstellen Sie eine TPU-VM mit einer v4-TPU, die die TPU-VM-Laufzeit für Pytorch 2.0 ausführt:

     gcloud compute tpus tpu-vm create your-tpu-name \
     --zone=us-central2-b \
     --accelerator-type=v4-8 \
     --version=tpu-vm-v4-pt-2.0

   Hinweis: Wenn Sie eine andere Version von PyTorch installieren möchten, ersetzen Sie tpu-vm-v4-pt-2.0 durch die TPU-Laufzeitversion, die Sie verwenden möchten. Eine Liste der TPU-Laufzeiten mit PyTorch finden Sie unter Von PyTorch unterstützte Versionen.

2. Stellen Sie über SSH eine Verbindung zur TPU-VM her:

     gcloud compute tpus tpu-vm ssh your-tpu-name \
     --zone=us-central2-b \
     --accelerator-type=v4-8

   Wichtiger Hinweis: In diesem Leitfaden bedeutet das Präfix (vm) $, dass Sie den Befehl auf der TPU-VM-Instanz ausführen sollten.

3. Legen Sie die PJRT- oder XRT-TPU-Gerätekonfiguration fest.

   Hinweis: Es gibt zwei PyTorch-/XLA-Laufzeitoptionen: PJRT und XRT. Wir empfehlen die Verwendung von PJRT, es sei denn, Sie haben einen Grund für die Verwendung von XRT. Weitere Informationen zu den verschiedenen Laufzeitkonfigurationen sowie einen Grund für die Verwendung von XRT finden Sie in der Dokumentation zur PJRT-Laufzeit.

   Logo: PJRT

       (vm)$ export PJRT_DEVICE=TPU
    

   Logo: XRT

       (vm)$ export XRT_TPU_CONFIG="localservice;0;localhost:51011"
    

4. Wenn Sie mit Cloud TPU v4 trainieren, legen Sie auch die folgende Umgebungsvariable fest:

     (vm)$ export TPU_NUM_DEVICES=4

Einfache Berechnung ausführen

1. Starten Sie den Python-Interpretierer auf der TPU-VM:

   (vm)$ python3

2. Importieren Sie die folgenden PyTorch-Pakete:

   import torch
   import torch_xla.core.xla_model as xm

3. Geben Sie das folgende Skript ein:

   dev = xm.xla_device()
   t1 = torch.randn(3,3,device=dev)
   t2 = torch.randn(3,3,device=dev)
   print(t1 + t2)

   Hinweis: Die Fehlermeldung OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op wird erwartet und in einem zukünftigen Release behoben.

   Die angezeigte Ausgabe sieht so aus:

   tensor([[-0.2121,  1.5589, -0.6951],
          [-0.7886, -0.2022,  0.9242],
          [ 0.8555, -1.8698,  1.4333]], device='xla:1')
   

Resnet auf einem TPU-Einzelgerät ausführen

An dieser Stelle können Sie beliebigen PyTorch-/XA-Code ausführen. Sie haben beispielsweise die Möglichkeit, ein ResNet-Modell mit fiktiven Daten auszuführen:

(vm)$ git clone --recursive https://github.com/pytorch/xla.git
(vm)$ python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1

Das ResNet-Beispiel wird für eine Epoche trainiert und dauert etwa 7 Minuten. Die Ausgabe sollte in etwa so aussehen:

Epoch 1 test end 20:57:52, Accuracy=100.00 Max Accuracy: 100.00%

Nach Abschluss des ResNet-Trainings löschen Sie die TPU-VM.

(vm)$ exit

$ gcloud compute tpus tpu-vm delete tpu-name \
--zone=zone

Der Löschvorgang kann einige Minuten dauern. Prüfen Sie, ob die Ressourcen gelöscht wurden. Führen Sie dazu gcloud compute tpus list --zone=${ZONE} aus.

Themen für Fortgeschrittene

Bei Modellen mit großen, häufigen Zuordnungen verbessert tcmalloc die Leistung im Vergleich zur C/C++-Laufzeitfunktion malloc. Die Standard-malloc auf der TPU-VM ist tcmalloc. Sie können erzwingen, dass die TPU-VM-Software die Standard-malloc verwendet, indem Sie die Umgebungsvariable LD_PRELOAD aufheben:

   (vm)$ unset LD_PRELOAD

In den vorherigen Beispielen (einfache Berechnung und ResNet50) startet das PyTorch/XLA-Programm den lokalen XRT-Server im selben Prozess wie den Python-Interpreter. Sie können den lokalen XRT-Dienst aber auch in einem separaten Prozess starten:

(vm)$ python3 -m torch_xla.core.xrt_run_server --port 51011 --restart

Hinweis: Der zuvor angegebene Port ist der Port, den Sie zuvor im export XRT_TPU_CONFIG="localservice;0;localhost:51011"-Befehl angegeben haben.

Der Vorteil dieses Ansatzes besteht darin, dass der Kompilierungs-Cache über alle Trainingsläufe hinweg beibehalten wird. Wenn der XLA-Server in einem separaten Prozess ausgeführt wird, werden serverseitige Logging-Informationen in /tmp/xrt_server_log geschrieben.

(vm)$ ls /tmp/xrt_server_log/

server_20210401-031010.log

TPU-VM-Leistungsprofilerstellung

Weitere Informationen zur Profilerstellung für Ihre Modelle auf der TPU-VM finden Sie unter PyTorch XLA-Leistungsprofilerstellung.

Beispiele für PyTorch-/XLA-TPU-Pods

Unter PyTorch TPU VM Pod finden Sie Informationen zur Einrichtung sowie Beispiele zum Ausführen von PyTorch/XLA auf einem TPU-VM-Pod.

Docker auf TPU-VM

In diesem Abschnitt wird gezeigt, wie Sie Docker auf einer TPU-VM mit vorinstalliertem PyTorch/XLA ausführen.

Hinweis: Derzeit wird das Ausführen von Docker auf einer TPU-VM nur auf Einzelgerät-TPUs (v2-8, v3-8) unterstützt.

Verfügbare Docker-Images

In der GitHub-README-Datei finden Sie alle verfügbaren TPU-VM-Docker-Images.

Docker-Images auf TPU-VM ausführen

Hinweis: Die letzte Befehlsausgabe, root:/#, ist die Eingabeaufforderung, die nach der Ausführung des Docker-Images angezeigt wird.

(tpuvm): sudo docker pull gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm
(tpuvm): sudo docker run --privileged  --shm-size 16G --name tpuvm_docker -it -d  gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm
(tpuvm): sudo docker exec --privileged -it tpuvm_docker /bin/bash
(pytorch) root:/#

Hinweis: Sie müssen den Container mit dem Flag --privileged ausführen, damit TPU-Geräte im Container sichtbar sind. In diesem Container ist eine Conda-Umgebung pytorch vorinstalliert. Jeder Befehl, der mit (pytorch) beginnt, sollte im Container ausgeführt werden.

libtpu überprüfen

Hinweis:libtpu sollte auf neueren TPUVM-Docker-Images vorinstalliert sein. Es ist jedoch immer besser, sie manuell zu prüfen.

Führen Sie folgenden Befehl aus, um zu prüfen, ob libtpu installiert ist:

(pytorch) root:/# ls /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/ | grep libtpu

Die Ausgabe sollte in etwa so aussehen:

libtpu
libtpu_nightly-0.1.dev20220518.dist-info

Wenn keine Ergebnisse angezeigt werden, können Sie die entsprechende libtpu manuell installieren:

(pytorch) root:/# pip install torch_xla[tpuvm]

tcmalloc bestätigen

tcmalloc ist das Standard- Malloc, das wir auf der TPU-VM verwenden. Weitere Informationen finden Sie in diesem Abschnitt. Diese Bibliothek sollte auf neueren TPU-VM-Docker-Images vorinstalliert sein. Es ist jedoch immer besser, sie manuell zu überprüfen. Mit dem folgenden Befehl können Sie prüfen, ob die Bibliothek installiert ist.

(pytorch) root:/# echo $LD_PRELOAD

Die Ausgabe sollte in etwa so aussehen:

/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4

Wenn LD_PRELOAD nicht festgelegt ist, können Sie Folgendes manuell ausführen:

(pytorch) root:/# sudo apt-get install -y google-perftools
(pytorch) root:/# export LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4"

Gerät bestätigen

Mit folgendem Befehl können Sie prüfen, ob das TPU-VM-Gerät verfügbar ist:

(pytorch) root:/# ls /dev | grep accel

Dies sollte zu den folgenden Ergebnissen führen.

accel0
accel1
accel2
accel3

Wenn keine Ergebnisse angezeigt werden, haben Sie den Container wahrscheinlich nicht mit dem Flag --privileged gestartet.

Modell ausführen

Mit dem folgenden Befehl können Sie prüfen, ob das TPU-VM-Gerät verfügbar ist:

(pytorch) root:/# export XRT_TPU_CONFIG="localservice;0;localhost:51011"
(pytorch) root:/# python3 pytorch/xla/test/test_train_mp_imagenet.py --fake_data --num_epochs 1