Profil für PyTorch-XLA-Arbeitslasten erstellen

Die Profilerstellung ist eine Möglichkeit, die Leistung von Modellen zu analysieren und zu verbessern. Obwohl geht es viel mehr darum. Manchmal hilft es, sich Profilerstellung als Timing Vorgänge und Codeteile, die auf beiden Geräten (TPUs) und Hosts ausgeführt werden (CPUs) In diesem Leitfaden erhalten Sie einen kurzen Überblick darüber, wie Sie Ihren Code für das Training oder die Inferenz profilieren. Weitere Informationen zur Analyse generierter Profile finden Sie in den folgenden Anleitungen.

• PyTorch XLA-Leistungsfehlerbehebung auf TPU-VMs – Teil 1
• PyTorch XLA-Leistungsfehlerbehebung auf TPU-VMs – Teil 2
• PyTorch XLA-Leistungsfehlerbehebung auf TPU-VMs – Teil 3

Jetzt starten

TPU erstellen

1. Exportieren Sie Umgebungsvariablen:

   $ export TPU_NAME=your_tpu_name
   $ export ZONE=us-central2-b
   $ export PROJECT_ID=project-id
   $ export ACCELERATOR_TYPE=v4-8
   $ export RUNTIME_VERSION=tpu-vm-v4-pt-2.0

   Variablenbeschreibungen exportieren

   TPU name

   Der Name, den Sie für Ihre Cloud TPU verwenden möchten.

   zone

   Die Zone, in der Sie die Cloud TPU erstellen möchten.

   project ID

   Die Projekt-ID, die Sie zum Trainieren und Profilieren Ihres Modells verwenden.

   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.

   version

   Die Version der Cloud TPU-Laufzeit. In der exportierten Datei wird ein Standardwert Sie können aber auch eine Variable aus der Liste unterstützten Konfigurationen.

2. TPU-Ressourcen starten

   $ gcloud compute tpus tpu-vm create ${TPU_NAME} \
   --zone us-central2-b \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --version ${RUNTIME_VERSION} \
   --project $PROJECT_ID \
   --subnetwork=tpusubnet

3. Verschieben Sie Ihren Code mit dem Befehl gcloud scp in Ihr Basisverzeichnis auf der TPU-VM. Beispiel:

   $ gcloud compute tpus tpu-vm scp my-code-file ${TPU_NAME}: --zone ${ZONE}

Profilerstellung

Hinweis: Ein v4-8 enthält 4 TPU-Geräte. Jeder Mitarbeiter profiliert nur ein TPU-Gerät aus den vier verfügbaren Geräten. Wenn Sie für jeden Worker ein Profil auf allen vier TPU-Geräten erstellen, wird bald unterstützt.

Ein Profil kann manuell über capture_profile.py oder programmatisch im Trainingsskript mithilfe der torch_xla.debug.profiler APIs erfasst werden.

Starten des Profilservers

Damit ein Profil erfasst werden kann, muss im Trainingsskript ein Profilerstellungs-Server ausgeführt werden. Starten Sie einen Server mit einer Portnummer Ihrer Wahl, z. B. 9012, wie im folgenden Befehl gezeigt.

Hinweis: In diesem Beispielscript sind die folgenden Schritte aufgeführt.

import torch_xla.debug.profiler as xp
server = xp.start_server(9012)

Der Server kann direkt am Anfang der main-Funktion gestartet werden.

Hinweis: Wenn Sie das Modul Lightning verwenden, übergeben Sie anstelle des Startens eines Servers profiler="xla" an die trainer. Der Server wird dann automatisch auf Port 9012 gestartet.

Sie können jetzt Profile wie im folgenden Abschnitt beschrieben erfassen. Das Skript wird ein Profil für alle Vorgänge auf einem TPU-Gerät erstellt.

Traces hinzufügen

Wenn Sie auch um Profile für Vorgänge auf dem Hostcomputer zu erstellen, können Sie xp.StepTrace oder xp.Trace in Ihren Code ein. Diese Funktionen verfolgen den Python-Code auf der Hostcomputer. Sie können sich das so vorstellen: Es wird gemessen, wie viel Zeit es dauert, den Python-Code auf dem Host (CPU) auszuführen, bevor der „Graph“ an das TPU-Gerät übergeben wird. Daher ist es hauptsächlich für die Analyse des Tracing-Overheads nützlich. Sie können fügen Sie dies in die Trainingsschleife ein, in der der Code Batches von Daten verarbeitet. zum Beispiel

for step, batch in enumerate(train_dataloader):
    with xp.StepTrace('Training_step', step_num=step): 
       ...

oder einzelne Teile des Codes mit

 with xp.Trace('loss'): 
    loss = ...

Wichtig: Vermeiden Sie es, xm.mark_step() in xp.Trace zu verpacken. In diesem Fall stürzt der Code mit einer Fehlermeldung ab, die in etwa so lautet: "RuntimeError: Expecting scope to be empty but it is train_loop." xp.StepTrace hat dieses Problem nicht, da xm.mark_step() beim Verlassen des Gültigkeitsbereichs automatisch hinzugefügt wird. Außerdem kann dieser Fehler auftreten, wenn Sie verwenden torchdynamo, weil damit intern xm.mark_step aufgerufen wird. Die aktuelle Problemumgehung besteht darin, nicht xp.Trace zu verwenden.

Wenn Sie Lighting verwenden, können Sie das Hinzufügen von Traces überspringen, da dies in einigen Teilen des Codes automatisch geschieht. Wenn Sie jedoch zusätzliche Traces hinzufügen möchten, können Sie sie in den Trainings-Loop einfügen.

Sie können die Geräteaktivitäten nach der ersten Kompilierung erfassen. Warten Sie, bis das Modell mit den Trainings- oder Inferenzschritten beginnt.

Manuelle Aufnahme

Das Skript capture_profile.py aus dem Pytorch XLA-Repository ermöglicht die schnelle Aufnahme eines Profils. Kopieren Sie dazu die Profildatei aufnehmen direkt zu Ihrer TPU-VM. Mit dem folgenden Befehl wird er in das Basisverzeichnis kopiert.

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone us-central2-b \
--worker=all \
--command="wget https://raw.githubusercontent.com/pytorch/xla/master/scripts/capture_profile.py"

Führen Sie während des Trainings Folgendes aus, um ein Profil zu erfassen:

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--zone us-central2-b \
--worker=all \
--command="python3 capture_profile.py --service_addr "localhost:9012" --logdir ~/profiles/ --duration_ms 2000"

Mit diesem Befehl werden .xplane.pb-Dateien in logdir gespeichert. Sie können Logging-Verzeichnis ~/profiles/ zu deinem bevorzugten Standort und Namen. Es ist auch möglich, direkt im Cloud Storage-Bucket speichern. Legen Sie dazu logdir wird zu gs://your_bucket_name/.

Programmatische Erfassung

Anstatt das Profil manuell durch Auslösen eines Skripts zu erfassen, können Sie Trainingsskript so konfigurieren, dass automatisch ein Profil ausgelöst wird mithilfe der Methode torch_xla.debug.profiler.trace_detached API im Trainingsskript.

Hinweis: Die trace_detached API ist nur in nächtlichen Builds oder danach verfügbar. Version 2.2. Für die programmatische Aufnahme in früheren Versionen von torch_xla müssen Sie Folgendes tun: kann den Text der Methode kopieren, um den Hintergrunderfassungs-Thread zu erstellen direkt ausführen.

Wenn Sie z. B. ein Profil in einer bestimmten Epoche und bei einem bestimmten Schritt automatisch erfassen möchten, können Sie Ihr Trainingsskript so konfigurieren, dass PROFILE_STEP, PROFILE_EPOCH- und PROFILE_LOGDIR-Umgebung Variablen:

import os
import torch_xla.debug.profiler as xp

# Within the training script, read the step and epoch to profile from the
# environment.
profile_step = int(os.environ.get('PROFILE_STEP', -1))
profile_epoch = int(os.environ.get('PROFILE_EPOCH', -1))
...

for epoch in range(num_epoch):
   ...
   for step, data in enumerate(epoch_dataloader):
      if epoch == profile_epoch and step == profile_step:
         profile_logdir = os.environ['PROFILE_LOGDIR']
         # Use trace_detached to capture the profile from a background thread
         xp.trace_detached('localhost:9012', profile_logdir)
      ...

Dadurch werden die .xplane.pb-Dateien im Verzeichnis gespeichert, das durch die Umgebungsvariable PROFILE_LOGDIR angegeben ist.

Analyse in TensorBoard

Für eine weitere Analyse von Profilen können Sie TensorBoard mit dem TPU TensorBoard-Plug-in
entweder auf demselben oder auf einem anderen Computer (empfohlen) verwenden.

Wenn Sie TensorBoard auf einem Remote-Computer ausführen möchten, stellen Sie über SSH eine Verbindung zu ihm her und aktivieren Sie Portweiterleitung Beispiel:

$ ssh -L 6006:localhost:6006 remote server address

oder

$ gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --ssh-flag="-4 -L 6006:localhost:6006"

Installieren Sie auf Ihrem Remote-Computer die erforderlichen Pakete und starten Sie TensorBoard. (vorausgesetzt, Sie haben auf diesem Computer Profile unter ~/profiles/). Wenn Sie Profilen in einem anderen Verzeichnis oder Cloud Storage-Bucket müssen Pfade korrekt angegeben werden, z. B. gs://your_bucket_name/profiles.

(vm)$ pip install tensorflow-cpu tensorboard-plugin-profile 

(vm)$ tensorboard --logdir ~/profiles/ --port 6006

Hinweis: .xplane-Dateien müssen sich in derselben verschachtelten Struktur befinden, die beim Erfassen eines Profils generiert wird. Beispiel: ~/profiles/plugins/profile/2023_04_10_21_40_22/localhost_9012.xplane.pb

Hinweis: Wenn die Fehlermeldung „Duplizierte Plugins“ angezeigt wird, kann das daran liegen, dass mehrere Pakete mit TensorBoard vorhanden sind. Deinstallieren Sie alle diese Pakete und führen Sie die pip install-Befehle noch einmal aus. Möglicherweise müssen Sie beispielsweise Folgendes ausführen:

(vm)$ pip uninstall tensorflow tf-nightly tensorboard tb-nightly tbp-nightly

TensorBoard ausführen

Rufen Sie in Ihrem lokalen Browser Folgendes auf: http://localhost:6006/ und wählen Sie PROFILE aus dem Drop-down-Menü aus, um Ihre Profile zu laden.

Weitere Informationen finden Sie unter TPU-Tools. TensorBoard-Tools und deren Interpretation der Ausgabe.