Einführung in Pathways in der Cloud

Pathways ist ein System, das die Erstellung von umfangreichen, dünnbesetzten Machine-Learning-Systemen mit mehreren Aufgaben ermöglicht. Sie ermöglicht die Verwendung von Tausenden oder Zehntausenden von Beschleunigern und die dynamische Zuweisung unterschiedlicher Rechenkapazitäten für verschiedene Aufgaben basierend auf ihren Verarbeitungsanforderungen.

Pathways vereinfacht umfangreiche Berechnungen für maschinelles Lernen, da ein einzelner JAX-Client Arbeitslasten über mehrere große TPU-Slices hinweg orchestrieren kann, die potenziell 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 von Pathways verwaltete TPUs 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 TPUs mit einem oder mehreren Slices in der Google Kubernetes Engine, und dass Sie allgemeine Erfahrung mit TPUs mit mehreren Slices haben.

Einzelner Controller und mehrere Controller

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

Feature

Einzelner Controller (Pathways)

Multi-Controller (JAX-Standard)

Kontrolle

Zentraler Kontrollpunkt: Ein einzelnes Clientprogramm fungiert als zentraler Controller.

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

Ansicht

Ganzheitliche Übersicht: 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: Nutzer interagieren mit einem einzigen Client, sodass das System wie eine einzige große Maschine mit vielen lokalen Beschleunigern erscheint.

SPMD: Verwendet hauptsächlich das SPMD-Paradigma, bei dem alle Geräte dasselbe Programm ausführen müssen.

Flexibilität

Unterstützt komplexere 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

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

Pathways-Ressourcenmanager

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

Pathways-Client

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

Pathways-Mitarbeiter

Dies sind die Prozesse, die auf den Beschleuniger-Maschinen (TPU-VMs) ausgeführt werden. 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

Dies ist eine OSS-Implementierung der Interim Framework Runtime (IFRT) API, die den Nutzercode von der zugrunde liegenden Laufzeit entkoppelt und die Portabilität und Transparenz des Codes 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. Es ist eine OSS-Implementierung der IFRT API. Diese Komponente benötigt nur CPU-Ressourcen.

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. Diese Komponente benötigt nur CPU-Ressourcen.

Sidecar-Server

Dieser gRPC-Server befindet sich zusammen mit dem Pathways-Worker auf der Beschleuniger-VM, um den vom Nutzer angegebenen Python-Code direkt auf der Beschleuniger-VM auszuführen und die Latenz bei der Datenübertragung vom Controller zu den Beschleunigern zu reduzieren. Der Sidecar-Server interagiert mit dem Pathways-Worker über ein benutzerdefiniertes versioniertes Protokoll im gRPC-Transport.

Zeigt die Beziehung der Pathways-Komponenten.
Pathways-Komponenten

Pathways-Komponenten in GKE

In diesem Abschnitt werden Pathways-Komponenten Google Kubernetes Engine-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/pathways/proxy_server:jax-<jax-version>

Pathways-Ressourcenmanager/-Worker

us-docker.pkg.dev/cloud-tpu-v2-images/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/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, um mit anderen Pathways-Komponenten zu kommunizieren.
  • --node_type: der Knotentyp. Dieser sollte für den Pathways-Ressourcenmanager auf „resource_manager“ festgelegt werden und ist für die anderen Container 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, um mit dem IFRT-Proxyclient zu kommunizieren.
  • --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, um mit dem Proxyserver und dem Pathways-Ressourcenmanager zu kommunizieren.
  • --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