PyTorch XLA-Leistungsprofil

Übersicht

In diesem Leitfaden erfahren Sie, wie Sie Leistungstools von Cloud TPU und die automatische Messwertanalyse von PyTorch verwenden. Diese Tools helfen Ihnen beim Debuggen und Optimieren der Leistung von Trainingsarbeitslasten.

Konzepte

Wenn Sie neu bei PyTorch / XL sind, finden Sie weitere Informationen in den Dokumenten zur PyTorch-API_GUIDE und Fehlerbehebung. Informationen zu Cloud TPU finden Sie im Konzeptdokument.

TPU-Knoten + PyTorch-/XLA-Profiling

Erforderliche Cloud TPU-Ressourcen erstellen und initialisieren

  1. Erstellen Sie Variablen für Ihre Projekt-ID, Ihren Cloud Storage-Bucket und die Zone, die für Ihre TPU-Ressourcen verwendet werden soll.

    export PROJECT_ID=PROJECT_ID
    export BUCKET_NAME=BUCKET_NAME
    export ZONE=ZONE
    
    gcloud --project=$PROJECT_ID compute project-info add-metadata \
    --metadata BUCKET_NAME=$BUCKET_NAME
    
  2. Compute Engine-VM-Instanz erstellen Hier werden alle Python-Skripts und -Modelle gespeichert.

    gcloud compute instances create profiler-tutorial-vm \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --machine-type=n1-standard-16 \
      --image-project=ml-images \
      --image-family=torch-xla \
      --boot-disk-size=300GB \
      --scopes=https://www.googleapis.com/auth/cloud-platform
    
  3. Erstellen Sie eine TPU-Ressource.

    gcloud compute tpus create profiler-tutorial-tpu \
     --project=${PROJECT_ID} \
     --zone=${ZONE} \
     --network=default \
     --version=pytorch-1.8 \
     --accelerator-type=v3-8
    
  4. Cloud Storage-Bucket erstellen

    Installieren Sie zuerst die gsutil-Befehlszeile, falls sie noch nicht installiert ist: Installationsanleitung.

  5. Erstellen Sie mit dem Befehl gsutil mb einen Cloud Storage-Bucket, in dem alle Profiling-Artefakte gespeichert sind. Ersetzen Sie die Regions- und Bucket-Namensvariablen durch die Werte, die Sie für das Training verwenden werden.

    gsutil mb -p ${PROJECT_ID} -c standard -l REGION gs://${BUCKET_NAME}
    

    wobei

    • REGION ist die Region, in der Sie die Cloud TPU erstellt haben, z. B. europe-west4.
  6. Erstellen Sie ein Dienstkonto für das Cloud TPU-Projekt.

    gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
    

    Der Befehl gibt ein Cloud TPU-Dienstkonto im folgenden Format zurück:

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    

    Beispiel: service-164006649440@cloud-tpu.iam.gserviceaccount.com

  7. Exportieren Sie das Dienstkonto und gewähren Sie Dienstkontoberechtigungen für den Storage-Bucket. Ersetzen Sie account-number durch die Projektausgabe, die in der Ausgabe der Dienstkontoerstellung zurückgegeben wurde.

    export SERVICE_ACCOUNT=service-ACCOUNT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    gsutil acl ch -u $SERVICE_ACCOUNT:READER gs://${BUCKET_NAME}
    gsutil acl ch -u $SERVICE_ACCOUNT:WRITER gs://${BUCKET_NAME}
    

TensorBoard einrichten

  1. ssh zu Ihrem VM-Weiterleitungsport 9001 auf Ihrer VM zu Port 9001 auf Ihrem lokalen Computer. Über diesen Port wird die TensorBoard-UI in Ihrem lokalen Browser geöffnet.

      gcloud compute ssh profiler-tutorial-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE} \
       --ssh-flag="-4 -L 9001:localhost:9001"
    
  2. Erstellen Sie eine Conda-Umgebung speziell für die TensorBoard-Installation:

    conda create -y -n tensorboard python=3.6
    conda activate tensorboard
    pip install tf-nightly==2.6.0.dev20210511 tb-nightly==2.6.0a20210512 tbp-nightly==2.5.0a20210511
    
  3. Testen Sie Ihre Installation, indem Sie den TensorBoard-Server auf Ihrer Compute Engine-VM ausführen und versuchen, dann über http://localhost:9001/#profile auf dem Server eine Verbindung zum Server herzustellen. Ihren lokalen Rechner:

     # Get bucket name
     BUCKET_NAME=$(curl "http://metadata.google.internal/computeMetadata/v1/project/attributes/BUCKET_NAME" -H "Metadata-Flavor: Google")
    
     tensorboard --logdir gs://${BUCKET_NAME} --port 9001
    

Wenn Sie auf Ihrem lokalen Computer http://localhost:9001/#profile aufrufen, sollten Sie in etwa Folgendes sehen:

Image

Wenn Sie http://localhost:9001 aufrufen, können Sie auch auf die obige Profilseite zugreifen. Dazu wählen Sie die Option "PROFILE" im Drop-down-Menü rechts oben neben der Schaltfläche "UPLOAD" aus. :

Image

Profil erstellen

Damit der TensorBoard-Server aktiv bleibt, starten Sie auf Ihrem lokalen Computer ein neues Terminalfenster und ssh noch einmal auf der GCE-VM (diesmal ohne -L Option für die Portweiterleitung).

  1. Exportieren Sie im neuen Terminalfenster Ihre Projekt-ID, die Storage-Bucket-Umgebung und die Zonenvariablen noch einmal, da sie sich in einer neuen Shell befindet.

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    
  2. ssh zur VM:

      gcloud compute ssh profiler-tutorial-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE}
    
    conda activate torch-xla-1.8
    PROJECT_ID=$(curl "http://metadata.google.internal/computeMetadata/v1/project/project-id" -H "Metadata-Flavor: Google")
    export TPU_IP_ADDRESS=$(gcloud compute tpus describe profiler-tutorial-tpu --zone=${ZONE} --project=${PROJECT_ID} \
    --format="value(ipAddress)")
    echo TPU_IP_ADDRESS=${TPU_IP_ADDRESS}
    export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
    
  3. Prüfen Sie, ob Integrationstests in Ihrer Umgebung durchgängig funktionieren:

    python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profiler.py # takes <1 min
    >
  4. Bearbeiten Sie vor dem Training die folgenden Zeilen in /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py:

    Änderung:

    accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
    
    An:
    accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
    

    Die beiden oben genannten Argumente in train_mnist verursachen künstliche dynamische Graphen und Tensorabrufe, die später im Abschnitt der Analyse automatischer Messwerte erläutert werden. Da Sie derzeit nur die TPU erstellen, wird das folgende Beispiel mit geringer Leistung ausgeführt.

  5. Starten Sie eine Trainingsausführung, die für die Serverprofilierung verwendet wird:

    PT_XLA_DEBUG=1 XLA_HLO_DEBUG=1 python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py --num_epochs 1000 --fake_data
    

TPU-Profile (Server)

Sobald das Training ausgeführt wurde, rufen Sie http://localhost:9001/#profile auf und erstellen Sie ein Profil gemäß der folgenden Anleitung:

Image

Die folgende Seite wird automatisch aktualisiert:

Image

Unterstützte Tools werden im Drop-down-Menü Tools im linken Bereich angezeigt:

  • Übersichtsseite (die Anzeige des Eingabe-Pipeline-Analysetools ist nicht enthalten, das sich zwischen TensorFlow/TPU und PyTorch / XLA TPU grundlegend unterscheidet)
  • Memory Viewer
  • Op Profile
  • Pod Viewer
  • TensorFlow-Statistiken (Statistiken auf Framework-Ebene, z.B. PyTorch-Statistiken)
  • Trace Viewer (Chrome-Browser erforderlich)

Übersichtsseite

Diese Seite bietet eine Übersicht über ein erfasstes Profil. Dieses Beispiel zeigt eine sehr hohe Leerlaufzeit, da ein kleines Modell für das MNIST-Dataset trainiert wird.

Memory Viewer

Zeigt den Arbeitsspeicher des Geräts (HBM), der pro Tensor und HLO PyTorch op verwendet wird. Der Arbeitsspeicherbetrachter erfasst eine Ansicht des Speichers pro HLO-Modul. Daher gibt es einige Module, z. B. Grafiken zu Gerätedaten (Eingabe und Labels im Batch gesendet) Um die Speichernutzung eines bestimmten HLO-Moduls anzusehen, wählen Sie im linken Drop-down-Menü Hosts das Modul aus:

Image

Wenn Sie das ausgewählte HLO-Modul aufrufen, erhalten Sie eine ganzheitliche Ansicht der HBM-Fußabdruckzeitachse für die Modulausführung. Die Reihenfolge ist nach Zuordnungsgröße, Programmausführungsgröße und Padding-Größe sortiert.

Image

Jede Pufferzuweisung kann dann weiter angezeigt werden, indem Sie den Mauszeiger darauf bewegen. So können Sie beispielsweise feststellen, welche Zuweisung die meisten Geräte von HBM beansprucht:

Image

Im obigen Beispiel entspricht (1) den vom Nutzercode hinzugefügten torch_xla.debug.profiler.Trace-Annotationen. Prüfen des test_profile_mp_mnist.py-Codes, der dieser Zeile entspricht:

   class MNIST(nn.Module):
   ...
     def forward(self, x):
       with xp.Trace('conv1'):
         x = F.relu(F.max_pool2d(self.conv1(x), 2))
         x = self.bn1(x)
   ...
   

Außerdem können Sie anhand des Vorgangs-Namespace test_mnist feststellen, dass dieses HLO-Modul der Bewertungsschleife entspricht, da es den Kontextmanager xp.trace('test_mnist') enthält.

XLA Op Profile

Op Profile ist ein Cloud TPU-Tool, das Leistungsstatistiken von XLA-Vorgängen anzeigt, die während eines Profilerstellungszeitraums ausgeführt wurden. Op Profile zeigt folgendes:

  • Wie gut Ihre Anwendung die Cloud TPU als Prozentsatz der Zeit verwendet, die für Vorgänge nach Kategorie und TPU-FLOPS-Auslastung aufgewendet wurde.
  • Vorgänge, die am zeitaufwendigsten waren. Diese Vorgänge sind potenzielle Ziele für die Optimierung.
  • Details zu einzelnen Vorgängen, einschließlich Form, Auffüllung und Ausdrücken, die den Vorgang verwenden.

Mithilfe des Vorgangsprofils finden Sie gute Optimierungsziele. Wenn Ihr Modell beispielsweise nur 5% der TPU-FLOPS erreicht, können Sie mit dem Tool ermitteln, welche XLA-Vorgänge die längste Ausführungszeit benötigen und wie viele TPU-FLOPS sie verbrauchen.

Image

Beschreibung der einzelnen:

  1. Übersicht. Zeigt die Cloud TPU-Auslastung und bietet Vorschläge zur Optimierung.
  2. Systemsteuerung. Enthält Steuerelemente, mit denen Sie die Anzahl der in der Tabelle angezeigten Vorgänge, die angezeigten Vorgänge und die Sortierung festlegen können.
  3. Vorgangstabelle. Eine Tabelle, in der die wichtigsten TensorFlow-Vorgangskategorien, die den XLA-Vorgängen zugeordnet sind, aufgeführt sind. Diese Vorgänge sind nach Prozentsatz der Cloud TPU-Nutzung sortiert.
  4. Karten mit Vorgangsdetails Details zum Vorgang, der angezeigt wird, wenn Sie den Mauszeiger in der Tabelle auf einen Vorgang bewegen. Dazu gehören die FLOPS-Auslastung, der Ausdruck, in dem der Vorgang verwendet wird, und das Vorgangslayout (passend).

Pod Viewer

Eine vollständige Beschreibung des Pod-Viewers finden Sie unter TPU-Tools.

Framework-Statistiken (TensorFlow-/PyTorch-Statistiken)

Framework-Statistiken bieten eine detaillierte Aufschlüsselung der Statistiken zu PyTorch und XRT-Vorgängen, die auf TPU-Geräten und TPU-Hosts ausgeführt werden.

Trace Viewer

Trace Viewer ist ein Cloud TPU-Tool zur Leistungsanalyse. Der Trace Viewer verwendet den Chrome Trace Event Profiler, daher muss der Chrome-Browser verwendet werden.

Trace Viewer enthält eine Zeitachse mit folgenden Informationen:

  • Dauer für die Vorgänge, die von Ihrem TensorFlow-Modell ausgeführt wurden.
  • Teil des Systems (TPU oder Hostcomputer), in dem ein Vorgang ausgeführt wurde. Bei PyTorch / XL arbeitet der Hostcomputer in der Regel mit der Kompilierung und der Zwischenspeicherzuweisung/-dealisierung, während die TPU das eigentliche Modelltraining ausführt.
  • Mit der Trace-Anzeige können Sie Leistungsprobleme in Ihrem Modell ermitteln und dann Maßnahmen ergreifen, um diese zu beheben. Sie können ermitteln, welche PyTorch-/XL-Vorgänge am längsten dauern.

Sie können Traces direkt hinzufügen, um zu messen, wie lange bestimmte Teile des Modells dauern, indem Sie xp.Trace(NAME)-Annotationen hinzufügen. Der folgende Trace zeigt beispielsweise:

Image

  1. Generiert durch explizite Nutzerannotationen im Modellcode von test_profile_mp_mnist.py.
  2. PyTorch-Vorgänge wurden ausgeführt (Vorabsenkung).
  3. Name des automatisch von PyTorch / XLA generierten HLO-Moduls.
  4. XLA-Vorgänge, die auf dem Gerät ausgeführt werden (zusammengeführt).

Ausführlichere Informationen finden Sie in der allgemeinen TPU-Dokumentation für den Trace Viewer. Allerdings ignorieren Sie die Abschnitte zur Eingabepipeline und zu anderen TensorFlow-spezifischen Teilen, da diese für die Kontext dieses Dokuments.

PyTorch-/XLA-Clientprofiling

Ähnlich wie bei der Profilerstellung für die TPU-Seite während der Modellausführung werden Sie nun während des Trainings ein Profil für die PyTorch-/XL-Clientseite erstellen. Das wichtigste Monitoringtool, das auf der Clientseite verwendet wird, ist die Trace-Anzeige.

Sie müssen den Profiler-Server in Ihrem Trainingsskript starten. Ein Beispiel finden Sie unter see. Sie können dann von TensorBoard abfragen, um ein Trace zu erfassen.

Image

Zum Erfassen von Traces aus mehreren Prozessen kann jeder Prozess Profilserver an verschiedenen Ports starten (z. B. durch Hinzufügen von xm.get_ordinal() zu einer Basisportnummer) und anschließend eine Liste von localhost:port verkettet durch ",". Da TensorBoard die Anzeige von Traces mehrerer Prozesse gleichzeitig nicht unterstützt, ist für jeden Prozess ein anderes Host-Drop-down-Menü zu sehen.

Das folgende Diagramm zeigt ein Beispiel-Trace:

Image

Ähnlich wie beim Hinzufügen unterschiedlicher Namespace-Traces für TPU-Traces können Sie sie mit derselben API zu den clientseitigen Traces (xp.Trace(NAME)) hinzufügen. Da dieses Modell klein ist und mit kleinen MNIST-Bildern verwendet wird, sind die Schrittzeiten kurz und nicht unbedingt einheitlich. Als Übung können Sie versuchen, die Traces hinzuzufügen und einen Profiler-Server wie den in unserem Beispiel zu starten.test_train_mp_imagenet.py --fake_data Traces von ResNet50 abrufen.

Die Traces enthalten zusätzliche Metadaten, die geprüft werden können. TransferToServer- und TransferFromServer-Traces enthalten beispielsweise die genaue Anzahl der gesendeten und empfangenen Tensoren sowie ihre Gesamtgröße:

Image

Für XLA-Grafikkompilierungen wird der Grafik-Hash angezeigt, der bei der Diagnose von Problemen hilfreich sein kann:

Image

Zusätzlich zur Profilerstellung über die TensorBoard-UI bieten wir eine API für die programmatische Profilerstellung sowohl der TPU als auch des Clients von PyTorch / XL: xp.trace().

Analyse automatischer Messwerte

In diesem Abschnitt erfahren Sie, wie Sie mithilfe des Debugging-Modus Leistungsprobleme wie die folgenden erkennen:

  • Dynamische Grafiken / kontinuierliche Zusammenstellungen
  • Sehr langsame Grafikkompilierung
  • Sehr langsame Ausführung von Grafiken
  • Häufige XLA-CPU-Übertragungen
  • Wiederholtes HBM auf dem Gerät zum Hostwechsel des RAM
  • Wiederholte HBM-Defragmentierung
  • aten:: Vorgänge ohne Limit

Setzen Sie vor dem Training die folgenden Zeilen in /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py zurück:

Änderung:

   accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
   
An:
   accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
   

Diese Änderungen verursachen künstliche Zusammenstellungen und Tensorabrufe. dynamic_graph=True ändert die Batchgröße für jeden Schritt künstlich, wodurch die XLA-Werte mit Senkung bei jedem Schritt und bei der Neukompilierung unterschiedlich sind. fetch_often=True fügt bei jedem Schritt loss.item()-Aufrufe ein, was dazu führt, dass in jedem Schritt Tensorwerte vom Gerät abgerufen werden, was die Leistung verlangsamt.

Beispiel-Trainingsskript ausführen:

   PT_XLA_DEBUG=1 python /usr/share/torch-xla-1.8/pytorch/xla/test/test_profile_mp_mnist.py --fake_data --num_cores=1
   

Die Fehlerbehebung hat sich bewährt, zur Fehlerbehebung mit --num_cores=1 auszuführen, da dies den Fehlerbehebungsprozess vereinfacht. Einige Beispielausgaben sehen so aus:

Epoch 1 train begin 01:18:05
| Training Device=xla:1/0 Step=0 Loss=0.00000 Rate=1905.00 GlobalRate=1904.72 Time=01:18:05
pt-xla-profiler: TransferFromServerTime too frequent: 3 counts during 3 steps
pt-xla-profiler: TransferFromServerTime too frequent: 4 counts during 4 steps
pt-xla-profiler: TransferFromServerTime too frequent: 5 counts during 5 steps
pt-xla-profiler: TransferFromServerTime too frequent: 6 counts during 6 steps
pt-xla-profiler: TransferFromServerTime too frequent: 7 counts during 7 steps
pt-xla-profiler: TransferFromServerTime too frequent: 8 counts during 8 steps
pt-xla-profiler: TransferFromServerTime too frequent: 9 counts during 9 steps
pt-xla-profiler: TransferFromServerTime too frequent: 10 counts during 10 steps
pt-xla-profiler: CompileTime too frequent: 21 counts during 11 steps
pt-xla-profiler: TransferFromServerTime too frequent: 11 counts during 11 steps
pt-xla-profiler: CompileTime too frequent: 23 counts during 12 steps

Zeilen mit dem Präfix pt-xla-profiler entsprechen der Ausgabe der Analyse zu automatischen Messwerten. In diesem Beispiel wird gezeigt, dass TransferFromServerTime einmal pro Schritt zu häufig gesehen wurde. Dies liegt daran, dass die Trainingsschleife bei jedem Schritt den Wert von loss.item() abruft. Außerdem wird möglicherweise die Warnung CompileTime zu häufig angezeigt, wenn Ihr Modell aufgrund von dynamischen Formen in der Grafik wiederholt neu kompiliert werden muss. Das folgende Code-Snippet verursacht diese Art von Problem: test_profile_mp_mnist.py

    for step, (data, target) in enumerate(loader):
      if dynamic_graph:
        # The batch dimension is different every step.
        index = max(-step, -flags.batch_size + 1)  # non-empty
        data, target = data[:-index, :, :, :], target[:-index]
      ...
      if fetch_often:
        # Fetch tensor value from XLA:TPU to CPU every step.
        loss_i = loss.item()

Wenn Sie dann das Trainingsskript mit Strg + C beenden, sollten Sie am Ende eine Zusammenfassung der nicht erhöhten Vorgänge sehen. Beachten Sie, dass aten::_local_scalar_dense ein spezieller Vorgang ist, der dem Abrufen von XLA-Tensoren nach dem CPU-Kontext entspricht.

In diesem Bericht sehen Sie, dass es zwei Hauptorte gibt, an denen die Operation aten::_local_scalar_dense aufgerufen wird. Beide entsprechen dem Quellcode von loss.item():

  • test/test_profile_mp_mnist.py:158
  • test/test_profile_mp_mnist.py:61
pt-xla-profiler: ================================================================================
pt-xla-profiler: Unlowered Op usage summary (more of these ops, lower performance)
pt-xla-profiler: Note: _local_scalar_dense typically indicates CPU context access
pt-xla-profiler: --------------------------------------------------------------------------------
pt-xla-profiler: FRAME (count=27):
pt-xla-profiler: Unlowered Op: "_local_scalar_dense"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   train_loop_fn (test/test_profile_mp_mnist.py:158)
pt-xla-profiler:   train_mnist (test/test_profile_mp_mnist.py:184)
pt-xla-profiler:   _mp_fn (test/test_profile_mp_mnist.py:206)
pt-xla-profiler:   _start_fn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:323)
pt-xla-profiler:   spawn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:386)
pt-xla-profiler:    (test/test_profile_mp_mnist.py:216)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: FRAME (count=2):
pt-xla-profiler: Unlowered Op: "_local_scalar_dense"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   _train_update (test/test_profile_mp_mnist.py:61)
pt-xla-profiler:    (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:700)
pt-xla-profiler:   _run_step_closures (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:709)
pt-xla-profiler:   mark_step (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/core/xla_model.py:723)
pt-xla-profiler:   __exit__ (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/debug/profiler.py:153)
pt-xla-profiler:   train_loop_fn (test/test_profile_mp_mnist.py:162)
pt-xla-profiler:   train_mnist (test/test_profile_mp_mnist.py:184)
pt-xla-profiler:   _mp_fn (test/test_profile_mp_mnist.py:206)
pt-xla-profiler:   _start_fn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:323)
pt-xla-profiler:   spawn (/home/jysohn/git/jysohn23/pytorch/xla/torch_xla/distributed/xla_multiprocessing.py:386)
pt-xla-profiler:    (test/test_profile_mp_mnist.py:216)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: ================================================================================

Führen Sie nun eine automatische Messwertanalyse für das folgende Skript aus, das einen nicht verringerten Vorgang (ab Version 1.8) _ctc_loss enthält:

PT_XLA_DEBUG=1 python <<EOF
import torch
import torch_xla.core.xla_model as xm

dev = xm.xla_device()
t = torch.randn(50, 16, 20).log_softmax(2).to(dev)
target = torch.randint(low=1, high=20, size=(16, 30), dtype=torch.long).to(dev)
input_lengths = torch.full(size=(16,), fill_value=50, dtype=torch.long).to(dev)
target_lengths = torch.randint(low=10, high=30, size=(16,), dtype=torch.long).to(dev)

for _ in range(10):
  loss = torch.nn.CTCLoss()(t, target, input_lengths, target_lengths)
  xm.mark_step()
EOF

Wenn Sie das Skript oben mit PT_XLA_DEBUG=1 ausführen, sollte es in etwa so aussehen:

…
pt-xla-profiler: TransferFromServerTime too frequent: 30 counts during 10 steps
pt-xla-profiler: Op(s) not lowered: aten::_ctc_loss,  Please open a GitHub issue with the above op lowering requests.
pt-xla-profiler: ================================================================================
pt-xla-profiler: Unlowered Op usage summary (more of these ops, lower performance)
pt-xla-profiler: Note: _local_scalar_dense typically indicates CPU context access
pt-xla-profiler: --------------------------------------------------------------------------------
pt-xla-profiler: FRAME (count=10):
pt-xla-profiler: Unlowered Op: "_ctc_loss"
pt-xla-profiler: Python Frames:
pt-xla-profiler:   ctc_loss (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/functional.py:2305)
pt-xla-profiler:   forward (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/modules/loss.py:1593)
pt-xla-profiler:   _call_impl (/anaconda3/envs/torch-xla-1.8/lib/python3.6/site-packages/torch/nn/modules/module.py:889)
pt-xla-profiler:    (:11)
pt-xla-profiler:
pt-xla-profiler:
pt-xla-profiler: ================================================================================

Das Analysetool für automatische Messwerte gibt an, dass die 11. Zeile des STDIN zu diesem nicht senkenden Vorgang führt (d. h. zur Zeile mit torch.nn.CTCLoss()). Derzeit ist der ctc_loss-Vorgang nicht gesunken. . Deshalb wird der obige Bericht angezeigt. Sie können auch einige Warnungen für TransferFromServerTime sehen, da sich die Tensoren vor der Ausführung in XLA:TPU befinden. Da der Vorgang jedoch nicht reduziert wird, müssen Sie zuerst die XLA-Tensoren zurück an übertragen. Führen Sie den Vorgang aten:: für die CPU aus und übertragen Sie ihn zurück.

Wenn Sie stattdessen die Ausgabe von pt-xla-profiler in eine Datei schreiben möchten, legen Sie PT_XLA_DEBUG=1 und PT_XLA_DEBUG_FILE=$PATH_TO_FILE fest.

Clean-up

Beenden Sie die VM und löschen Sie dann die TPU, die VM und den Cloud Storage-Bucket mit den folgenden Befehlen:

(vm)$ exit
gcloud compute instances delete profiler-tutorial-vm \
  --zone=${ZONE} \
  --project=${PROJECT_ID}
gcloud compute tpus delete profiler-tutorial-tpu \
  --zone=${ZONE} \
  --project=${PROJECT_ID} \
  --async
gsutil rm -fr gs://${BUCKET_NAME}

TPU VM + PyTorch/XLA-Profiling

In diesem Abschnitt können Sie mithilfe der TPU-VM-Architektur ein Profil für PyTorch/XLA erstellen.

Umgebungsvariablen exportieren

  1. Erstellen Sie Variablen für Ihre Projekt-ID und die Zone, die für Ihre TPU-Ressourcen verwendet werden sollen.

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    

Cloud TPU erstellen

Weitere Informationen finden Sie im TPU-VM-Nutzerhandbuch. Erstellen Sie nach der Einrichtung eine v3-8-TPU-VM, die im Lieferumfang von torch, torch_xla und { enthalten ist. 101}torchvision und tensorboard sind vorinstalliert.

  1. Erstellen Sie eine TPU-Ressource.

    gcloud alpha compute tpus tpu-vm create profiler-tutorial-tpu-vm \
     --project=${PROJECT_ID} \
     --zone=${ZONE} \
     --version=v2-alpha \
     --accelerator-type=v3-8
    

Start des TensorBoard-Servers

  1. Stellen Sie eine SSH-Verbindung zur VM her, installieren Sie tensorboard-plugin-profile und starten Sie einen TensorBoard-Server.

      gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE} \
       --ssh-flag="-4 -L 9001:localhost:9001"
    
      pip3 install tf-nightly==2.6.0.dev20210511 tb-nightly==2.6.0a20210512 tbp-nightly==2.5.0a20210511
    
      tensorboard --logdir ./tensorboard --port 9001
    

Wenn Sie die TensorBoard-Ausgabe unter http://localhost:9001 auf Ihrem lokalen Computer ansehen, sollten Sie in etwa Folgendes sehen:

Image

Wenn Sie die TensorBoard-Ausgabe unter http://localhost:9001 ansehen, können Sie auch auf die obige Profilseite zugreifen. Dazu wählen Sie die Option "PROFILE" im Drop-down-Menü rechts oben neben der Schaltfläche "UPLOAD" aus.

Image

Profil erstellen

Exportieren Sie in einem neuen Terminalfenster in Ihrer Entwicklungsumgebung dieselben Umgebungsvariablen wie oben und ssh auf Ihre TPU-VM:

  1. Exportieren Sie im neuen Terminalfenster Ihre Projekt-ID und Zonenvariablen noch einmal, da sie sich in einer neuen Shell befinden.

    export PROJECT_ID=PROJECT_ID
    export ZONE=ZONE
    
  2. ssh zur VM:

      gcloud alpha compute tpus tpu-vm ssh profiler-tutorial-tpu-vm \
       --project ${PROJECT_ID} \
       --zone ${ZONE}
    
  3. Klonen Sie das PyTorch/XLA-Repository und führen Sie unseren e2e-Test aus:

      git clone -b r1.8 https://github.com/pytorch/xla
      export XRT_TPU_CONFIG="localservice;0;localhost:51011"
      python3 xla/test/test_profiler.py  # takes <1 min
    >
  4. Bearbeiten Sie vor dem Start des Trainings die folgenden Zeilen in xla/test/test_profile_mp_mnist.py:

    Änderung:

        accuracy = train_mnist(flags, dynamic_graph=True, fetch_often=True)
    
    An:
        accuracy = train_mnist(flags, dynamic_graph=False, fetch_often=False)
    

    Die beiden oben genannten Argumente in train_mnist verursachen künstliche dynamische Graphen und Tensorabrufe, die später im Abschnitt der Analyse automatischer Messwerte erläutert werden. Da Sie derzeit nur die TPU erstellen, wird das folgende Beispiel mit geringer Leistung ausgeführt.

  5. Starten Sie einen Trainingslauf:

     XLA_HLO_DEBUG=1 python3 xla/test/test_profile_mp_mnist.py --num_epochs 1000 --fake_data
    

TPU + Client-Profiling

Wenn das Training ausgeführt wird, sehen Sie sich die TensorBoard-Ausgabe unter http://localhost:9001 an und erfassen Sie ein Profil mithilfe der folgenden Anweisungen:

Image

Daraufhin sollte die folgende Seite aktualisiert werden:

Image

Derzeit ist in der TPU-VM-Einrichtung nur das Trace-Viewer-Tool ausgewählt. Wählen Sie im Drop-down-Menü Tools die Option trace_viewer aus und prüfen Sie die Traces. In der TPU-VM-Einrichtung sehen Sie sowohl die clientseitigen als auch die TPU-Geräte-Traces in einer vollständigen Ansicht:

Image

Clean-up

  1. Beenden Sie die VM und löschen Sie dann die TPU, die VM und den Cloud Storage-Bucket mit den folgenden Befehlen:

    (vm)$ exit
    

Löschen Sie die von Ihnen erstellte TPU-VM:

  1. Löschen Sie Ihre Cloud TPU- und Compute Engine-Ressourcen.

    $ gcloud alpha compute tpus tpu-vm delete profiler-tutorial-tpu-vm \
      --project ${PROJECT_ID} --zone=${ZONE}
    
  2. Prüfen Sie mit dem folgenden Befehl, ob die Ressourcen gelöscht wurden. Der Löschvorgang kann einige Minuten dauern. Eine Antwort wie die folgende gibt an, dass Ihre Instanzen erfolgreich gelöscht wurden.

    $ gcloud alpha compute tpus tpu-vm list --project ${PROJECT_ID} --zone=${ZONE}
    
    Listed 0 items.