JAX-Arbeitslasten zu Pathways portieren

Aufgrund der verteilten Natur von JAX mit Pathways lassen sich einige Vorgänge aufgrund des Kommunikationsaufwands möglicherweise nicht gut skalieren. Pathways minimiert diese Overheads mit Funktionen wie dem asynchronen Versand. Es gibt jedoch einige Dinge, die Sie beachten müssen, wenn Sie JAX-Arbeitslasten zu Pathways portieren oder eine JAX-Arbeitslast mit Pathways auf eine große Anzahl von Beschleunigern skalieren.

Hinweis

Sie benötigen Folgendes:

Prozessindex

JAX mit Pathways behandelt alle Geräte in Ihrem Pathways-Cluster als lokal. Dies vereinfacht die Geräteverwaltung und ermöglicht es JAX, alle verfügbaren Ressourcen zu nutzen. In der Praxis bedeutet das:

  • jax.process_index() ist für alle Geräte immer 0.
  • jax.devices() und jax.local_devices() geben alle TPU-Geräte im gesamten Job zurück.

Hardwaretyp und Colocation

Für eine optimale Leistung sollten Sie alle Pathways-Komponenten und den Nutzerjob in derselben Google Cloud -Cloud-Zone platzieren. Verwenden Sie eine große CPU wie den IFRT-Proxy und den Ressourcenmanager. Wir empfehlen mindestens einen dedizierten n2-standard-64 mit 64 vCPUs und 256 GB Arbeitsspeicher.

PathwaysUtils

Pathways-utils ist ein Python-basiertes GitHub-Repository mit wichtigen Dienstprogrammen und Tools, mit denen Sie die Bereitstellung und Ausführung von JAX-Arbeitslasten in der Pathways on Cloud-Architektur optimieren können. Dieses Paket übernimmt die erforderlichen Anpassungen für die Cloud-Umgebung, sodass sich JAX-Entwickler mit minimaler plattformspezifischer Konfiguration auf ihre Kernworkflows für maschinelles Lernen konzentrieren können. Konkret bietet sie Folgendes:

  • Ein „Proxy“-JAX-Backend: Mit diesem benutzerdefinierten Backend kann Ihre JAX-Anwendung die Pathways-Infrastruktur verwenden, indem die Umgebungsvariable JAX_PLATFORMS=proxy festgelegt wird.
  • Integrierte Profilerstellungstools: Profilerstellungsfunktionen, mit denen Sie die Leistung Ihrer Anwendung nachvollziehen können. Mit Standard-JAX-Profilerstellungs-APIs wie jax.profiler.start_trace und jax.profiler.start_server können Sie nicht nur Ihren JAX-Code, sondern auch die zugrunde liegenden Pathways-Komponenten profilieren. So erhalten Sie einen ganzheitlichen Überblick über die Ausführung in der Cloud-Umgebung.
  • Verteiltes Checkpointing mit Orbax: Ein benutzerdefinierter Orbax-Checkpoint-Handler, mit dem Sie verteilte Checkpoints verwenden und Ihre Checkpoints wiederherstellen können, wenn Sie die Orbax-Bibliothek in der Pathways-Umgebung verwenden. Diese Integration soll ohne Änderungen an Ihrem vorhandenen Orbax-Prüfpunktcode funktionieren, solange pathwaysutils importiert wird.
  • Primitive für elastisches Training: Bietet grundlegende Primitive für elastisches Training, mit denen Sie robuste und skalierbare Trainingsworkflows mit Pathways erstellen können. Mit diesen Primitiven können sich Ihre Trainingsjobs dynamisch an Änderungen der verfügbaren Ressourcen anpassen, was die Effizienz und Ausfallsicherheit in Cloud-Umgebungen verbessert.

Prüfpunkte

Orbax wurde mit Pathways für die verteilte Prüfpunkterstellung und Wiederherstellung mit Cloud Storage gründlich getestet. Wenn Sie die Umgebungsvariable ENABLE_PATHWAYS_PERSISTENCE=1 festlegen und import pathwaysutils; pathwaysutils.initialize() in Ihrem train.py aufrufen, wird ein benutzerdefinierter ArrayHandler registriert, der Prüfpunktvorgänge effizient über den IFRT-Proxy verarbeitet. So können Pathways-Worker auf Beschleunigern Daten direkt speichern und wiederherstellen.

Zusammengeführte Python-Anwendung

Colocated Python ist eine Open-Source-JAX-API, mit der Sie benutzerdefinierten Python-Code direkt auf den TPU- oder GPU-Hosts ausführen können. Dies ist in JAX mit mehreren Controllern einfacher. Mit diesem Ansatz können Sie die Datenübertragung zwischen dem Client und den TPU-Maschinen für rechenintensive Aufgaben wie das Laden von Daten und das Erstellen von Prüfpunkten vermeiden. Wenn Sie Ihren Pathways-Cluster für die Ausführung der colocated Python JAX API konfigurieren möchten, erstellen Sie ein Sidecar-Container-Image mit dem Beispiel-Sidecar-Dockerfile und übertragen Sie es per Push.

Laden der Daten

Während des Trainings werden wiederholt Batches aus einem Dataset geladen, um sie in das Modell einzuspeisen. Ein effizienter, asynchroner Data Loader, der den Batch auf Hosts verteilt, ist wichtig, um unterausgelastete Beschleuniger zu vermeiden. Wenn Sie ein Training mit Pathways ausführen, wird der Daten-Loader auf einer CPU-VM ausgeführt (im Gegensatz zu einer TPU-VM, die bei Setups mit mehreren Controllern verwendet wird) und sendet Daten an TPU-VMs. Dies führt zu einer höheren Latenz beim Lesen von Daten, die jedoch teilweise durch das Vorauslesen einer bestimmten Anzahl von Batches auf dem CPU-Host und das asynchrone Senden der gelesenen Daten an die TPUs verringert wird. Diese Lösung ist für kleine bis mittlere Umgebungen ausreichend.

Wenn Sie einen Multiprocess-Datenlader (z. B. Grain) im Head-Arbeitslastcontainer in GKE ausführen, teilt sich der Container den Head-Knoten mit den Sidecars pathways-proxy und pathways-rm. Aufgrund dieser gemeinsam genutzten Architektur müssen Sie Folgendes konfigurieren:

  • Explizite Worker-Anzahl festlegen: In einem Container meldet os.cpu_count() die Gesamtzahl der vCPUs der Host-VM und nicht das CPU-Limit des Containers. Übergeben Sie eine explizite Worker-Anzahl (nicht mehr als das CPU-Limit des Containers), anstatt sie von os.cpu_count() abzuleiten.
  • Speichergestütztes /dev/shm-Volume bereitstellen: Multiprozess-Datenlader verwenden POSIX-Shared-Memory (/dev/shm) für die prozessübergreifende Kommunikation. Achten Sie darauf, dass ein speichergestütztes emptyDir-Volume (medium: Memory) unter /dev/shm bereitgestellt ist, damit die im Voraus abgerufenen Batches das Standardlimit für den gemeinsam genutzten Arbeitsspeicher des Containers nicht überschreiten.

Für eine optimale Leistung im großen Maßstab empfehlen wir dringend, die Eingabedatenpipeline mit colocated Python zu platzieren, um die Datenpipeline direkt auf den Beschleunigern auszuführen. Dadurch wird der CPU-Engpass beseitigt und die schnellen Verbindungen der TPU für die Datenübertragung genutzt.

Eine Referenzimplementierung für die Migration einer TFDS-basierten Eingabepipeline finden Sie in der RemoteIterator-Implementierung in multihost_dataloading.py. Diese Implementierung funktioniert sowohl auf JAX mit mehreren Controllern als auch auf Pathways verteilt mit der colocated Python JAX API.

Jax-Versionsverwaltung

Pathways-Releases sind eng mit JAX-Versionen gekoppelt, um Kompatibilität und Stabilität zu gewährleisten. Prüfen Sie, ob Ihre Pathways-Artefakte und Ihre JAX-Version übereinstimmen, um potenzielle Probleme zu vermeiden. In jedem Pathways-Release werden die kompatiblen JAX-Versionen durch ein Tag im Format jax-<version> angegeben.

Kompilierungs-Cache

Der persistente Kompilierungscache von Pathways ist eine Funktion, mit der Pathways-Server kompilierte ausführbare XLA-Dateien an einem persistenten Speicherort wie Cloud Storage speichern können, um redundante Kompilierungen zu vermeiden. Dieses Feature ist standardmäßig aktiviert. Der Speicherort des Cache wird als --gcs_scratch_location-Flag an den Ressourcenmanager und die Pathways-Worker-Container übergeben. Damit die damit verbundenen Speicherkosten so gering wie möglich bleiben, wird dem Cloud Storage-Speicherort eine Lebenszyklusrichtlinie zugewiesen. Pro Cloud Storage-Bucket gilt ein Limit von 50 Richtlinien. Daher empfehlen wir, für alle Arbeitslasten einen gemeinsamen Cloud Storage-Speicherort zu verwenden.

Wenn Sie das Flag --gcs_scratch_location weglassen oder leer lassen, greift Pathways automatisch auf einen In-Memory-Kompilierungs-Cache zurück. In diesem Cache werden kompilierte XLA-Ausführungsdateien nicht über Pathways-Bereitstellungen hinweg persistent gespeichert.

Dieser Cache ähnelt dem JAX-Kompilierungscache, der für Pathways-Arbeitslasten durch pathwaysutils.initialize() deaktiviert wird.

Die folgenden Cloud Storage-Berechtigungen sind für den Kompilierungscache erforderlich:

  • storage.buckets.get: Zum Abrufen von Bucket-Metadaten.
  • storage.buckets.update: Erforderlich, damit Pathways Richtlinien für den Objektlebenszyklus einrichten kann, um die TTL für die Cache-Entfernung zu erzwingen.
  • storage.objects.list: Zum Auflisten vorhandener Cache-Objekte im Bucket.
  • storage.objects.create: Zum Schreiben neuer kompilierter ausführbarer Dateien in den Cache.
  • storage.objects.get: Zum Lesen im Cache gespeicherter ausführbarer Dateien aus dem Bucket.

Profilerstellung

Mit dem JAX-Profiler können Sie Traces eines JAX-Programms generieren. Es gibt zwei gängige Methoden, die von Pathways unterstützt werden:

  • Programmatisch
    • Profile programmatisch aus Ihrem JAX-Code erfassen
  • Manuell
    • Profile nach dem Starten des Profiler-Servers über Ihren JAX-Code on demand erfassen

In beiden Fällen werden die Profile in einen Cloud Storage-Bucket geschrieben. Im Cloud Storage-Bucket werden mehrere Trace-Dateien erstellt, möglicherweise in verschiedenen Zeitstempelordnern, z. B.:

  • Haupt-Python-Prozess, der den Trace aufgerufen hat (normalerweise Ihre Notebook-VM): <jax-client-vm-name>.xplane.pb
  • Pathways-IFRT-Proxy: client.<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
  • Pathways Resource Manager: server.*<var>PATHWAYS_HEAD_NODE_NAME</var>.xplane.pb
  • Pathways-Worker: server.*<tpu-node-name>.xplane.pb

Diese Trace-Dateien können mit TensorBoard analysiert werden. Führen Sie dazu den folgenden Befehl aus. Weitere Informationen zu TensorBoard und allen Profilerstellungstools finden Sie unter TensorFlow-Leistung mit dem Profiler optimieren.

# verify trace files are present
gsutil ls -l -r gs://BUCKET/PREFIX

# View on tensorboard
tensorboard --logdir=gs://BUCKET/PREFIX

Ersetzen Sie Folgendes:

  • BUCKET : ein Cloud Storage-Bucket zum Speichern der Trace-Dateien
  • PREFIX: ein Pfad in Ihrem Cloud Storage-Bucket, in dem die Trace-Dateien gespeichert werden sollen

Programmatische Profilerfassung

Erfassen Sie ein Profil aus Ihrem Code. Die Profile werden in gs://<bucket>/<prefix> in einem Zeitstempelverzeichnis gespeichert.

import jax
import pathwaysutils

pathwaysutils.initialize()

jax.profiler.start_trace("gs://BUCKET/PREFIX")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

Manuelle Profilerfassung

Wenn Sie Profilinformationen manuell erfassen möchten, müssen Sie den Profiler-Server über Ihren Python-Code starten:

import jax
import pathwaysutils

pathwaysutils.initialize()
jax.profiler.start_server(jax_profiler_port)

# Your JAX code
jax.profiler.stop_server() # this is functionally a no-op

Während der Profiler-Server ausgeführt wird, können Sie ein Profil erfassen und die Daten in den Cloud Storage-Zielspeicherort exportieren:

export DURATION_IN_SECS=6
export OUTPUT_DIR=gs://<var>BUCKET</var>/<var>PREFIX</var>

curl -d "{\"duration_ms\":\"${DURATION_IN_SECS} * 1000 }}\", \"repository_path\":\"${OUTPUT_DIR}\"}" -H "Content-Type: application/json" -X POST http://localhost:<jax_profiler_port>

Informationen zum Timing für IFRT-Proxy-Client-Methoden wie Compile und Execute finden Sie im Trace Ihres Programms. Diese Ereignisse, die die Interaktionen mit dem IFRT-Proxy-gRPC-Server während der Kompilierung und Ausführung beschreiben, werden im Thread mit dem Namen GrpcClientSessionUserFuturesWorkQueue angezeigt. Wenn Sie diesen Thread in Ihrem Trace untersuchen, können Sie Einblicke in die Leistung dieser Vorgänge erhalten.

XLA-Flags

Wenn Sie Pathways verwenden, müssen Sie die XLA-Flags im pathways-proxy-Container festlegen.

Wenn Sie das Cluster Toolkit (den Befehl gcluster job submit) verwenden, legen Sie XLA-Flags mit --pathways-proxy-args fest:

--pathways-proxy-args="--xla_flag_1=value1 --xla_flag_2=value2"

HLO-Dump

Wenn Sie die HLO-Eingaben (High Level Optimizer) prüfen möchten, die an den XLA-Compiler übergeben werden, können Sie Pathways so konfigurieren, dass die HLO mit --pathways-proxy-args an einem bestimmten Cloud Storage-Speicherort abgelegt wird:

--pathways-proxy-args="--xla_dump_to=gs://BUCKET_NAME/HLO_PREFIX/"

Nächste Schritte