Profil für PyTorch-XLA-Arbeitslasten erstellen
Die Profilerstellung ist eine Möglichkeit, die Leistung von Modellen zu analysieren und zu verbessern. Obwohl dies viel mehr ist, kann es hilfreich sein, sich die Profilerstellung als das Timing von Vorgängen und Codeteilen vorzustellen, die auf beiden Geräten (TPUs) und Hosts (CPUs) ausgeführt werden. Dieser Leitfaden bietet einen kurzen Überblick über die Erstellung von Codeprofilen für Training oder Inferenz. Weitere Informationen zum Analysieren generierter Profile finden Sie in den folgenden Anleitungen.
- Fehlerbehebung bei der Leistung von PyTorch XLA auf TPU-VMs – Teil 1
- Fehlerbehebung bei der Leistung von PyTorch XLA auf TPU-VMs – Teil 2
- Fehlerbehebung bei der Leistung von PyTorch XLA auf TPU-VMs – Teil 3
Jetzt starten
TPU erstellen
Umgebungsvariablen exportieren:
$ 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 jede TPU-Version finden Sie unter TPU-Versionen.
version
- Die Version der Cloud TPU-Laufzeit. In der exportierten Variablen wird eine Standardeinstellung angezeigt. Sie können aber auch eine aus der Liste der unterstützten Konfigurationen verwenden.
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
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
Ein Profil kann manuell über capture_profile.py
oder programmatisch aus dem Trainingsskript mithilfe der torch_xla.debug.profiler
APIs erfasst werden.
Starten des Profilservers
Damit ein Profil erfasst werden kann, muss im Trainingsskript ein Profilserver ausgeführt werden. Starten Sie einen Server mit einer Portnummer Ihrer Wahl, z. B. 9012
, wie im folgenden Befehl gezeigt.
import torch_xla.debug.profiler as xp server = xp.start_server(9012)
Der Server kann direkt am Anfang der main
-Funktion gestartet werden.
Sie können jetzt Profile wie im folgenden Abschnitt beschrieben erfassen. Das Skript erstellt ein Profil für alle Vorgänge auf einem TPU-Gerät.
Traces hinzufügen
Wenn Sie auch ein Profil für Vorgänge auf dem Hostcomputer erstellen möchten, können Sie xp.StepTrace
oder xp.Trace
in Ihren Code einfügen. Diese Funktionen verfolgen den Python-Code auf dem Hostcomputer.
Damit wird gemessen, wie lange es dauert, den Python-Code auf dem Host (CPU) auszuführen, bevor der "Graph" an das TPU-Gerät übergeben wird. Daher ist es vor allem nützlich, um den Tracing-Overhead zu analysieren. Sie können dies in die Trainingsschleife einfügen, in der der Code Batches von Daten 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 = ...
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 die Trainingsschleife einfügen.
Sie können nach der ersten Kompilierung die Geräteaktivität erfassen. Warten Sie, bis das Modell mit den Trainings- oder Inferenzschritten beginnt.
Manuelle Erfassung
Mit dem Skript capture_profile.py
aus dem Pytorch XLA-Repository kann ein Profil schnell erfasst werden. Kopieren Sie dazu die Erfassungsprofildatei direkt auf Ihre 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 den 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, 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 ein Profil automatisch ausgelöst wird. Dazu verwenden Sie die API torch_xla.debug.profiler.trace_detached in Ihrem Trainingsskript.
Wenn Sie beispielsweise ein Profil in einer bestimmten Epoche und bei 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 in dem Verzeichnis gespeichert, das in der Umgebungsvariable PROFILE_LOGDIR
angegeben ist.
Analyse in TensorBoard
Für die weitere Analyse von Profilen können Sie TensorBoard mit dem TPU TensorBoard-Plug-in
entweder auf demselben oder einem anderen Computer verwenden (empfohlen).
Wenn Sie TensorBoard auf einem Remote-Computer ausführen möchten, stellen Sie über SSH eine Verbindung zu ihm 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 (sofern Sie auf dieser Maschine unter ~/profiles/
Profile haben). 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
(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.
Weitere Informationen zu den TensorBoard-Tools und zur Interpretation der Ausgabe finden Sie unter TPU-Tools.