Modell auf Cloud TPU-VMs profilieren

Die Profilerstellung ist eines der wichtigsten Tools, mit denen Sie die Leistung Ihrer Modelle auf Cloud TPU optimieren können. Das Hauptprofilerstellungstool heißt XProf und ist im GitHub-Repository OpenXLA/XProf verfügbar. XProf unterstützt das Profiling aller XLA-basierten Frameworks, einschließlich JAX, Pytorch XLA und Tensorflow/Keras.

Profile erfassen

Um XProf zu verwenden, müssen Sie zuerst ein Profil Ihres Modells erfassen. Es gibt zwei Möglichkeiten, Profile zu erfassen:

  1. Programmatische Erfassung
  2. On-Demand-Erfassung (auch manuelle Erfassung)

Bei der programmatischen Erfassung müssen Sie Ihren Modellcode mit Anmerkungen versehen, um anzugeben, an welcher Stelle im Code Profile erfasst werden sollen. Normalerweise erfassen Sie ein Profil für einige Trainingsschritte oder für einen bestimmten Codeblock in Ihrem Modell. Es gibt verschiedene Möglichkeiten, Traces in den verschiedenen Frameworks JAX, Pytorch XLA und TensorFlow zu erfassen – entweder API-basiert oder kontextmanagerbasiert. In Frameworks auf höherer Ebene wie MaxText können Sie die Profilerstellung aktivieren, indem Sie beim Ausführen von MaxText einfach ein Flag profiler=xplane aktivieren.

Sie können die On-Demand-Profilerfassung verwenden, wenn Sie Profile ad hoc erfassen möchten oder die programmatische Profilerfassung nicht aktiviert haben. Das kann hilfreich sein, wenn während des Laufs ein Problem mit den Modellmesswerten auftritt und Sie Profile für einen bestimmten Zeitraum erfassen möchten, um das Problem zu diagnostizieren.

Wenn Sie die Profilerstellung auf Abruf aktivieren möchten, müssen Sie den XProf-Server weiterhin in Ihrem Code starten. Wenn Sie profiler.start_server aktivieren, wird ein XProf-Server für Ihre ML-Arbeitslast gestartet, der auf den On-Demand-Erfassungstrigger wartet, um mit der Erfassung von Profilen zu beginnen. Sie können die Erfassung von On-Demand-Profilen entweder über die TensorBoard-Benutzeroberfläche oder über die Befehlszeile mit dem XProfiler-Tool auslösen.

Informationen zum Aktivieren der programmatischen und On-Demand-Erfassung für verschiedene Frameworks finden Sie unter:

Sowohl für die programmatische als auch für die On-Demand-Profilerstellung müssen Sie den Speicherort für die erfassten Profile angeben. Sie können die Profile in einem Verzeichnis auf Ihrer TPU-VM mit einem Pfad wie /profiles/run-name oder in Cloud Storage mit einem Pfad wie gs://bucket-name/run-name/ speichern.

Bei den oben genannten Pfaden werden Ihre Profile unter folgenden Pfaden erfasst:

/profiles/run-name/plugins/profile/session1/

oder

gs://bucket-name/run-name/plugins/profile/session1/.

Eine Sitzung stellt eine Erfassung von Profilen während eines Laufs Ihres Modells dar. Wenn Sie beispielsweise in einem Trainingslauf Profile aus den Schritten 1 bis 3 und später Profile aus den Schritten 8 bis 10 erfassen, sind das Profile für denselben Lauf. Die erste Erfassung aus den Schritten 1 bis 3 ist session1 und die zweite Erfassung aus den Schritten 8 bis 10 ist session2.

Die verschiedenen Sitzungen werden unter jedem Lauf mit unterschiedlichen Zeitstempeln gekennzeichnet. Sie können Profile in verschiedenen Sitzungen entweder programmatisch, auf Anfrage oder in einer Kombination aus beidem erfassen. Das ist nützlich, wenn Sie Profile aus verschiedenen Teilen des Modelltrainings vergleichen möchten. Sie können beispielsweise Profile vom Beginn des Trainings mit Profilen vom Ende des Trainings vergleichen.

Profile mit TensorBoard visualisieren

Wenn Sie XProf-Profile in TensorBoard ansehen möchten, müssen Sie das Plug-in „tensorboard-plugin-profile“ installieren.

Damit Sie Ihre erfassten Profile in Google Cloudvisualisieren können, empfehlen wir, TensorBoard entweder auf einer TPU-VM oder das XProfiler-Tool zum Hosten von TensorBoard auf einer VM oder einem GKE-Pod zu verwenden.

Wenn Sie Ihre Profile auf der TPU-VM erfasst haben, können Sie sie mit TensorBoard auf dieser TPU-VM ansehen.

So installieren Sie XProf und TensorBoard auf einer TPU-VM:

pip install tensorboard_plugin_profile tensorboard
pip install xprof

So rufen Sie die mit TensorBoard erfassten Profile auf:

$ tensorboard --logdir=profiles/run-name

oder

$ tensorboard --logdir=profiles

Wenn Sie mehrere Profile aus mehreren Läufen laden möchten, verweisen Sie mit TensorBoard auf das Stammverzeichnis, das alle Ihre Läufe und Profile enthält, und nicht auf das Unterverzeichnis mit Profilen eines bestimmten Laufs.

Cloud-Diagnostics-XProf-Bibliothek zum Visualisieren von Tensorboard-Profilen auf Google Cloud

Die cloud-diagnostics-xprof (auch XProfiler Library genannt) erleichtert das Hosten von TensorBoard und das Visualisieren von Profilen auf Google Cloud. Weitere Informationen finden Sie im GitHub-Repository cloud-diagnostics-xprof.

Das Repository „cloud-diagnostics-xprof“ bietet die folgenden Verbesserungen gegenüber der lokalen Ausführung von XProf und TensorBoard:

  • Einrichtung und Verpackung von XProf- und TensorBoard-Abhängigkeiten
  • Speichern Sie Ihre Profile in Cloud Storage. Das kann für die langfristige Aufbewahrung und die Analyse nach dem Lauf nützlich sein. Lokale Profile werden gelöscht, nachdem der Researcher den Lauf beendet hat.
  • Schnelles Laden großer und mehrerer Profile durch Bereitstellung von TensorBoard auf einer Compute Engine-VM oder einem GKE-Pod. Der Maschinentyp kann je nach Nutzeranforderungen an Ladegeschwindigkeit und Kosten geändert werden.
  • Link zum Teilen von Profilen und zur Zusammenarbeit mit Teammitgliedern und Google-Entwicklern erstellen
  • Einfacheres On-Demand-Profiling von Arbeitslasten in GKE und Compute Engine, um einen beliebigen Host auszuwählen, auf dem Ihre Arbeitslast ausgeführt wird, um Profile zu erfassen

Bevor Sie die XProfiler-Bibliothek verwenden, müssen Sie Profile für Ihren Arbeitslastcode entweder programmatisch oder durch Starten des Profilservers erfassen, damit Sie später Profile bei Bedarf erfassen können.

Um die XProfiler-Bibliothek einzurichten, müssen Sie die Google Cloud CLI und eine virtuelle Python-Umgebung eingerichtet haben. Dann müssen Sie nur noch Folgendes ausführen:

pip install cloud-diagnostics-xprof

Dadurch werden alle erforderlichen XProf- und TensorBoard-Abhängigkeiten installiert.

Führen Sie als Nächstes den folgenden Befehl aus, um eine VM oder einen GKE-Pod zum Hosten von TensorBoard zu erstellen:

xprofiler create -z $ZONE -l $GCS_PATH

oder

xprofiler create --GKE -z $ZONE -l $GCS_PATH

Ersetzen Sie $ZONE durch eine beliebige Zone und $GCS_PATH durch den Pfad zu Ihren Profil-Traces. Sie können das Stammverzeichnis angeben, das alle Profil-Traces für mehrere Läufe enthält, oder einen bestimmten Satz von Profil-Traces für einen Lauf.

Angenommen, Sie erfassen Profile in den folgenden Verzeichnissen:

gs://bucket-name/run1/plugins/profile/session1/<profile.xplane.pb>
gs://bucket-name/run1/plugins/profile/session2/<profile.xplane.pb>
gs://bucket-name/run2/plugins/profile/session1/<profile.xplane.pb>

Sie können den GCS-Pfad auf das Stammverzeichnis (gs://bucket-name) festlegen:

xprofiler create -z $ZONE -l gs://bucket-name/

In der TensorBoard-UI werden alle Profile für alle Läufe und Sitzungen als run1/session1, run1/session2 und run2/session1 angezeigt.

Standardmäßig wird mit xprofiler create eine Compute Engine-VM erstellt, genauer gesagt eine c4-highmem-8-VM. Sie können den Maschinentyp mit dem Flag -m ändern. Wenn Sie eine TensorBoard-Instanz in einem GKE-Pod anstelle einer Compute Engine-VM erstellen möchten, können Sie das Flag –GKE an xprofiler create übergeben. Wenn Sie eine TensorBoard-Instanz in einem GKE-Pod hosten, kann es einfacher sein, TensorBoard zusammen mit dem Rest Ihrer in GKE bereitgestellten Arbeitslast zu verwalten.

Auf der VM oder dem GKE-Pod, auf dem TensorBoard gehostet wird, können große Profile und mehrere Profile auf Google Cloud viel schneller geladen werden als in einer lokal gehosteten TensorBoard-Instanz. Gemäß unseren Benchmarks werden Sie bei der ersten Last mit der Standard-VM c4-highmem-8 innerhalb weniger Minuten Profile mit einer Größe von etwa 1 GB sehen. Außerdem können Sie den richtigen Maschinentyp basierend auf Ihren Leistungs- und Kostenanforderungen auswählen.

Nachdem Sie xprofiler create ausgeführt haben, sollte die Ausgabe in etwa so aussehen:

Instance for gs://<bucket> has been created.
You can access it via the following:
1. https://<id>-dot-us-<region>.notebooks.googleusercontent.com.
2. xprofiler connect -z <zone> -l gs://bucket-name -m ssh
Instance is hosted at xprof-97db0ee6-93f6-46d4-b4c4-6d024b34a99f VM.

Die erste Option ist ein Link, den Sie anklicken können, um Ihre XProf-Profile in einem TensorBoard anzusehen. Über diesen Link können Sie Profile mit Ihrem Team und mit Google-Entwicklern teilen, die bei der Leistungsoptimierung auf Google Cloudhelfen.

Sie legen fest, wer Zugriff auf den Link hat, indem Sie die Berechtigungen für den Cloud Storage-Bucket festlegen, in dem Ihre Profildaten gespeichert sind. Falls der Link aus irgendeinem Grund nicht funktioniert, können Sie auch über SSH eine Verbindung zur TensorBoard-Instanz herstellen, um Ihre Profile mit dem Befehl xprofiler connect aufzurufen.

Wenn Sie den Cloud Profiler-Server in Ihrem Arbeitslastcode aktiviert haben (was sowohl für programmatische als auch für On-Demand-Erfassungen erforderlich ist) und On-Demand-Profiling durchführen möchten, haben Sie zwei Möglichkeiten:

a. Klicken Sie in TensorBoard auf die Schaltfläche Profil erfassen. Sie können jetzt einen beliebigen Gerätehost auswählen, auf dem Ihre Arbeitslast ausgeführt wird, um ein Profil zu erstellen. Wir unterstützen die On-Demand-Erfassung für Arbeitslasten, die in GKE oder Compute Engine ausgeführt werden.

b. Verwenden Sie den XProfiler-Erfassungsbefehl und geben Sie Informationen wie Zone, Cloud Storage-Bucket, Framework, Namen der Host-VM/des Host-Pods und Dauer der Erfassung in Millisekunden an. Das sind dieselben Informationen, die Sie in der TensorBoard-UI eingeben müssen.

Weitere Informationen zur Verwendung der cloud-diagnostics-xprof-Bibliothek finden Sie auf dieser GitHub-Seite.