PyTorch XLA-Arbeitslasten profilieren

Mit dem Profiling können Sie die Leistung von Modellen analysieren und verbessern. Das Profilerstellungsverfahren ist zwar noch viel komplexer, aber manchmal hilft es, sich das Timing von Vorgängen und Codeteilen vorzustellen, die sowohl auf Geräten (TPUs) als auch auf Hosts (CPUs) ausgeführt werden. 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 Variablen wird ein Standardwert angezeigt. Sie können aber auch einen 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 umfasst 4 TPU-Geräte. Jeder Mitarbeiter profiliert nur ein TPU-Gerät aus den vier verfügbaren Geräten. Das Profiling aller vier TPU-Geräte für jeden Worker 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.

Profilserver starten

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 Script erstellt ein Profil für alles, was auf einem TPU-Gerät passiert.

Traces hinzufügen

Wenn Sie auch Vorgänge auf dem Hostcomputer erfassen möchten, können Sie xp.StepTrace oder xp.Trace in Ihren Code einfügen. Diese Funktionen zeichnen den Python-Code auf dem Hostcomputer auf. 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 dies in die Trainingsschleife einfügen, in der der Code Datensätze verarbeitet, z. B.

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

oder einzelne Codeteile in

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

Wichtig :Umschließen Sie xm.mark_step() nicht mit xp.Trace. 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 beim Verlassen des Gültigkeitsbereichs automatisch xm.mark_step() hinzugefügt wird. Außerdem kann dieser Fehler auftreten, wenn Sie torchdynamo verwenden, da torchdynamo intern xm.mark_step aufruft. Die aktuelle Umgehungslösung besteht darin, xp.Trace nicht zu verwenden.

Wenn Sie Lighting verwenden, können Sie das Hinzufügen von Traces überspringen, da dies in einigen Codeteilen 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

Mit dem capture_profile.py-Script aus dem Pytorch XLA-Repository können Sie schnell ein Profil erstellen. Kopieren Sie dazu die Datei mit dem Aufnahmeprofil direkt auf Ihre TPU-VM. Mit dem folgenden Befehl wird es 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 das Protokollverzeichnis ~/profiles/ an einen beliebigen Speicherort und mit einem beliebigen Namen ändern. Es ist auch möglich, die Daten direkt im Cloud Storage-Bucket zu speichern. Legen Sie dazu logdir auf gs://your_bucket_name/ fest.

Programmatisch erfassen

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

Hinweis: Die trace_detached API ist nur in Nightly-Builds oder nach der 2.2-Version verfügbar. Wenn Sie in früheren Versionen von torch_xla programmatisch erfassen möchten, können Sie den Methodenkörper kopieren, um den Hintergrund-Aufnahme-Thread direkt zu erstellen.

Wenn Sie beispielsweise ein Profil bei einer bestimmten Epoche und einem bestimmten Schritt automatisch erfassen möchten, 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 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 Remotecomputer ausführen möchten, stellen Sie eine SSH-Verbindung 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 die erforderlichen Pakete auf dem Remote-Computer und starten Sie TensorBoard. Angenommen, Sie haben Profile auf diesem Computer unter ~/profiles/ gespeichert. Wenn Sie die Profile in einem anderen Verzeichnis oder Cloud Storage-Bucket gespeichert haben, müssen Sie die Pfade korrekt 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 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 http://localhost:6006/ auf und wählen Sie im Drop-down-Menü PROFILE aus, um Ihre Profile zu laden.

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