TPU-Arbeitslasten in GKE Autopilot bereitstellen

Auf dieser Seite wird beschrieben, wie Sie Arbeitslasten für maschinelles Lernen (ML) beschleunigen, indem Sie Cloud TPU-Beschleuniger (TPUs) in Google Kubernetes Engine (GKE) Autopilot-Clustern bereitstellen. Machen Sie sich vor dem Lesen dieser Seite mit den folgenden Konzepten vertraut:

  1. Einführung in Cloud TPU
  2. Cloud TPU-Systemarchitektur
  3. TPUs in GKE

Funktionsweise von TPUs in Autopilot

Wenn Sie TPUs in Autopilot-Arbeitslasten verwenden möchten, fordern Sie in Ihrem Arbeitslastmanifest eine TPU-Version und eine unterstützte Topologie für diese TPU-Version an. Anschließend verwenden Sie die Kubernetes-Felder resources.requests und resources.limits, um die Anzahl der TPU-Chips anzugeben, die für die Arbeitslast verwendet werden sollen. Wenn Sie die Arbeitslast bereitstellen, stellt GKE Knoten mit der angeforderten TPU-Konfiguration bereit und plant Ihre Pods auf den Knoten. GKE platziert die einzelnen Arbeitslasten auf je eigenen Knoten, damit jeder Pod bei minimalem Störungsrisiko auf die vollständigen Ressourcen des Knotens zugreifen kann.

TPUs in Autopilot sind mit den folgenden Funktionen kompatibel:

  1. Spot-Pods
  2. Spezifische Kapazitätsreservierungen
  3. Pods mit verlängerter Laufzeit

TPU-Konfiguration planen

Bevor Sie TPUs anfordern, sollten Sie sich anhand Ihrer CPU- und Arbeitsspeicheranforderungen für eine Konfiguration entscheiden. Folgende Entscheidungen sind erforderlich:

  • TPU-Version: Die spezifische Cloud TPU-Version, z. B. v5e.
  • Topologie für die ausgewählte TPU-Version: die Anordnung und die Anzahl der TPUs.

Die von Ihnen ausgewählte TPU-Version und Topologie bestimmen, ob GKE Knoten als Slices mit einem einzelnen Host oder Slices mit mehreren Hosts bereitstellt. In Slices mit einem einzelnen Host ist jeder Knoten unabhängig von anderen TPU-Slice-Knoten. In Slices mit mehreren Hosts erstellt GKE eine Gruppe von Knoten mit verbundenen VMs in einem TPU-Slice. Slices mit mehreren Hosts sind atomar, d. h., GKE skaliert die gesamte verbundene Gruppe von Knoten als eine Einheit hoch oder herunter.

Informationen zu den verfügbaren TPU-Versionen, den entsprechenden Topologien, der CPU- und Speicherkapazität und dem resultierenden Slice-Typ finden Sie unter Autopilot-TPU-Konfiguration auswählen.

Preise

Preisinformationen finden Sie unter Autopilot-Preise.

Hinweise

Führen Sie die folgenden Schritte durch, bevor Sie beginnen:

  • Aktivieren Sie die Google Kubernetes Engine API.
    • Google Kubernetes Engine API aktivieren
  • Wenn Sie die Google Cloud CLI für diese Aufgabe verwenden möchten, müssen Sie die gcloud CLI installieren und dann initialisieren. Wenn Sie die gcloud CLI bereits installiert haben, rufen Sie die neueste Version mit gcloud components update ab.
  • Sie benötigen einen Autopilot-Cluster, auf dem GKE-Version 1.29.2-gke.1521000 oder höher ausgeführt wird.
  • Wenn Sie reservierte TPUs verwenden möchten, benötigen Sie eine spezifische Kapazitätsreservierung. Eine Anleitung finden Sie unter Reservierte zonale Ressourcen nutzen.

TPU-Kontingent prüfen

Zum Erstellen von TPU-Slice-Knoten benötigen Sie ein TPU-Kontingent, sofern Sie keine vorhandene Kapazitätsreservierung verwenden. Wenn Sie reservierte TPUs verwenden, überspringen Sie diesen Abschnitt.

Zum Erstellen von TPU-Scheibenknoten in GKE ist ein Kontingent für die Compute Engine API (compute.googleapis.com) erforderlich, kein Cloud TPU API-Kontingent (tpu.googleapis.com). Der Name des Kontingents unterscheidet sich in regulären Autopilot-Pods und in Spot-Pods.

Führen Sie die folgenden Schritte aus, um das Limit und die aktuelle Nutzung Ihres Compute Engine API-Kontingents für TPUs zu prüfen:

  1. Rufen Sie in der Google Cloud Console die Seite Kontingente auf.

    Kontingente aufrufen

  2. Gehen Sie im Feld Filter folgendermaßen vor:

    1. Wählen Sie das Attribut Dienst aus, geben Sie Compute Engine API ein und drücken Sie die Eingabetaste.

    2. Wählen Sie das Attribut Typ und dann Kontingent aus.

    3. Wählen Sie das Attribut Name aus und geben Sie einen Kontingentnamen basierend auf dem gewünschten TPU-Typ ein:

      • TPU v5p (tpu-v5p-slice): TPU v5p-Chips
      • TPU v5e (tpu-v5-lite-podslice): TPU v5 Lite PodSlice-Chips
      • TPU v5e (tpu-v5-lite-device): TPU v5 Lite-Gerätechips
      • TPU v4 (tpu-v4-podslice): TPU v4 PodSlice-Chips

      Wählen Sie für Spot-Pods das entsprechende Kontingent "Auf Abruf" aus.

    4. Wählen Sie das Attribut Dimensionen (z.B. Standorte) aus und geben Sie region: gefolgt vom Namen der Region ein, in der Sie TPUs in GKE erstellen möchten. Geben Sie beispielsweise region:us-west4 ein, wenn Sie TPU-Slice-Knoten in der Zone us-west4-a erstellen möchten. Das TPU-Kontingent ist regional, d. h. alle Zonen innerhalb derselben Region nutzen dasselbe TPU-Kontingent.

Wenn keine Kontingente vorhanden sind, die Sie eingegeben haben, wurde dem Projekt keines der angegebenen Kontingente für die gewünschte Region zugewiesen und Sie müssen eine TPU-Kontingenterhöhung anfordern.

TPU-Anwendung vorbereiten

Für TPU-Arbeitslasten gelten die folgenden Vorbereitungsanforderungen.

  1. Frameworks wie JAX, PyTorch und TensorFlow greifen auf TPU-VMs mit der freigegebenen Bibliothek libtpu zu. libtpu enthält den XLA-Compiler, die TPU-Laufzeitsoftware und den TPU-Treiber. Jede Version von PyTorch und JAX erfordert eine bestimmte Version von libtpu.so. Für die Verwendung von TPUs in GKE benötigen Sie die folgenden Versionen:
    TPU-Typ libtpu.so-Version
    TPU v5e
    tpu-v5-lite-podslice
    tpu-v5-lite-device
    TPU v5p
    tpu-v5p-slice
    • Empfohlene jax[tpu]-Version: 0.4.19 oder höher.
    • Empfohlene torchxla[tpuvm]-Version: nächtliche Version verwenden, die am 23. Oktober 2023 erstellt wurde.
    TPU v4
    tpu-v4-podslice
  2. Legen Sie die folgenden Umgebungsvariablen für den Container fest, der die TPU-Ressourcen anfordert:
    • TPU_WORKER_ID: Eine eindeutige Ganzzahl für jeden Pod. Diese ID bezeichnet eine eindeutige Worker-ID im TPU-Slice. Die unterstützten Werte für dieses Feld reichen von null bis zur Anzahl der Pods minus eins.
    • TPU_WORKER_HOSTNAMES: Eine durch Kommas getrennte Liste von TPU-VM-Hostnamen oder IP-Adressen, die innerhalb des Slices miteinander kommunizieren müssen. Für jede TPU-VM im Slice sollte ein Hostname oder eine IP-Adresse vorhanden sein. Die Liste der IP-Adressen oder Hostnamen ist geordnet und ab null nach der TPU_WORKER_ID indexiert.

      • GKE fügt diese Umgebungsvariablen automatisch ein. Dazu wird ein mutierender Webhook verwendet, wenn ein Job mit completionMode: Indexed, subdomain, parallelism > 1 erstellt wird, und es werden google.com/tpu-Attribute angefordert. GKE fügt einen monitorlosen Service hinzu, sodass die DNS-Einträge für die Pods, die den Service unterstützen, hinzugefügt werden.

Nachdem Sie die Arbeitslastvorbereitung abgeschlossen haben, können Sie einen Job ausführen, der TPUs verwendet.

TPUs in einer Arbeitslast anfordern

In diesem Abschnitt erfahren Sie, wie Sie einen Job erstellen, der TPUs in Autopilot anfordert. In jeder Arbeitslast, die TPUs benötigt, müssen Sie Folgendes angeben:

  • Knotenselektoren für die TPU-Version und -Topologie
  • Die Anzahl der TPU-Chips für einen Container in Ihrer Arbeitslast

Eine Liste der unterstützten TPU-Versionen, -Topologien und der entsprechenden Anzahl von TPU-Chips und -Knoten in einem Slice finden Sie unter Autopilot-TPU-Konfiguration auswählen.

Hinweise zu TPU-Anfragen in Arbeitslasten

Nur ein Container in einem Pod kann TPUs verwenden. Die Anzahl der TPU-Chips, die ein Container anfordern muss, muss der Anzahl der TPU-Chips entsprechen, die an einen Knoten im Slice angehängt sind. Wenn Sie beispielsweise TPU v5e (tpu-v5-lite-podslice) mit einer 2x4-Topologie anfragen, können Sie eine der folgenden Optionen anfordern:

  • 4-Chips, die zwei Knoten mit mehreren Hosts und jeweils 4 Chips erstellen
  • 8-Chips, die einen Knoten mit einem einzelnen Host und 8 TPU-Chips erstellen

Als Best Practice zur Maximierung Ihrer Kosteneffizienz sollten Sie immer alle TPUs in dem von Ihnen angeforderten Slice nutzen. Wenn Sie ein Slice mit mehreren Hosts und zwei Knoten mit jeweils 4 TPU-Chips anfordern, sollten Sie eine Arbeitslast bereitstellen, die auf beiden Knoten ausgeführt wird und alle 8 TPU-Chips im Slice verbraucht.

Arbeitslast erstellen, die TPUs anfordert

Mit den folgenden Schritten wird ein Job erstellt, der TPUs anfordert. Wenn Sie Arbeitslasten haben, die in TPU-Slices mit mehreren Hosts ausgeführt werden, müssen Sie auch einen monitorlosen Dienst erstellen, der Ihre Arbeitslast anhand des Namens auswählt. Mit diesem monitorlosen Service können Pods auf verschiedenen Knoten im Slice mit mehreren Hosts miteinander kommunizieren. Dazu wird die Kubernetes-DNS-Konfiguration so aktualisiert, dass sie auf die Pods in der Arbeitslast verweist.

  1. Speichern Sie das folgende Manifest als tpu-autopilot.yaml:

    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: TPU_TYPE
        cloud.google.com/gke-tpu-topology: TOPOLOGY
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS

    Ersetzen Sie Folgendes:

    • TPU_TYPE: Der zu verwendende TPU-Typ, z. B. tpu-v4-podslice. Muss ein von GKE unterstützter Wert sein.
    • TOPOLOGY: die Anordnung der TPU-Chips im Slice, z. B. 2x2x4. Muss eine unterstützte Topologie für den ausgewählten TPU-Typ sein.
    • NUMBER_OF_CHIPS: die Anzahl der TPU-Chips, die der Container verwenden soll. Muss für limits und requests derselbe Wert sein.

  2. Stellen Sie den Job bereit:

    kubectl create -f tpu-autopilot.yaml

Wenn Sie diesen Job erstellen, führt GKE automatisch Folgendes aus:

  1. Stellt Knoten zum Ausführen der Pods bereit. Abhängig vom angegebenen TPU-Typ, der Topologie und den Ressourcenanfragen sind diese Knoten entweder Slices mit einem einzelnen Host oder mit mehreren Hosts.
  2. Fügt den Pods Markierungen und den Knoten Toleranzen hinzu, um zu verhindern, dass andere Arbeitslasten auf denselben Knoten wie TPU-Arbeitslasten ausgeführt werden.

Beispiel: Gesamtzahl der TPU-Chips in einem Slice mit mehreren Hosts anzeigen

Die folgende Arbeitslast gibt die Anzahl der TPU-Chips für alle Knoten in einem TPU-Slice mit mehreren Hosts zurück. Zum Erstellen eines Slices mit mehreren Hosts hat die Arbeitslast die folgenden Parameter:

  • TPU-Version: TPU v4
  • Topologie: 2x2x4

Diese Versions- und Topologieauswahl führt zu einem Slice mit mehreren Hosts.

  1. Speichern Sie das folgende Manifest als available-chips-multihost.yaml: 
    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
        cloud.google.com/gke-tpu-topology: 2x2x4
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
  2. Stellen Sie das Manifest bereit: 
    kubectl create -f available-chips-multihost.yaml

    GKE führt ein TPU v4-Slice mit vier VMs aus (TPU-Slice mit mehreren Hosts). Das Slice hat 16 miteinander verbundene TPU-Chips.

  3. Prüfen Sie, ob der Job vier Pods erstellt hat: 
    kubectl get pods

    Die Ausgabe sieht in etwa so aus:

    NAME                       READY   STATUS      RESTARTS   AGE
tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s
  4. Rufen Sie die Logs eines der Pods ab: 
    kubectl logs POD_NAME

    Ersetzen Sie POD_NAME durch den Namen eines der erstellten Pods. Beispiel: tpu-job-podslice-0-5cd8r.

    Die Ausgabe sieht in etwa so aus: 

    TPU cores: 16

Beispiel: TPU-Chips in einem einzelnen Knoten anzeigen

Die folgende Arbeitslast ist ein statischer Pod, der die Anzahl der TPU-Chips anzeigt, die mit einem bestimmten Knoten verbunden sind. Zum Erstellen eines Knotens mit einem einzelnen Host hat die Arbeitslast folgende Parameter:

  • TPU-Version: TPU v5e
  • Topologie: 2x4

Diese Versions- und Topologieauswahl führt zu einem Slice mit einem einzelnen Host.

  1. Speichern Sie das folgende Manifest als available-chips-singlehost.yaml: 
    apiVersion: v1
kind: Pod
metadata:
  name: tpu-job-jax-v5
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
    cloud.google.com/gke-tpu-topology: 2x4
  containers:
  - name: tpu-job
    image: python:3.10
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    command:
    - bash
    - -c
    - |
      pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8
      limits:
        google.com/tpu: 8
  2. Stellen Sie das Manifest bereit: 
    kubectl create -f available-chips-singlehost.yaml

    GKE stellt Knoten mit acht TPU-Slices mit einem einzelnen Host bereit, die TPU v5e verwenden. Jeder TPU-Knoten hat acht TPU-Chips (TPU-Slice mit einzelnem Host).

  3. Rufen Sie die Logs des Pods ab: 
    kubectl logs tpu-job-jax-v5

    Die Ausgabe sieht in etwa so aus: 

    Total TPU chips: 8

Beobachtbarkeit und Messwerte

Dashboard

Auf der Seite Kubernetes-Cluster in der Google Cloud Console werden auf dem Tab Beobachtbarkeit die TPU-Beobachtbarkeitsmesswerte angezeigt. Weitere Informationen finden Sie unter GKE-Beobachtbarkeitsmesswerte.

Das TPU-Dashboard wird nur ausgefüllt, wenn in Ihrem GKE-Cluster Systemmesswerte aktiviert sind.

Laufzeitmesswerte

In GKE Version 1.27.4-gke.900 oder höher exportieren TPU-Arbeitslasten, die JAX Version 0.4.14 oder höher verwenden und containerPort: 8431 angeben, TPU-Auslastungsmesswerte als GKE-Systemmesswerte. Die folgenden Messwerte sind in Cloud Monitoring verfügbar, um die Laufzeitleistung Ihrer TPU-Arbeitslast zu überwachen:

  • Arbeitszyklus: Prozentsatz der Zeit im vergangenen Stichprobenzeitraum (60 Sekunden), während der TensorCores auf einem TPU-Chip aktiv mit Verarbeitung beschäftigt waren. Ein höherer Prozentsatz bedeutet eine bessere TPU-Auslastung.
  • Arbeitsspeichernutzung: Menge des dem Beschleuniger zugeteilten Arbeitsspeichers in Byte. Alle 60 Sekunden wird eine Stichprobe erstellt.
  • Speicherkapazität: Gesamter Arbeitsspeicher des Beschleunigers in Byte. Alle 60 Sekunden wird eine Stichprobe erstellt.

Diese Messwerte befinden sich im Kubernetes-Knotenschema (k8s_node) und im Kubernetes-Containerschema (k8s_container).

Kubernetes-Container:

  • kubernetes.io/container/accelerator/duty_cycle
  • kubernetes.io/container/accelerator/memory_used
  • kubernetes.io/container/accelerator/memory_total

Kubernetes-Knoten:

  • kubernetes.io/node/accelerator/duty_cycle
  • kubernetes.io/node/accelerator/memory_used
  • kubernetes.io/node/accelerator/memory_total

Hostmesswerte

In GKE-Version 1.28.1-gke.1066000 oder höher exportieren VMs in einem TPU-Slice TPU-Auslastungsmesswerte als GKE-Systemmesswerte. Die folgenden Messwerte sind in Cloud Monitoring verfügbar, um die Leistung Ihres TPU-Hosts zu überwachen:

  • TensorCore-Auslastung: Aktueller Prozentsatz des verwendeten TensorCore. Der TensorCore-Wert entspricht der Summe der Matrixmultiplikationseinheiten (MXUs) plus der Vektoreinheit. Der TensorCore-Auslastungswert ist der Quotient aus den TensorCore-Vorgängen, die im letzten Stichprobenzeitraum (60 Sekunden) ausgeführt wurden, und der unterstützten Anzahl von TensorCore-Vorgängen im selben Zeitraum. Ein größerer Wert bedeutet eine bessere Auslastung.
  • Auslastung der Arbeitsspeicherbandbreite: Aktueller Prozentsatz der verwendeten Arbeitsspeicherbandbreite des Beschleunigers. Wird berechnet, indem die in einem Stichprobenzeitraum (60 s) verwendete Arbeitsspeicherbandbreite durch die maximal unterstützte Bandbreite im selben Stichprobenzeitraum geteilt wird.

Diese Messwerte befinden sich im Kubernetes-Knotenschema (k8s_node) und im Kubernetes-Containerschema (k8s_container).

Kubernetes-Container:

  • kubernetes.io/container/accelerator/tensorcore_utilization
  • kubernetes.io/container/accelerator/memory_bandwidth_utilization

Kubernetes-Knoten:

  • kubernetes.io/container/node/tensorcore_utilization
  • kubernetes.io/container/node/memory_bandwidth_utilization

Weitere Informationen finden Sie unter Kubernetes-Messwerte und GKE-Systemmesswerte.

Logging

Logs, die von Containern ausgegeben werden, die auf GKE-Knoten ausgeführt werden, einschließlich TPU-VMs, werden vom GKE-Logging-Agent erfasst und an Logging gesendet und sind in Logging sichtbar.

Empfehlungen für TPU-Arbeitslasten in Autopilot

Die folgenden Empfehlungen können die Effizienz Ihrer TPU-Arbeitslasten verbessern:

  • Verwenden Sie Pods mit verlängerter Laufzeit für einen Kulanzzeitraum von bis zu sieben Tagen, bevor GKE Ihre Pods wegen Herunterskalierungen oder Knotenupgrades beendet. Sie können Wartungsfenster und -ausschlüsse mit Pods mit verlängerter Laufzeit verwenden, um automatische Knotenupgrades weiter zu verzögern.
  • Verwenden Sie Kapazitätsreservierungen, damit Ihre Arbeitslasten angeforderte TPUs erhalten, ohne in Verfügbarkeits-Warteschlangen gestellt zu werden.