Profil für PyTorch XLA-Arbeitslasten erstellen

Die Profilerstellung ist eine Möglichkeit, die Leistung von Modellen zu analysieren und zu verbessern. Es ist zwar viel mehr, aber manchmal ist es hilfreich, sich die Profilerstellung als zeitliche Abläufe und Teile des Codes zu vorstellen, die sowohl auf Geräten (TPUs) als auch auf Hosts (CPUs) ausgeführt werden. Dieser Leitfaden bietet einen kurzen Überblick darüber, wie Sie für Ihren Code ein Profil für Training oder Inferenz erstellen. Weitere Informationen zum Analysieren generierter Profile finden Sie in den folgenden Leitfäden.

• Fehlerbehebung für die Leistung von PyTorch XLA auf TPU-VMs – Teil 1
• Fehlerbehebung für die Leistung von PyTorch XLA auf TPU-VMs – Teil 2
• Debugging von PyTorch XLA 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 Ihre Cloud TPU erstellen möchten.

   project ID

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

   accelerator-type

   Der Beschleunigertyp gibt die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Weitere Informationen zu unterstützten Beschleunigertypen für die einzelnen TPU-Versionen finden Sie unter TPU-Versionen.

   version

   Die Version der Cloud TPU-Laufzeit. Die exportierte Variable enthält einen Standardwert, aber Sie können auch eine Konfiguration aus der Liste der unterstützten Konfigurationen verwenden.

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 Worker erstellt nur für ein TPU-Gerät von 4 verfügbaren Geräten ein Profil. Die Profilerstellung für alle 4 TPU-Geräte für jeden Worker wird bald unterstützt.

Ein Profil kann manuell über capture_profile.py oder programmatisch aus dem Trainingsskript mit den torch_xla.debug.profiler APIs erfasst werden.

Profilserver starten

Um ein Profil zu erfassen, muss innerhalb des Trainingsskripts ein Profilserver ausgeführt werden. Starten Sie einen Server mit einer Portnummer Ihrer Wahl, z. B. 9012, wie im folgenden Befehl gezeigt.

Hinweis: Sie können auf dieses Beispielskript zurückgreifen, das die folgenden Schritte enthält.

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

Der Server kann direkt zu Beginn der main-Funktion gestartet werden.

Hinweis: Wenn Sie das Modul Lightning verwenden, anstatt einen Server zu starten, übergeben Sie profiler="xla" an trainer. Der Server wird dann automatisch bei Port 9012 gestartet.

Sie können jetzt Profile wie im folgenden Abschnitt beschrieben erfassen. Das Skript erstellt Profile für alles, was auf einem TPU-Gerät passiert.

Traces hinzufügen

Wenn Sie auch ein Profil von Vorgängen auf dem Hostcomputer erstellen möchten, können Sie xp.StepTrace oder xp.Trace in Ihren Code einfügen. Mit diesen Funktionen wird der Python-Code auf dem Hostcomputer verfolgt. (Sie können sich dies so vorstellen, als zu messen, wie viel Zeit für die Ausführung des Python-Codes auf dem Host (CPU) benötigt wird, bevor die „Grafik“ an das TPU-Gerät übergeben wird. Daher ist es besonders nützlich, um den Tracing-Overhead zu analysieren.) Sie können dies in die Trainingsschleife einfügen, in der der Code Datenbatches verarbeitet, z. B.

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 fassen. Der Code stürzt in diesem Fall mit einer Fehlermeldung wie dieser ab: "RuntimeError: Expecting scope to be empty but it is train_loop." Bei xp.StepTrace tritt das Problem nicht auf, da beim Verlassen des Bereichs automatisch xm.mark_step() hinzugefügt wird. Außerdem kann dieser Fehler auftreten, wenn Sie torchdynamo verwenden, da torchdynamo xm.mark_step intern aufrufen kann. Verwenden Sie derzeit nicht xp.Trace, um das Problem zu umgehen.

Wenn Sie Lighting verwenden, können Sie das Hinzufügen von Traces überspringen, da dies in einigen Teilen des Codes automatisch erfolgt. Wenn Sie jedoch weitere Traces hinzufügen möchten, können Sie diese in der Trainingsschleife einfügen.

Sie können Geräteaktivitäten nach der ersten Kompilierung erfassen. Warten Sie, bis das Modell seine Trainings- oder Inferenzschritte beginnt.

Manuelle Erfassung

Mit dem Skript capture_profile.py aus dem Pytorch XLA-Repository kann schnell ein Profil erfasst werden. Kopieren Sie dazu die Erfassungsprofildatei direkt auf Ihre TPU-VM. Mit dem folgenden Befehl wird sie 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 folgenden Befehl 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 im logdir gespeichert. Sie können das Logging-Verzeichnis ~/profiles/ in Ihren bevorzugten Speicherort und Namen ändern. Es ist auch möglich, Daten direkt im Cloud Storage-Bucket zu speichern. Dazu setzen Sie logdir auf gs://your_bucket_name/.

Programmatische Erfassung

Anstatt das Profil manuell durch Auslösen eines Skripts zu erfassen, können Sie Ihr Trainingsskript so konfigurieren, dass es automatisch ein Profil auslöst. Verwenden Sie dazu die API torch_xla.debug.profiler.trace_detached in Ihrem Trainingsskript.

Hinweis: Die trace_detached API ist nur in nächtlichen Builds oder nach Version 2.2 verfügbar. Wenn Sie in früheren „torch_xla“-Versionen programmatisch erfassen möchten, können Sie den Text der Methode kopieren, um den Hintergrunderfassungsthread direkt zu erstellen.

Wenn beispielsweise ein Profil in einer bestimmten Epoche und einem bestimmten Schritt automatisch erfasst werden soll, können Sie Ihr Trainingsskript so konfigurieren, dass die Umgebungsvariablen PROFILE_STEP, PROFILE_EPOCH und PROFILE_LOGDIR verwendet werden:

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 in dem Verzeichnis gespeichert, das von der Umgebungsvariablen PROFILE_LOGDIR angegeben wird.

Analyse in TensorBoard

Zur weiteren Analyse von Profilen können Sie TensorBoard mit dem TPU TensorBoard-Plug-in
auf derselben oder auf einem anderen Computer verwenden (empfohlen).

Wenn Sie TensorBoard auf einem Remote-Computer ausführen möchten, stellen Sie über SSH eine Verbindung zu diesem her und aktivieren Sie die 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, dass Sie auf dieser Maschine Profile unter ~/profiles/ haben). Wenn Sie die Profile in einem anderen Verzeichnis oder Cloud Storage-Bucket gespeichert haben, müssen Sie die Pfade richtig angeben, 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 in derselben verschachtelten Struktur gespeichert sein, die generiert wird, wenn Sie ein Profil erfassen. Beispiel: ~/profiles/plugins/profile/2023_04_10_21_40_22/localhost_9012.xplane.pb

Hinweis: Wenn Sie die Fehlermeldung „Doppelte Plug-ins“ erhalten, kann dies daran liegen, dass TensorBoard mehrere Pakete enthält. Deinstallieren Sie alle diese Pakete und führen Sie die pip install-Befehle noch einmal aus. Möglicherweise müssen Sie z. B. Folgendes ausführen:

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

TensorBoard ausführen

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

Informationen zu den TensorBoard-Tools und zur Interpretation der Ausgabe finden Sie unter TPU-Tools.