Einführung in Pathways in der Cloud

Pathways ist ein System, das die Entwicklung von umfangreichen, auf mehrere Aufgaben ausgerichteten und spärlich aktivierten Machine-Learning-Systemen ermöglicht. Es ermöglicht die Verwendung von Tausenden oder Zehntausenden von Beschleunigern und die dynamische Zuweisung unterschiedlicher Mengen an Rechenleistung für verschiedene Aufgaben basierend auf ihren Verarbeitungsanforderungen.

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

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

Hinweis

Sie benötigen Folgendes:

In diesem Dokument wird beschrieben, wie Sie verwaltete TPUs von Pathways in Google Kubernetes Engine (GKE) für Batch-, Echtzeit- und interaktive Arbeitslasten verwenden. Es wird davon ausgegangen, dass Sie bereits mit der Verwendung von TPUs mit GKE vertraut sind, einschließlich Einzel- und Multislice-TPUs in Google Kubernetes Engine, sowie allgemeine Erfahrung mit Multislice-TPUs haben.

Einzelner Controller und mehrere Controller

Es gibt hauptsächlich zwei verschiedene Möglichkeiten, Berechnungen auf mehreren Geräten zu verwalten und zu orchestrieren:

Feature

Einzelner Controller (Pathways)

Mehrere Controller (JAX-Standard)

Kontrolle

Zentraler Kontrollpunkt: Ein einzelnes Clientprogramm fungiert als zentraler Controller.

Verteilte Steuerung: Mehrere Prozesse sind beteiligt, jeder mit einer eigenen Python-Interpreterinstanz.

Ansicht

Ganzheitliche Übersicht: Der Client sieht alle Geräte als ein einziges, einheitliches System.

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

Programmierung

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

SPMD: Hier wird hauptsächlich das SPMD-Paradigma verwendet, bei dem auf allen Geräten dasselbe Programm ausgeführt werden muss.

Flexibilität

Unterstützt komplexere Rechenmuster als SPMD, einschließlich asymmetrischer Pipeline-Parallelität und Rechensparsamkeit.

Die Ressourcenverwaltung kann weniger flexibel sein, insbesondere bei verschiedenen TPU-Slices.

Pathways-Komponenten

Im folgenden Abschnitt werden die Hauptkomponenten der Pathways-Architektur beschrieben.

Pathways Resource Manager

Dies ist die zentrale Steuerungsebene des Pathways-Systems. Sie verwaltet alle Beschleunigerressourcen und ist für die Koordinierung der Zuweisung von Beschleunigern für Nutzerjobs verantwortlich. Sie überwacht den Zustand der Worker und kümmert sich um die Planung, das Pausieren und das Fortsetzen von Jobs. Es dient als zentrale Anlaufstelle für Fehler und Systemstatus. Für diese Komponente sind nur CPU-Ressourcen erforderlich.

Pathways-Client

Dies ist eine Implementierung der Interim Framework Runtime (IFRT), die als Einstiegspunkt in das Pathways-System dient. Sie empfängt High-Level Operations (HLOs) von Ihrem Programm. Der Pathways-Client ist dafür verantwortlich, mit dem Pathways-Ressourcenmanager abzustimmen, wo kompilierte Programme basierend auf dem Nutzercode zur Ausführung platziert werden sollen. Es bietet einem bestimmten JAX-Client eine einheitliche Ansicht des Systems. Für diese Komponente sind nur CPU-Ressourcen erforderlich.

Pathways-Mitarbeiter

Dies sind die Prozesse, die auf den Beschleuniger-Maschinen (TPU-VMs) ausgeführt werden. Sie empfangen 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 zurück an Ihr Programm. Für diese Komponente sind Accelerator-Ressourcen erforderlich.

IFRT-Proxyclient

Dies ist eine OSS-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 Laufzeit mit mehreren Controllern. 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. Es handelt sich um eine OSS-Implementierung der IFRT API. Für diese Komponente sind nur CPU-Ressourcen erforderlich.

IFRT-Proxyserver

Dieser gRPC-Server empfängt Anfragen vom IFRT-Proxy-Client und leitet sie an den Pathways-Client weiter, der die eigentliche Verteilung der Arbeit übernimmt. Für diese Komponente sind nur CPU-Ressourcen erforderlich.

Sidecar-Server

Dieser gRPC-Server befindet sich zusammen mit dem Pathways-Worker auf der Accelerator-VM, um nutzerspezifischen Python-Code direkt auf der Accelerator-VM auszuführen. So wird die Latenz für die Datenübertragung vom Controller zu den Accelerators reduziert. Der Sidecar-Server interagiert mit dem Pathways-Worker über ein benutzerdefiniertes versioniertes Protokoll auf dem gRPC-Transport.

Zeigt die Beziehung der Pathways-Komponenten.
Pathways-Komponenten

Pathways-Komponenten in GKE

In diesem Abschnitt werden die Komponenten von Pathways den Komponenten von Google Kubernetes Engine wie Containern und Pods zugeordnet.

Sie finden die Container-Images für Pathways an den folgenden Speicherorten.

Containertyp

Standort

IFRT-Proxyserver

us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:jax-<jax-version>

Pathways-Ressourcenmanager/-Worker

us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:jax-<jax-version>

Pathways Resource Manager

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

  - name: pathways-rm
    image: us-docker.pkg.dev/cloud-tpu-v2-images/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

Argumentbeschreibungen:

  • --server_port: Der Pathways-Ressourcenmanager verwendet diesen Port für die Kommunikation mit anderen Pathways-Komponenten.
  • --node_type: der Knotentyp. Für den Pathways-Ressourcenmanager sollte dieser Wert auf „resource_manager“ festgelegt werden. Für die anderen Container ist er nicht erforderlich.
  • --instance_count: die Anzahl der TPU-Slices.
  • --instance_type: Der TPU-Typ und die Topologie des Slice. Im Format tpu{TPU type}:{TPU 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/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

Argumentbeschreibungen:

  • --resource_manager_address: Der Hostname und der Port, die der Proxyserver für die Kommunikation mit dem Pathways-Ressourcenmanager verwendet. Der Port sollte mit dem --server_port-Wert übereinstimmen, der für den Pathways-Ressourcenmanager-Container verwendet wird.
  • --server_port: Der IFRT-Proxyserver verwendet diesen Port für die Kommunikation mit dem IFRT-Proxyclient.
  • --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/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"

Argumentbeschreibungen:

  • --resource_manager_address: Der Hostname und der Port, die die TPU-Worker für die Kommunikation mit dem Pathways-Ressourcenmanager verwenden. Der Port sollte mit dem --server_port-Wert übereinstimmen, der für den Pathways-Ressourcenmanager-Container verwendet wird.
  • --server_port: Die Worker verwenden diesen Port für die Kommunikation mit dem Proxyserver und dem Pathways-Ressourcenmanager.
  • --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.

Nächste Schritte