Einführung in Pathways in der Cloud

Pathways ist ein System, das die Erstellung von umfangreichen, dünnbesetzten ML-Systemen (maschinelles Lernen) mit mehreren Aufgaben ermöglicht. Sie können Tausende oder Zehntausende von Beschleunigern verwenden und je nach Verarbeitungsanforderungen dynamisch unterschiedliche Rechenmengen für verschiedene Aufgaben zuweisen.

Pathways vereinfacht umfangreiche ML-Berechnungen, da ein einzelner JAX-Client Arbeitslasten über große TPU-Slices hinweg orchestrieren kann, die potenziell Tausende von TPU-Chips umfassen.

Google verwendet Pathways intern, um große Modelle wie Gemini zu trainieren. Pathways on Cloud bietet Kunden von Google Cloud dieselben Vorteile.

Hinweis

Hinweis

Sie können Pathways in der Google Kubernetes Engine (GKE) oder in einem Cluster Director-Cluster mit Slurm ausführen. Wählen Sie den Tab für Ihre Plattform aus.

GKE

Erfüllen Sie die folgenden Voraussetzungen:

In diesem Dokument wird beschrieben, wie Sie von Pathways verwaltete TPUs in GKE für Batch-, Echtzeit- und interaktive Arbeitslasten verwenden. Es wird davon ausgegangen, dass Sie bereits mit der Verwendung von TPUs in GKE vertraut sind, einschließlich Single-Slice- und Multi-Slice-TPUs in GKE, und dass Sie allgemeine Erfahrung mit Multi-Slice-TPUs haben.

Cluster Director

Erfüllen Sie die folgenden Voraussetzungen:

In diesem Dokument wird beschrieben, wie Sie Pathways mit TPUs in einem Cluster Director-Cluster für Batch- und interaktive Arbeitslasten verwenden. Es wird davon ausgegangen, dass Sie bereits mit Cluster Director, TPU-Maschinen in Cluster Director und dem Ausführen von Jobs mit Slurm vertraut sind und allgemeine Erfahrung mit Multislice-TPUs haben.

Informationen zum Erstellen eines Clusters finden Sie unter Cluster Director-Cluster mit Pathways erstellen.

Einzelner Controller und mehrere Controller

Es gibt zwei Hauptmethoden, um Berechnungen auf mehreren Geräten zu verwalten und zu orchestrieren:

Funktion Einzelner Controller (Pathways) Multi-Controller (JAX-Standard)

Kontrolle

Zentraler Kontrollpunkt: Ein einzelnes Clientprogramm fungiert als zentraler Controller.

Verteilte Steuerung: Es sind mehrere Prozesse beteiligt, die jeweils eine eigene Python-Interpreter-Instanz haben.

Ansicht

Einheitliche Ansicht: Der Client sieht alle Geräte als ein einziges, einheitliches System.

Lokalisierte Ansicht: Jeder Python-Prozess sieht nur die mit ihm verbundenen Geräte.

Programmierung

Vereinfachte Programmierung: Sie interagieren mit einem einzelnen Client, sodass das System als eine einzige große Maschine mit vielen lokalen Beschleunigern erscheint.

Single Program, Multiple Data (SPMD): Verwendet hauptsächlich das SPMD-Paradigma, bei dem alle Geräte dasselbe Programm ausführen müssen.

Flexibilität

Unterstützt komplexe Berechnungsmuster über SPMD hinaus, einschließlich asymmetrischer Pipeline-Parallelität und dünnbesetzter Berechnungen.

Kann bei der Ressourcenverwaltung weniger flexibel sein, insbesondere bei verschiedenen TPU-Slices.

Pathways-Komponenten

In den folgenden Abschnitten werden die Hauptkomponenten der Pathways-Architektur beschrieben.

Pathways-Ressourcenmanager

Diese Komponente ist die zentrale Steuerungsebene des Pathways-Systems. Er verwaltet alle Beschleunigerressourcen und koordiniert die Zuweisung von Beschleunigern für Ihre Jobs. Er überwacht den Zustand der Worker und übernimmt die Planung, das Pausieren und das Fortsetzen von Jobs. Es dient als zentrale Anlaufstelle für Fehler und den Systemstatus. Diese Komponente benötigt nur CPU-Ressourcen.

Pathways-Client

Diese Komponente ist eine Implementierung der Interim Framework Runtime (IFRT), die als Einstiegspunkt in das Pathways-System dient. Er empfängt High-Level-Vorgänge (HLOs) von Ihrem Programm. Der Pathways-Client stimmt sich mit dem Pathways-Ressourcenmanager ab, um anhand Ihres Codes zu bestimmen, wo kompilierte Programme zur Ausführung platziert werden sollen. Sie bietet einem bestimmten JAX-Client eine einheitliche Ansicht des Systems. Diese Komponente benötigt nur CPU-Ressourcen.

Pathways-Mitarbeiter

Pathways-Worker-Prozesse werden auf den Beschleunigermaschinen (TPU-VMs) ausgeführt. Sie erhalten kompilierte ausführbare Dateien Ihres Programms vom IFRT-Proxyserver und führen die Berechnungen auf den TPUs aus. Pathways-Worker senden Daten über den IFRT-Proxyserver an Ihr Programm zurück. Für diese Komponente sind Beschleunigerressourcen erforderlich.

IFRT-Proxy-Client

Der IFRT-Proxy-Client ist eine Open-Source-Implementierung der Interim Framework Runtime (IFRT)-API, die den Nutzercode von der zugrunde liegenden Laufzeit entkoppelt und die Codeportabilität und -transparenz verbessert. JAX verwendet diese Implementierung als Alternative zur standardmäßigen Multi-Controller-Laufzeit. Der IFRT-Proxy-Client fungiert als Kommunikationsbrücke zwischen Ihrem Programm und den Pathways-Komponenten. Sie sendet Anfragen an den IFRT-Proxyserver und empfängt Ergebnisse von ihm. Diese Komponente benötigt nur CPU-Ressourcen.

IFRT-Proxyserver

Der IFRT-Proxyserver ist ein gRPC-Server, der Anfragen vom IFRT-Proxyclient empfängt und an den Pathways-Client weiterleitet, der die Verteilung der Arbeit übernimmt. Diese Komponente benötigt nur CPU-Ressourcen.

Sidecar-Server

Der Sidecar-Server ist ein gRPC-Server, der sich zusammen mit dem Pathways-Worker auf der Beschleuniger-VM befindet. Der von Nutzern angegebene Python-Code wird direkt auf der Beschleuniger-VM ausgeführt, um die Latenz bei der Datenübertragung zwischen dem Controller und den Beschleunigern zu reduzieren. Der Sidecar-Server interagiert mit dem Pathways-Worker über ein benutzerdefiniertes versioniertes Protokoll im gRPC-Transport.

Architekturdiagramm, das zeigt, wie der JAX-Client, der IFRT-Proxy, der Pathways-Ressourcenmanager und die Pathways-Worker interagieren.
Pathways-Komponenten

Pathways-Komponenten in GKE

In diesem Abschnitt werden Pathways-Komponenten GKE-Komponenten wie Containern und Pods zugeordnet.

Pathways-Container-Images finden Sie an den folgenden Speicherorten:

Containertyp Standort

IFRT-Proxyserver

us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:jax-JAX_VERSION

Pathways-Ressourcenmanager/-Worker

us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:jax-JAX_VERSION

Pathways-Ressourcenmanager

Nachdem Sie einen GKE-Cluster erstellt haben, können Sie den Pathways-Ressourcenmanager mit dem folgenden containerSpec bereitstellen:

  - name: pathways-rm
    image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:latest
    imagePullPolicy: Always
    env:
    - name: HOST_ADDRESS
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    - name: TPU_SKIP_MDS_QUERY
      value: "true"
    args:
    - --server_port=29001
    - --node_type=resource_manager
    - --instance_count=WORKLOAD_NODEPOOL_COUNT
    - --instance_type=SLICE_TOPOLOGY
    - --gcs_scratch_location=gs://BUCKET_NAME

Der Container verwendet die folgenden Argumente:

  • --server_port: der Port, den der Pathways-Ressourcenmanager für die Kommunikation mit anderen Pathways-Komponenten verwendet.
  • --node_type: der Knotentyp. Legen Sie dieses Argument für den Pathways-Ressourcenmanager auf resource_manager fest. Für die anderen Container ist es nicht erforderlich.
  • --instance_count: die Anzahl der TPU-Slices.
  • --instance_type: Der TPU-Typ und die Topologie des Slice im Format tpuTPU_TYPE:TOPOLOGY, z. B. tpuv5e:4x4.
  • --gcs_scratch_location: ein Cloud Storage-Bucket, der für temporäre Dateien verwendet wird.

IFRT-Proxyserver

Sie können den folgenden containerSpec verwenden, um einen IFRT-Proxyserver bereitzustellen:

  - name: pathways-proxy
    image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/proxy_server:latest
    imagePullPolicy: Always
    env:
    - name: PATHWAYS_HEAD
      valueFrom:
        fieldRef:
          fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
    args:
    - --resource_manager_address=$(PATHWAYS_HEAD):29001
    - --server_port=29000
    - --gcs_scratch_location=gs://BUCKET_NAME
    ports:
    - containerPort: 29000

Der Container verwendet die folgenden Argumente:

  • --resource_manager_address: Der Hostname und der Port, die der Proxyserver für die Kommunikation mit dem Pathways-Ressourcenmanager verwendet. Der Port muss mit dem --server_port-Wert übereinstimmen, der für den Pathways-Ressourcenmanager-Container verwendet wird.
  • --server_port: Der Port, den der IFRT-Proxyserver für die Kommunikation mit dem IFRT-Proxyclient verwendet.
  • --gcs_scratch_location: Ein Cloud Storage-Bucket, der für temporäre Dateien verwendet wird.

Pathways-Mitarbeiter

Sie können die folgenden containerSpec verwenden, um Pathways-Worker bereitzustellen:

- name: worker
  image: us-docker.pkg.dev/cloud-tpu-v2-images/cloud-pathways/server:latest
  imagePullPolicy: Always
  env:
  - name: PATHWAYS_HEAD
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/coordinator']"
  - name: MEGASCALE_NUM_SLICES
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas']"
  - name: MEGASCALE_SLICE_ID
    valueFrom:
      fieldRef:
        fieldPath: "metadata.labels['jobset.sigs.k8s.io/job-index']"
  - name: MEGASCALE_COORDINATOR_ADDRESS
    value: "$(PATHWAYS_HEAD)"
  args:
  - --server_port=29001
  - --resource_manager_address=$(PATHWAYS_HEAD):29001
  - --gcs_scratch_location=gs://BUCKET_NAME
  ports:
  - containerPort: 29001
  resources:
    limits:
      google.com/tpu: "4"

Der Container verwendet die folgenden Argumente:

  • --resource_manager_address: der Hostname und der Port, über den die TPU-Worker mit dem Pathways-Ressourcenmanager kommunizieren. Der Port muss mit dem --server_port-Wert übereinstimmen, der für den Pathways-Ressourcenmanager-Container verwendet wird.
  • --server_port: der Port, den die Worker für die Kommunikation mit dem Proxyserver und dem Pathways-Ressourcenmanager verwenden.
  • --gcs_scratch_location: Ein Cloud Storage-Bucket, der für temporäre Dateien verwendet wird.

Der Pathways-Ressourcenmanager, der IFRT-Proxyserver und die Pathways-Worker können alle unterschiedliche Ports haben. In diesem Beispiel verwenden der Pathways-Ressourcenmanager und der Pathways-Worker jedoch denselben Port.

Pathways-Komponenten in Cluster Director

In diesem Abschnitt werden Pathways-Komponenten Slurm-Jobs in einem Cluster Director-Cluster zugeordnet.

Die Pathways-Skripts führen die Pathways-Komponenten mit Enroot aus. Das Script setup_env.sh verwendet die folgenden Container-Images, die der JAX-Version (0.11.1) der von ihm erstellten Python-Umgebung entsprechen:

Containertyp Standort

IFRT-Proxyserver

us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:20261103-jax_0.11.1

Pathways-Ressourcenmanager/-Worker

us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:20261103-jax_0.11.1

Sie starten die Pathways-Komponenten, indem Sie das Skript start_pathways.sh als Pathways-Dienstjob senden. Sie schreiben die folgenden Befehle nicht selbst: Die Skripts leiten die Argumente aus den Umgebungsvariablen ab, die Sie beim Senden von start_pathways.sh festlegen, z. B. SLICE_SHAPE, MIN_SLICES, MAX_SLICES und GCS_BUCKET.

Pathways-Ressourcenmanager

Der Pathways-Dienstjob führt den Pathways-Ressourcenmanager auf einem CPU-Knoten aus. Im Pathways-Servercontainer wird der folgende Befehl ausgeführt:

/pathways/server \
    --server_port=29001 \
    --node_type=resource_manager \
    --instance_type=tpuTPU_TYPE:TOPOLOGY \
    --instance_count=MAX_SLICES \
    --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \
    --enable_gcs_olm=false

Der Befehl verwendet die folgenden Argumente:

  • --server_port: der Port, den der Pathways-Ressourcenmanager überwacht. Der IFRT-Proxyserver und die Pathways-Worker stellen eine Verbindung zu diesem Port her.
  • --node_type: der Knotentyp, resource_manager. Dieses Argument wird nur vom Pathways-Ressourcenmanager festgelegt.
  • --instance_type: Der TPU-Typ und die Topologie jedes Slice, aus SLICE_SHAPE. Beispiel: SLICE_SHAPE=2x4 legt tpuv6e:2x4 fest.
  • --instance_count: die Anzahl der TPU-Slices, von MAX_SLICES.
  • --gcs_scratch_location: ein Pfad in Ihrem Cloud Storage-Bucket aus GCS_BUCKET, der für temporäre Dateien wie den Kompilierungscache verwendet wird.
  • --enable_gcs_olm: Gibt an, ob Pathways Cloud Storage-Objektlebenszyklusregeln erstellt, um alte Kompilierungscache-Einträge zu bereinigen. Das Skript deaktiviert diese Funktion, da dafür zusätzliche Berechtigungen für den Bucket erforderlich sind.

IFRT-Proxyserver

Der Pathways-Dienstjob führt den IFRT-Proxyserver auf demselben CPU-Knoten wie den Pathways-Ressourcenmanager aus. Im IFRT-Proxyserver-Container wird der folgende Befehl ausgeführt:

/pathways/proxy_server \
    --server_port=29000 \
    --resource_manager_address=RM_IP:29001 \
    --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \
    --num_elastic_instances=MAX_SLICES

Der Befehl verwendet die folgenden Argumente:

  • --server_port: der Port, den der IFRT-Proxyserver überwacht. Ihr JAX-Programm stellt eine Verbindung zu diesem Port her.
  • --resource_manager_address: die Adresse des Pathways-Ressourcenmanagers. RM_IP ist die interne IP-Adresse des CPU-Knotens.
  • --gcs_scratch_location: derselbe Cloud Storage-Pfad wie der Pathways-Ressourcenmanager.
  • --num_elastic_instances: die Anzahl der Slices, die nicht verfügbar sein können, ohne die IFRT-Proxy-Sitzung zu beenden, von MAX_SLICES. Ihre Arbeitslast erzwingt MIN_SLICES.

Wenn mindestens MIN_SLICES Slices ausgeführt werden, schreibt der Dienstjob die IFRT-Proxyserveradresse in ~/pathways/cluster.env. Ihr JAX-Programm, der IFRT-Proxy-Client, liest diese Datei, um eine Verbindung zu Pathways herzustellen. Das JAX-Programm wird als Arbeitslast ausgeführt, z. B. run_maxtext.sh, oder in einem Jupyter-Notebook.

Pathways-Mitarbeiter

Der Pathways-Dienstjob sendet ein Slurm-Job-Array (run_worker_slice.sh) mit einer Array-Aufgabe für jeden TPU-Slice. Jede Array-Aufgabe führt den folgenden Befehl im Pathways-Servercontainer auf jeder TPU-VM in ihrem Slice aus:

/pathways/server \
    --server_port=29001 \
    --resource_manager_address=RM_IP:29001 \
    --gcs_scratch_location=gs://BUCKET_NAME/pathways-scratch/shared \
    --enable_gcs_olm=false

Der Befehl verwendet die folgenden Argumente:

  • --server_port: der Port, den der Pathways-Worker überwacht. Der Pathways-Ressourcenmanager und der IFRT-Proxyserver stellen eine Verbindung zu diesem Port her.
  • --resource_manager_address: die Adresse des Pathways-Ressourcenmanagers, die mit der des IFRT-Proxyservers übereinstimmt.
  • --gcs_scratch_location, --enable_gcs_olm: wie beim Pathways-Ressourcenmanager.

Das Skript legt auch die TPU-Topologie und die Umgebungsvariablen MEGASCALE_* für jeden Worker aus den Knoten des Slices und dem Array-Aufgabenindex fest.

Wenn ein TPU-Slice vorzeitig beendet wird oder fehlschlägt, stellt Slurm die Array-Aufgabe wieder in die Warteschlange. Wenn Sie auf den Clusterknoten passwortloses sudo haben, werden die Knoten des Slices durch den Dienstjob ebenfalls heruntergefahren, damit Slurm sie neu erstellen kann. Der Pathways-Ressourcenmanager und der IFRT-Proxyserver laufen weiter, sodass Ihr JAX-Programm keine Verbindung zu einer neuen Adresse herstellen muss.

Nächste Schritte