Einführung in Pathways in der Cloud

Pathways ist ein System, mit dem sich umfangreiche, komplexe und nur selten aktivierte Machine-Learning-Systeme erstellen lassen. Es ermöglicht die Verwendung von Tausenden oder Zehntausenden von Beschleunigern, wobei je nach Verarbeitungsanforderungen dynamisch unterschiedliche Mengen an Rechenleistung für verschiedene Aufgaben zugewiesen werden können.

Pathways vereinfacht umfangreiche Machine-Learning-Berechnungen, da ein einzelner JAX-Client Arbeitslasten auf mehreren großen TPU-Slices orchestrieren kann, die potenziell Tausende von TPU-Chips umfassen.

Pathways wird intern bei Google zum Trainieren großer Modelle wie Gemini verwendet. 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 TPUs mit einem und mehreren Slices in Google Kubernetes Engine, sowie über allgemeine Erfahrung mit TPUs mit mehreren Slicesverfügen.

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

Einzelner 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, wodurch das System als eine einzelne große Maschine mit vielen lokalen Beschleunigern erscheint.

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

Flexibilität

Unterstützt komplexere Berechnungsmuster über SPMD hinaus, einschließlich asymmetrischer Pipeline-Parallelität und Berechnungssparsamkeit.

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 der Worker und verwaltet die Jobplanung, das Anhalten und das Fortsetzen. Er 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. Er empfängt High-Level Operations (HLOs) von Ihrem Programm. Der Pathways-Client ist für die Koordinierung mit dem Pathways-Ressourcenmanager verantwortlich, um basierend auf dem Nutzercode zu bestimmen, wo kompilierte Programme zur Ausführung platziert werden sollen. Er bietet einem bestimmten JAX-Client eine einheitliche Ansicht des Systems. Für diese Komponente sind nur CPU-Ressourcen erforderlich.

Pathways-Worker

Dies sind die Prozesse, die auf den Beschleunigermaschinen (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 an Ihr Programm zurück. Für diese Komponente sind Beschleunigerressourcen erforderlich.

IFRT-Proxyclient

Dies ist eine Open-Source-Implementierung der Interim Framework Runtime (IFRT) API, die den Nutzercode von der zugrunde liegenden Laufzeit entkoppelt und die Code portabilität und ‑transparenz verbessert. JAX verwendet diese Implementierung als Alternative zur standardmäßigen Laufzeit mit mehreren Controllern. Der IFRT-Proxyclient fungiert als Kommunikationsbrücke zwischen Ihrem Programm und den Pathways-Komponenten. Er sendet Anfragen an den IFRT-Proxyserver und empfängt Ergebnisse von ihm. Es ist eine Open-Source-Implementierung der IFRT API. Für diese Komponente sind nur CPU-Ressourcen erforderlich.

IFRT-Proxyserver

Dieser gRPC-Server empfängt Anfragen vom IFRT-Proxyclient 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 Beschleuniger-VM, um nutzerspezifischen Python-Code direkt auf der Beschleuniger-VM auszuführen und so 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 auf dem gRPC-Transport.

Zeigt die Beziehung der Pathways-Komponenten.
Pathways-Komponenten

PathwaysJob API

Die PathwaysJob API ist eine Kubernetes-native Open-Source-API, mit der Sie ML-Trainings- und Batch-Inferenz-Arbeitslasten bereitstellen können. Der Controller für den PathwaysJob nutzt die JobSet API, um den Lebenszyklus und die Koordinierung aller Pathways-Komponenten zu verwalten. Diese benutzerdefinierte Ressourcendefinition (Custom Resource Definition, CRD) bietet eine Schnittstelle auf hoher Ebene zum Definieren Ihrer Pathways-Arbeitslasten, sodass Sie einzelne Pod-Spezifikationen für gängige Szenarien nicht direkt verwalten müssen. Eine umfassende Liste aller Parameter und ihrer spezifischen Bedeutungen finden Sie in der Dokumentation zur PathwaysJob API auf GitHub.

apiVersion: pathways-job.pathways.domain/v1
kind: PathwaysJob
metadata:
  name: pathways-USER
spec:
  maxRestarts: MAX_RESTARTS
  pathwaysVersion: jax-JAX_VERSION
  workers:
  - type: $(TPU_MACHINE_TYPE)
    topology: $(TOPOLOGY)
    numSlices: $(WORKLOAD_NODEPOOL_COUNT)
    maxSliceRestarts:  # Optional
  customComponents: # This section is completely optional
  - componentType: proxy_server
    image: CUSTOM_PROXY_SERVER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: pathways_server
    image: CUSTOM_PATHWAYS_SERVER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: worker
    image: CUSTOM_WORKER
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  - componentType: colocated_python_sidecar
    image: CUSTOM_SIDECAR_IMAGE
    customFlags:
    - --flag_name_1=value_1
    customEnv:
    - name: key_1
      value: value_1
  pathwaysDir: "gs://BUCKET_NAME" # Pre-create this bucket.
  controller:
    deploymentMode: default # Default mode deploys pathways cpu resources (resource
# manager and proxy server) on a dedicated CPU node, recommended for training
    elasticSlices: ELASTIC_SLICES
    template:
      spec:
        containers:
        - name: main
          image: python:3.11
        command:
          - bash
          - -c
          - |
            pip install pathwaysutils
            python3 -c 'import pathwaysutils; import jax; pathwaysutils.initialize(); print(jax.devices())'

In der folgenden Tabelle werden die Einstellungen für die PathwaysJob API beschrieben:

Attribut Beschreibung
apiVersion Gibt die API-Version für die benutzerdefinierte Ressourcendefinition (CRD) PathwaysJob an: pathways-job.pathways.domain/v1.
kind Identifiziert das Kubernetes-Objekt als PathwaysJob.
metadata.name Der Name des PathwaysJob-Objekts in Kubernetes, der in der Regel dem Muster pathways-folgt.
spec Definiert den gewünschten Status und die Konfiguration für den PathwaysJob.
spec.maxRestarts Die maximale Anzahl von Malen, die der PathwaysJob automatisch vom System neu gestartet werden kann, wenn Fehler auftreten.
spec.pathwaysVersion (Optional) Gibt die gewünschte Version des JAX-Frameworks an, die für diesen Job in der Pathways-Umgebung verwendet werden soll (z. B. jax-0.5.3).
spec.workers Ein Array, das die Konfiguration für den Worker-Pool des PathwaysJob definiert und in der Regel TPU-Ressourcen verwendet.
spec.workers[].type Der Typ der TPU-Maschine, die für die Worker-Knoten verwendet werden soll (z. B. könnte $TPU_MACHINE_TYPE ct6e-standard-4t sein).
spec.workers[].topology Die Topologie der TPU-Slices, die den Workern zugewiesen sind (z. B. könnte $TOPOLOGY 2x2, 4x4 oder 2x2x2 sein).
spec.workers[].numSlices Die Anzahl der TPU-Slices, die für den Worker-Pool bereitgestellt werden sollen (z. B. könnte $WORKLOAD_NODEPOOL_COUNT 2 sein).
spec.workers[].maxSliceRestarts (Optional) Die maximale Anzahl von Malen, die ein einzelner Worker innerhalb eines Slice neu gestartet werden kann, wenn er fehlschlägt.
spec.customComponents (Optional) Ein Array, mit dem Sie benutzerdefinierte Komponenten (z. B. Proxyserver, Pathways-Server oder zusätzliche Worker) neben dem Hauptjob definieren und bereitstellen können.
spec.customComponents[].componentType Gibt den Typ der benutzerdefinierten Komponente an, die definiert wird (z. B. proxy_server, pathways_server, worker, colocated_python_sidecar).
spec.customComponents[].image Das Docker-Image, das für den Container dieser benutzerdefinierten Komponente verwendet werden soll.
spec.customComponents[].customFlags Ein Array benutzerdefinierter Befehlszeilen-Flags, die beim Start an den Container übergeben werden.
spec.customComponents[].customEnv Ein Array benutzerdefinierter Umgebungsvariablen, die im Container festgelegt werden sollen. Jedes Element hat einen Namen und einen Wert.
spec.pathwaysDir Der Cloud Storage-Bucket, der vom PathwaysJob zum Speichern von Kompilierungsartefakten und anderen temporären Daten verwendet wird. Dieser Bucket muss erstellt werden, bevor Sie Ihre Arbeitslast ausführen.
spec.controller Konfigurationseinstellungen für den Pathways-Controller, der die gesamte Jobausführung verwaltet.
spec.controller.deploymentMode Gibt an, wie die CPU-Ressourcen des Pathways-Controllers (Pathways-Ressourcenmanager und Proxyserver) bereitgestellt werden. Im Standardmodus werden sie auf einem dedizierten CPU-Knoten bereitgestellt, während colocate_head_with_workers sie zusammen mit einem TPU Worker bereitstellt.
spec.controller.elasticSlices (Optional) Die maximale Anzahl von TPU-Slices, die während der Ausführung des Jobs nicht verfügbar sein können bevor er als fehlerhaft eingestuft wird.
spec.controller.template (Optional) Definiert die Pod-Vorlage für den Nutzerjob. Dies ist erforderlich für Batch-Arbeitslasten, aber nicht für interaktive Arbeitslasten.
spec.controller.template.spec Die Spezifikation des Pods für den Nutzerjob.
spec.controller.template.spec.containers Ein Array, das die Container definiert, die im Nutzerjob ausgeführt werden.
spec.controller.template.spec.containers[].name Der Name des Containers im Nutzerjob (in diesem Beispiel ist es „main“).
spec.controller.template.spec.containers[].image Das Docker-Image, das für den Container im Hauptcontainer verwendet werden soll (in diesem Beispiel ist es „python:3.11“).
spec.controller.template.spec.containers[].command Der Befehl, der beim Start des Hauptcontainers ausgeführt werden soll. In diesem Beispiel werden `pathwaysutils` installiert, Pathways initialisiert und JAX-Geräte ausgegeben.

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

Speicherort

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 mit der folgenden containerSpec den Pathways-Ressourcenmanager 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. Dieser sollte für den Pathways-Ressourcenmanager auf „resource_manager“ gesetzt 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

Mit der folgenden containerSpec können Sie einen IFRT-Proxyserver bereitstellen:

 - 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 Wert --server_port übereinstimmen, der für den Container des Pathways-Ressourcenmanagers 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-Worker

Mit der folgenden containerSpec können Sie Pathways-Worker bereitstellen:

- 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 Wert --server_port übereinstimmen, der für den Container des Pathways-Ressourcenmanagers verwendet wird.
  • --server_port: Die Worker verwenden diesen Port für die Kommunikation mit dem Proxy server 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