Profil für Modell auf Cloud TPU-VMs erstellen

Die Profilerstellung ist eines der wichtigsten Tools zur Optimierung der Leistung Ihrer Cloud TPU-Modelle. Das zentrale Tool zur Profilerstellung heißt XProf und ist im GitHub-Repository OpenXLA/XProf verfügbar. XProf unterstützt die Profilerstellung für alle XLA-basierten Frameworks, darunter JAX, Pytorch XLA und Tensorflow/Keras.

Profile erfassen

Um XProf verwenden zu können, müssen Sie zunächst 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 versehen Sie Ihren Modellcode mit Anmerkungen und geben damit an, an welcher Stelle im Code Profile erfasst werden sollen. In der Regel erfassen Sie ein Profil für einige Trainingsschritte oder für einen bestimmten Codeblock in Ihrem Modell. Die Frameworks JAX, Pytorch XLA und TensorFlow bieten mehrere Möglichkeiten zur Erfassung von Traces – API-basiertes Starten oder Beenden von Traces oder kontextmanagerbasiert. In übergeordneten Frameworks wie MaxText können Sie die Profilerfassung aktivieren, indem Sie beim Ausführen von MaxText einfach das Flag profiler=xplane aktivieren.

Verwenden Sie die On-Demand-Profilerfassung, wenn Sie Profile ad hoc erfassen möchten oder wenn Sie die programmatische Profilerfassung nicht aktiviert haben. Dies kann hilfreich sein, wenn Sie während der Ausführung ein Problem mit Ihren Modellmesswerten feststellen und sofort Profile für einen bestimmten Zeitraum erfassen möchten, um das Problem zu diagnostizieren.

Um die On-Demand-Profilerfassung zu aktivieren, müssen Sie den XProf-Server weiterhin in Ihrem Code starten. Durch Aktivieren von profiler.start_server wird ein XProf-Server für Ihre ML-Arbeitslast gestartet. Dieser wartet auf den Trigger für die On-Demand-Erfassung, der die Erfassung von Profilen startet. Sie können die On-Demand-Profilerfassung entweder über die Tensorboard-Benutzeroberfläche oder mit dem XProfiler-Tool über die Befehlszeile auslösen.

Informationen dazu, wie Sie die programmatische und die On-Demand-Erfassung für verschiedene Frameworks aktivieren können, finden Sie unter:

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

Mit den oben genannten Pfaden werden Ihre Profile unter folgendem Pfad erfasst:

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

oder

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

Eine Sitzung entspricht einer Erfassung von Profilen während einer Ausführung Ihres Modells. Wenn Sie beispielsweise in einem Trainingslauf Profile aus den Schritten 1 bis 3 und später Profile aus den Schritten 8 bis 10 erfassen, gehören diese Profile zum selben Lauf, aber 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 Datumsstempeln gekennzeichnet. Sie können Profile in verschiedenen Sitzungen entweder programmgesteuert, on demand oder kombiniert erfassen. Das ist nützlich, wenn Sie Profile aus verschiedenen Teilen des Modelltrainings vergleichen möchten. Beispielsweise können Sie 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.

Zur Visualisierung Ihrer erfassten Profile in Google Cloud, empfehlen wir Ihnen, TensorBoard entweder auf einer TPU-VM zu verwenden oder das XProfiler-Tool zu nutzen, mit dem TensorBoard auf einer VM oder einem GKE-Pod gehostet werden kann.

Wenn Sie Ihre Profile auf der TPU-VM erfasst haben, können Sie die erfassten Profile mit TensorBoard auf dieser TPU-VM aufrufen.

Mit den folgenden Befehlen 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 in Google Cloud

Mit der Bibliothek cloud-diagnostics-xprof (auch XProfiler) ist es einfacher, TensorBoard zu hosten und Profile in Google Cloudzu visualisieren. Weitere Informationen finden Sie im GitHub-Repository „cloud-diagnostics-xprof“.

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

  • Einrichtung und Paketerstellung von XProf- und TensorBoard-Abhängigkeiten
  • Speichern von Profilen in Cloud Storage. Das kann für die langfristige Aufbewahrung und die Analyse nach dem Lauf nützlich sein. Lokal erfasste Profile werden nach Beendigung des Laufs gelöscht.
  • 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 den Anforderungen des Nutzers hinsichtlich Ladegeschwindigkeit und Kosten geändert werden.
  • Erstellung eines Links zum Teilen von Profilen und zur Zusammenarbeit mit Teammitgliedern und Google-Entwicklern
  • Einfachere On-Demand-Profilerstellung von Arbeitslasten auf der GKE und Compute Engine: Sie können einen beliebigen Host auswä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 eine On-Demand-Profil-Erfassung durchführen können.

Um die XProfiler-Bibliothek einzurichten, müssen Sie die Google Cloud CLI und eine virtuelle Python-Umgebung eingerichtet haben. Anschließend müssen Sie nur noch den folgenden Befehl ausführen:

pip install cloud-diagnostics-xprof

Damit 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.

Beispiel: 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 legen den GCS-Pfad auf das Stammverzeichnis (gs://bucket-name) fest:

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

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

Standardmäßig erstellt xprofiler create eine Compute Engine-VM, 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. Das Hosten einer TensorBoard-Instanz auf einem GKE-Pod kann die Verwaltung von TensorBoard zusammen mit dem Rest Ihrer auf der GKE bereitgestellten Arbeitslast vereinfachen.

Große und mehrere Profile können viel schneller in Google Cloud geladen werden, wenn TensorBoard nicht lokal, sondern auf einer VM oder einem GKE-Pod gehostet wird. Unseren Benchmarks zufolge werden Profile mit einer Größe von etwa 1 GB bei der ersten Ladung mit der Standard-VM c4-highmem-8 innerhalb weniger Minuten geladen. Darüber hinaus können Sie den richtigen Maschinentyp entsprechend Ihren Anforderungen an Leistung und Kosten auswählen.

Nachdem Sie xprofiler create ausgeführt haben, wird eine Ausgabe angezeigt, die in etwa so aussieht:

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. Sie können darauf klicken, um Ihre XProf-Profile auf einem TensorBoard aufzurufen. Über diesen Link können Sie Profile mit Ihrem Team und auch mit Google-Entwicklern teilen, die bei der Leistungsoptimierung in Google Cloudhelfen.

Sie können anhand der Berechtigungen für den Cloud Storage-Bucket, in dem Ihre Profildaten gespeichert sind, steuern, wer Zugriff auf den Link hat. Falls der Link nicht funktioniert, bieten wir auch eine Möglichkeit, über SSH eine Verbindung zur TensorBoard-Instanz herzustellen und Ihre Profile mit dem Befehl xprofiler connect aufzurufen.

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

a. Klicken Sie in TensorBoard auf die Schaltfläche Profil erfassen. Wählen Sie nun ein beliebiges Gerät aus, auf dem Ihre Arbeitslast ausgeführt wird, um ein Profil zu erstellen. Wir unterstützen die On-Demand-Erfassung für Arbeitslasten, die in der GKE oder Compute Engine ausgeführt werden.

b. Verwenden Sie den Befehl XProfiler-und geben Sie Informationen wie Zone, Cloud Storage-Bucket, Framework, Namen der Host-VM/des Host-Pods und Dauer der Erfassung in Millisekunden an. Diese Informationen müssen Sie auch in der TensorBoard-Benutzeroberfläche eingeben.

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