TPU-Multiplikationen in GKE bereitstellen

Auf dieser Seite erfahren Sie, wie Sie Arbeitslasten in Google Kubernetes Engine (GKE) mithilfe der Cloud TPU-Multislice-Konfiguration für ein kostengünstiges, umfangreiches Training bereitstellen.

Diese Anleitung richtet sich an Entwickler von maschinellem Lernen (ML) sowie an Plattformadministratoren und ‑operatoren, die die Kubernetes-Containerorchestrierung verwenden möchten, um umfangreiche Arbeitslasten für das Training, die Optimierung und die Inferenz von Modellen mit TPUs zu verwalten. Weitere Informationen zu gängigen Rollen und Beispielaufgaben, auf die in Google Cloud -Inhalten verwiesen wird, finden Sie unter Häufig verwendete GKE-Nutzerrollen und ‑Aufgaben.

Bevor Sie Multislice in GKE konfigurieren, sollten Sie mit den folgenden Konzepten vertraut sein:

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

Was ist ein TPU-Multisegment?

TPU Multislice ist die Architekturorganisation von VMs in einem TPU-Slice, bei der zwei oder mehr Cloud TPU-Slices über das Data Center Network (DCN) kommunizieren. Multislices ermöglicht ein kostengünstiges Full-Stack-Training mit nahezu linearer Skalierung bis zu Zehntausenden von TPU-Chips. In einer Multislice-Konfiguration stellt GKE eine Multislice-Arbeitslast auf mehreren TPU-Slices bereit. Die Kommunikation zwischen den TPU-Chips innerhalb eines Segments erfolgt über Inter-Chip-Interconnect-Verbindungen (ICI). Die Kommunikation zwischen den Segmenten erfolgt über das DCN.

Wir empfehlen die Verwendung des Multislices, wenn der Job zu groß für ein einzelnes TPU-Slice ist.

Verfügbarkeit mehrerer Segmente in GKE

  • GKE unterstützt Multislice in Version 1.27.4-gke.900 und höher.
  • Autopilot unterstützt Multislice in Version 1.29.2-gke.1521000 und höher.
  • Multislice unterstützt JAX- und PyTorch-Frameworks. Die unterstützte Mindestversion der JAX ist 2.1.
  • Multi-Slice unterstützt nur TPU-Slice-Knotenpools mit mehreren Hosts. Sie können Multislice beispielsweise nicht mit einem ct4p-hightpu-4t mit einer 2x2x1-Topologie oder einem ct5lp-hightpu-4t mit einer 2x2-Topologie verwenden, da es sich um TPU-Slice-Knotenpools mit einem Host handelt.
  • Multislices unterstützt nur das synchrone Multicontroller-Training.
  • Mehrteilige Arbeitslasten können nur über TPU-Slices ausgeführt werden, die denselben TPU-Typ, dieselbe Größe und dieselbe Topologie haben.
  • Multislice unterstützt keine TPU v3.

Hinweise

Führen Sie die folgenden Aufgaben aus, bevor Sie beginnen:

  • Aktivieren Sie die Google Kubernetes Engine API.
  • Google Kubernetes Engine API aktivieren
  • Wenn Sie die Google Cloud CLI für diesen Task 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 dem Befehl gcloud components update ab. In früheren gcloud CLI-Versionen werden die Befehle in diesem Dokument möglicherweise nicht unterstützt.

Arbeitslast auf einem Multislice ausführen

In diesem Abschnitt erfahren Sie, wie Sie eine Arbeitslast auf einem Multislice ausführen. Wenn Sie den GKE-Autopilot-Modus verwenden, fahren Sie mit dem Abschnitt Multislice-Arbeitslast ausführen fort. Autopilot-Cluster, die Version 1.29.2-gke.1521000 oder höher ausführen, aktivieren standardmäßig TPUs.

Knotenpool im Standardmodus vorbereiten

Dieser Abschnitt behandelt die folgenden Schritte:

  1. Erstellen Sie drei TPU-Slice-Knotenpools mit mehreren Hosts.
  2. Knotenpoolstatus prüfen

TPU-Slice-Knotenpool erstellen

Sie können mehrere Slice-TPU-Knotenpools mit mehreren Hosts erstellen. Erstellen Sie für diese Anleitung drei TPU-Slice-Knotenpools mit mehreren Hosts, um eine Arbeitslast mit mehreren Segmenten auszuführen. Sie können ein TPU-Slice mit mehreren Hosts über die Google Cloud CLI, Terraform oder die Google Cloud Console erstellen.

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONES \
    --machine-type=MACHINE_TYPE \
    --tpu-topology=TPU_TOPOLOGY \
    [--num-nodes=NUM_NODES] \
    [--spot \]
    [--flex-start \]
    [--enable-autoscaling \
      --max-nodes MAX_NODES]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME] \
    [--node-labels cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME,cloud.google.com/gke-workload-type=HIGH_AVAILABILITY]
    [--placement-type=COMPACT]

Ersetzen Sie Folgendes:

  • POOL_NAME: Der Name des neuen Knotenpools.
  • LOCATION: Der Name der Zone basierend auf der zu verwendenden TPU-Version. Informationen dazu, wie Sie einen verfügbaren Standort ermitteln, finden Sie unter TPU-Verfügbarkeit in GKE.
  • CLUSTER_NAME ist der Name des Clusters.
  • NODE_ZONES: Die kommagetrennte Liste mit einer oder mehreren Zonen, in denen GKE den Knotenpool erstellt.
  • MACHINE_TYPE: Der Maschinentyp für Knoten. Weitere Informationen zu den verfügbaren Maschinentypen finden Sie unter TPU-Version auswählen.
  • TPU_TOPOLOGY: Die physische Topologie für das TPU-Slice. Das Format der Topologie hängt von der TPU-Version ab. Weitere Informationen zu TPU-Topologien finden Sie in der Tabelle unter Topologie auswählen.

    Weitere Informationen finden Sie unter Topologie.

Optional können Sie auch die folgenden Flags verwenden:

  • NUM_NODES: Die maximale Anzahl von Knoten im Knotenpool. Sie muss null oder das Produkt der in TPU_TOPOLOGY definierten Werte ({A}x{B}x{C}) geteilt durch die Anzahl der Chips in jeder VM sein. Für die Multi-Host-Version von TPU v4 und TPU v5e beträgt die Anzahl der Chips in jeder VM vier. Wenn Ihre TPU_TOPOLOGY also 2x4x4 ist (TPU v4 mit vier Chips in jeder VM), beträgt NUM_NODES 32/4, was 8 entspricht. Wenn Sie dieses Flag weglassen, wird die Anzahl der Knoten anhand der Topologie und des Maschinentyps berechnet und standardmäßig festgelegt.
  • RESERVATION_NAME: Der Name der Reservierung, die GKE beim Erstellen des Knotenpools verwendet. Wenn Sie dieses Flag weglassen, verwendet GKE die verfügbaren TPU-Slice-Knotenpools. Weitere Informationen zu TPU-Reservierungen finden Sie unter TPU-Reservierung.
  • --spot: Legt den Knotenpool so fest, dass Spot-VMs für die TPU-Slice-Knoten verwendet werden. Dies kann nach dem Erstellen des Knotenpools nicht mehr geändert werden. Weitere Informationen finden Sie unter Spot-VMs.
  • --flex-start: Legt den Knotenpool so fest, dass Flex-Start-VMs verwendet werden. Flex-Start-VMs werden mit der Verbrauchsoption flex-start erstellt, die in GKE-Version 1.33.0-gke.1712000 oder höher unterstützt wird.
  • --enable-autoscaling: Knotenpool mit aktiviertem Autoscaling erstellen. Wenn GKE einen TPU-Slice-Knotenpool mit mehreren Hosts skaliert, wird der Knotenpool atomar von null auf die maximale Größe hochskaliert.

    • MAX_NODES: Maximale Größe des Knotenpools. Das Flag --max-nodes ist erforderlich, wenn --enable-autoscaling angegeben ist und dem Produkt aus den in TPU_TOPOLOGY definierten Werten ({A}x{B}x{C}) geteilt durch die Anzahl der Chips in jeder VM entsprechen muss.
  • --node-label=cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME, cloud.google.com/gke-workload-type=HIGH_AVAILABILITY: Gibt GKE, dass der Knotenpool für TPU-Slices mit mehreren Hosts eine Sammlung ist. Verwenden Sie dieses Flag, wenn die folgenden Bedingungen erfüllt sind:

    • Der Knotenpool führt Inferenzarbeitslasten im neuen Knotenpool aus.
    • Der Knotenpool verwendet TPU Trillium.
    • Spot-VMs unterstützen keine Zeitplanung für die Erfassung.

    Weitere Informationen zur Verwaltung der Planung von Sammlungen finden Sie unter Planung von Sammlungen in TPU-Slices mit mehreren Hosts verwalten.

  • --placement-type=COMPACT: Knotenpool mit aktivierter kompakter Platzierung erstellen. Diese Option muss mit dem Flag --tpu-topology verwendet werden. Weitere Informationen finden Sie unter Richtlinie für kompakte Platzierung erstellen und TPU-Topologie.

Terraform

  1. Prüfen Sie, ob Sie Version 4.84.0 oder höher des Anbieters google verwenden.
  2. Fügen Sie Ihrer Terraform-Konfiguration den folgenden Block hinzu:

    resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
      provider           = google
      project            = PROJECT_ID
      cluster            = CLUSTER_NAME
      name               = POOL_NAME
      location           = CLUSTER_LOCATION
      node_locations     = [NODE_ZONES]
      initial_node_count = NUM_NODES
    
      autoscaling {
        max_node_count = MAX_NODES
        location_policy      = "ANY"
      }
      node_config {
        machine_type = MACHINE_TYPE
        reservation_affinity {
          consume_reservation_type = "SPECIFIC_RESERVATION"
          key = "compute.googleapis.com/reservation-name"
          values = [RESERVATION_LABEL_VALUES]
        }
        spot = true
        flex_start = false
      }
    
      placement_policy {
        type = "COMPACT"
        tpu_topology = TPU_TOPOLOGY
      }
    }
    

    Ersetzen Sie Folgendes:

    • NODE_POOL_RESOURCE_NAME: Der Name der Knotenpoolressource in der Terraform-Vorlage.
    • PROJECT_ID: Ihre Projekt-ID.
    • CLUSTER_NAME: Der Name des vorhandenen Clusters, dem der Knotenpool hinzugefügt werden soll.
    • POOL_NAME: Der Name des Knotenpools, der erstellt werden soll.
    • CLUSTER_LOCATION: Der Compute-Standort des Clusters. Wir empfehlen einen regionalen Cluster für eine höhere Zuverlässigkeit der Kubernetes-Steuerungsebene. Sie können auch einen zonalen Cluster verwenden. Weitere Informationen finden Sie unter TPU-Version und Topologie auswählen.
    • NODE_ZONES: Die durch Kommas getrennte Liste mit einer oder mehreren Zonen, in denen GKE den Knotenpool erstellt.
    • NUM_NODES: Die maximale Anzahl von Knoten im Knotenpool. Sie muss null oder das Produkt der Anzahl der TPU-Chips geteilt durch vier sein, da in TPU-Slices mit mehreren Hosts jeder TPU-Slice-Knoten vier Chips hat. Wenn TPU_TOPOLOGY beispielsweise 4x8 ist, gibt es 32 Chips, d. h., NUM_NODES muss 8 sein. Weitere Informationen zu TPU-Topologien finden Sie in der Tabelle unter TPU-Version auswählen.
    • TPU_TOPOLOGY: Dies gibt die gewünschte physische Topologie für das TPU-Slice an. Das Format der Topologie hängt von der verwendeten TPU-Version ab. Weitere Informationen zu TPU-Topologien finden Sie in der Tabelle unter Topologie auswählen.

    Optional können Sie auch die folgenden Variablen verwenden:

    • RESERVATION_NAME: Wenn Sie die TPU-Reservierung verwenden, ist dies die Liste der Labels der Reservierungsressourcen, die beim Erstellen des Knotenpools verwendet werden sollen. Weitere Informationen zum Ausfüllen der RESERVATION_LABEL_VALUES im Feld reservation_affinity finden Sie unter Terraform-Anbieter.
    • autoscaling: Knotenpool mit aktiviertem Autoscaling erstellen. Wenn GKE einen TPU-Slice-Knotenpool mit mehreren Hosts skaliert, wird der Knotenpool atomar von null auf die maximale Größe hochskaliert.
      • MAX_NODES: Die maximale Größe des Knotenpools. Sie muss gleich dem Produkt aus den Werten in TPU_TOPOLOGY definierten Werte ({A}x{B}x{C}) geteilt durch die Anzahl der Chips in jeder VM sein.
    • spot: Ermöglicht dem Knotenpool, Spot-VMs für die TPU-Slice-Knoten zu verwenden. Dies kann nach dem Erstellen des Knotenpools nicht mehr geändert werden. Weitere Informationen finden Sie unter Spot-VMs.
    • flex_start: Legt fest, dass der Knotenpool die Verbrauchsoption flex-start verwendet. Kann nicht auf true gesetzt werden, wenn spot aktiviert ist.

Console

So erstellen Sie einen Knotenpool mit TPUs:

  1. Öffnen Sie in der Google Cloud Console die Seite Google Kubernetes Engine.

    Zur Seite "Google Kubernetes Engine"

  2. Klicken Sie in der Clusterliste auf den Namen des Clusters, den Sie ändern möchten.

  3. Klicken Sie auf Knotenpool hinzufügen .

  4. Klicken Sie im Abschnitt Knotenpooldetails das Kästchen Knotenstandorte angeben an.

  5. Wählen Sie den Namen der Zone basierend auf der TPU-Version aus, die Sie verwenden möchten. Informationen dazu, wie Sie einen verfügbaren Standort ermitteln, finden Sie unter TPU-Verfügbarkeit in GKE.

  6. Klicken Sie im Navigationsbereich auf Knoten.

  7. Wählen Sie im Abschnitt Maschinenkonfiguration die Option TPUs aus.

  8. Wählen Sie im Drop-down-Menü Serie eine der folgenden Optionen aus:

    • CT3P: Für TPU v3.
    • CT4P: Für TPU v4.
    • CT5LP: Für TPU v5e.
  9. Wählen Sie im Drop-down-Menü Maschinentyp den Namen des Computers aus, der für Knoten verwendet werden soll. In der Tabelle TPU-Version auswählen erfahren Sie, wie Sie den Maschinentyp und die TPU-Topologie definieren, mit denen ein TPU-Slice-Knotenpool mit mehreren Hosts erstellt wird.

  10. Wählen Sie im Drop-down-Menü TPU-Topologie die physische Topologie für das TPU-Slice aus.

  11. Klicken Sie im Dialogfeld Änderungen erforderlich auf Änderungen vornehmen.

  12. Achten Sie darauf, dass als Bootlaufwerktyp entweder Nichtflüchtiger Standardspeicher oder Nichtflüchtiger SSD-Speicher angegeben ist.

  13. Klicken Sie optional das Kästchen Knoten auf Spot-VMs aktivieren an, um Spot-VMs für die Knoten im Knotenpool zu verwenden.

  14. Klicken Sie auf Erstellen.

Knotenpoolstatus prüfen

  1. Rufen Sie die Anmeldedaten ab, damit Sie kubectl für den Zugriff auf den Cluster verwenden können:

    gcloud container clusters get-credentials CLUSTER_NAME \
        --project=PROJECT_ID \
        --location=CONTROL_PLANE_LOCATION
    

    Ersetzen Sie Folgendes:

    • CLUSTER_NAME: Der Name des Clusters.
    • PROJECT_ID: Ihre Projekt-ID.
    • CONTROL_PLANE_LOCATION: Der Compute Engine-Standort der Steuerungsebene des Clusters. Geben Sie für regionale Cluster eine Region und für zonale Cluster eine Zone an.
  2. Verwenden Sie kubectl in Cloud Shell, um Ihre TPU-Slice-Knoten anzuzeigen:

    kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=ACCELERATOR_TYPE \
       -l cloud.google.com/gke-tpu-topology=TPU_TOPOLOGY
    

    Ersetzen Sie Folgendes:

    • TPU_ACCELERATOR: Der Typ des TPU-Beschleunigers, den Sie beim Erstellen der Knotenpools verwendet haben. Beispiel: tpu-v4-podslice oder tpu-v5-lite-podslice.
    • TPU_TOPOLOGY: Die physische Topologie für das TPU-Slice.

    Die Ausgabe sieht in etwa so aus:

     NAME                                    STATUS   ROLES    AGE    VERSION
     gke-tpu-20ee2cce-5tv6                   Ready    <none>   34h     v1.28.1-gke.1066000
    

Multislice-Arbeitslast ausführen

In diesem Abschnitt führen Sie eine JAX-Arbeitslast aus, die die globale Anzahl von TPU-Chips im TPU-Slice anzeigt und dann beendet wird.

So führen Sie eine JAX-Arbeitslast aus:

  1. Erstellen Sie das folgende tpu-multislice.yaml-Manifest:

    Autopilot

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: multislice-job
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
    spec:
      failurePolicy:
        maxRestarts: 4
      replicatedJobs:
        - name: slice
          replicas: NUM_SLICES
          template:
            spec:
              parallelism: NUM_NODES
              completions: NUM_NODES
              backoffLimit: 0
              template:
                spec:
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
                    cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
                  containers:
                  - name: jax-tpu
                    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    - containerPort: 8431
                    command:
                    - bash
                    - -c
                    - |
                      python -c 'import jax; print("Global device count:", jax.device_count())'
                      sleep 60
                    resources:
                     limits:
                        google.com/tpu: NUM_CHIPS
    

    Standard

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: multislice-job
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
    spec:
      failurePolicy:
        maxRestarts: 4
      replicatedJobs:
        - name: slice
          replicas: NUM_SLICES
          template:
            spec:
              parallelism: NUM_NODES
              completions: NUM_NODES
              backoffLimit: 0
              template:
                spec:
                  hostNetwork: true
                  dnsPolicy: ClusterFirstWithHostNet
                  nodeSelector:
                    cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE
                    cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
                  containers:
                  - name: jax-tpu
                    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                    ports:
                    - containerPort: 8471
                    - containerPort: 8080
                    - containerPort: 8431
                    securityContext:
                      privileged: true
                    command:
                    - bash
                    - -c
                    - |
                      python -c 'import jax; print("Global device count:", jax.device_count())'
                      sleep 60
                    resources:
                      limits:
                       google.com/tpu: NUM_CHIPS
    

    Ersetzen Sie Folgendes:

    • NUM_SLICES: Die Anzahl der TPU-Slice-Knotenpools. In diesem Fall ist NUM_SLICES gleich 3.
    • ACCELERATOR_TYPE: Der Typ des TPU-Beschleunigers, den Sie beim Erstellen der Knotenpools verwendet haben. Beispiel: tpu-v4-podslice oder tpu-v5-lite-podslice.
    • TPU_TOPOLOGY: Die physische Topologie für das TPU-Slice. Beispiel: 4x4x4 oder 2x2, je nach TPU-Version.
    • NUM_NODES: Die maximale Anzahl von Knoten im Knotenpool. Sie muss null oder das Produkt der in TPU_TOPOLOGY definierten Werte ({A}x{B}x{C}) geteilt durch die Anzahl der TPU-Chips in jeder VM sein. Bei Multi-Host-TPU v4 beträgt die Anzahl der TPU-Chips in jeder VM vier. Bei Multi-Host-TPU v5e beträgt die Anzahl der TPU-Chips in jeder VM eins, vier oder acht. Wenn Ihre TPU_TOPOLOGY also 2x4x4 ist (TPU v4 mit vier TPU-Chips in jeder VM), beträgt NUM_NODES 32/4, was 8 entspricht.
    • NUM_CHIPS: Bei Multi-Host-TPU v4 beträgt die Anzahl der TPU-Chips in jeder VM vier. Bei Multi-Host-TPU v5e beträgt die Anzahl der TPU-Chips in jeder VM eins, vier oder acht. Weitere Informationen finden Sie unter TPU-Chips auf der VM in einem TPU-Slice.

    In diesem Manifest:

    • Der JobSet ist ein monitorloser Service mit dem gleichen Namen wie der JobSet-Name. In diesem Fall lautet er multislice-job.
    • Mit der Annotation alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool wird die Pod-Affinität konfiguriert, um sicherzustellen, dass alle Pods auf demselben Slice geplant werden.
    • Das Feld image ist auf das neueste JAX AI-Image festgelegt. Wenn Sie eine andere Version festlegen möchten, sehen Sie sich die Liste der aktuellen JAX AI-Images an.
    • maxRestarts: 4 gibt an, wie oft GKE den JobSet neu startet, wenn ein untergeordneter Job fehlschlägt. Wenn die JobSet-Neustarts den maximal definierten Wert erreichen, wird der JobSet als fehlgeschlagen markiert.
    • Die Felder parallelism und completions entsprechen der Anzahl der Knoten in jedem Knotenpool.
    • backoff ist 0, da Multislice nur das synchrone Multi-Controller-Training unterstützt. Für dieses Feld muss 0 festgelegt werden. Wenn bei einem Pod ein Fehler auftritt, schlägt der Job fehl.
    • Die Werte im Abschnitt „Affinität“ sorgen dafür, dass nur eine TPU-Multislic-Arbeitslast in einer Gruppe von Multislices ausgeführt wird.
    • containerPort: 8080 ist der Port für den MXLA-Koordinator.
    • containerPort: 8431 ist der Port zum Exportieren der TPU-Nutzungsmesswerte.
    • Der securityContext: privileged: true gibt an, dass Knoten den privilegierten Modus für den Zugriff auf TPUs aktiviert haben. Für Knoten in GKE-Version 1.28 oder höher muss der privilegierte Modus für den Zugriff auf TPUs nicht aktiviert sein. Weitere Informationen finden Sie unter Container ohne privilegierten Modus ausführen.
  2. Wenden Sie das Manifest an:

    kubectl apply -f tpu-multislice.yaml
    
  3. Prüfen Sie, ob die Arbeitslast zugelassen ist:

    kubectl get jobsets
    

    Die Ausgabe sieht in etwa so aus:

    NAME            RESTARTS   COMPLETED   AGE
    multislice-job                         3s
    
  4. Überwachen Sie den Status der bereitgestellten Pods:

    kubectl get pods
    

    Die Ausgabe sieht etwa so aus:

     NAME                                READY   STATUS      RESTARTS   AGE
     multislice-job-slice-0-0-wzq9t      0/1     Completed   0          2m31s
     multislice-job-slice-0-1-zf4dp      0/1     Completed   0          2m30s
     multislice-job-slice-1-0-hbfn5      0/1     Completed   0          2m31s
     multislice-job-slice-1-1-45fgl      0/1     Completed   0          2m30s
     multislice-job-slice-2-0-wjbp4      0/1     Completed   0          2m30s
     multislice-job-slice-2-1-lwnvs      0/1     Completed   0          2m30s
    

    Der Jobsatz multislice-job plant, erstellt und führt die Pods vollständig aus. Die Pod-Namen haben das Format <jobsetName>-<jobName>-<jobReplicaIndex>-<randomSuffix>. Das Präfix jobsetName bestimmt, zu welchem JobSet der Pod gehört.

  5. Optional: JAX-Arbeitslast entfernen:

    kubectl delete -f tpu-multislice.yaml
    

Zusätzliche Einstellungen konfigurieren

In den folgenden Abschnitten werden die zusätzlichen Konfigurationen beschrieben, die Sie auf Ihr Multislice anwenden können.

Netzwerkleistung mit hostNetwork verbessern

Zur Verbesserung der Netzwerkleistung zwischen TPU-Slices sollten Sie hostNetworking aktivieren. Verwenden Sie hostNetwork: true in Ihrer Pod-Spezifikation, um den gesamten Kubernetes-Netzwerkstack zu überspringen und Ihre Kubernetes-Pods das Hostnetzwerk direkt für die VM-zu-VM-Kommunikation zu verwenden.

Fügen Sie der Pod-Spezifikation die folgenden beiden Zeilen hinzu, um hostNetworking zu aktivieren:

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

Wenn Sie podHostnames weiterhin für die Worker-Knoten-Erkennung mit hostNetwork verwenden möchten, legen Sie dnsPolicy: ClusterFirstWithHostNet fest. Dies ist wichtig, wenn Sie Trainingsjobs automatisch fortsetzen und Sie dieselben Namen für das Aktualisieren derselben Prüfpunkte benötigen.

Wenn Sie TPU Trillium (v6e) verwenden und Ihre Pods hostNetworking nutzen, installieren Sie das folgende DaemonSet zum Optimieren von /proc/sys/net/ipv4/tcp_rmem auf dem Knoten.

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/ai-on-gke/51bf3dcab6ff658cf62cc32867f96860bf58dfdc/scripts/network-setup/v6e-increase-rmem.yaml

Netzwerkleistung ohne hostNetwork auf TPU Trillium verbessern

Wenn Sie TPU Trillium verwenden und Ihre Pods hostNetworking nicht nutzen können, aktivieren Sie die Multi-Netzwerk-Funktion mit dem Modus netdevice, um die beste Netzwerkleistung zu erzielen. Die NIC-Unterstützung im netdevice-Modus mit Multi-Netzwerk übergibt die VM-NIC direkt an den Pod und umgeht Kubernetes und GKE Dataplane V2.

Der Maschinentyp ct6e-standard-4t wird von zwei physischen NICs unterstützt. Kubernetes erfordert eine vNIC, die nicht an Pods übergeben werden kann. Daher muss jeder Knoten drei vNICs haben, damit Pods direkten Zugriff auf zwei vNICs haben und die beste Leistung beider physischen NICs erzielt werden kann.

So aktivieren Sie den netdevice-Modus für ct6e-standard-4t:

  1. Zwei zusätzliche VPCs erstellen, die den netdevice-Modus unterstützen
  2. GKE-Cluster mit Multi-Netzwerk-Funktionen erstellen
  3. Zwei netdevice-Netzwerke konfigurieren Sie können beispielsweise die folgenden GKENetworkParamSet- und Network-Objekte verwenden (SECOND_VPC und THIRD_VPC sind die VPCs, die im vorherigen Schritt erstellt wurden):

    apiVersion: networking.gke.io/v1
    kind: GKENetworkParamSet
    metadata:
      name: tpu-second
    spec:
      vpc: SECOND_VPC
      vpcSubnet: SECOND_VPC_SUBNET
      deviceMode: NetDevice
    ---
    apiVersion: networking.gke.io/v1
    kind: GKENetworkParamSet
    metadata:
      name: tpu-third
    spec:
      vpc: THIRD_VPC
      vpcSubnet: SECOND_VPC_SUBNET
      deviceMode: NetDevice
    ---
    apiVersion: networking.gke.io/v1
    kind: Network
    metadata:
      name: tpu-second
    spec:
      provider: "GKE"
      type: "Device"
      parametersRef:
        group: networking.gke.io
        kind: GKENetworkParamSet
        name: tpu-second
    ---
    apiVersion: networking.gke.io/v1
    kind: Network
    metadata:
      name: tpu-third
    spec:
      provider: "GKE"
      type: "Device"
      parametersRef:
        group: networking.gke.io
        kind: GKENetworkParamSet
        name: tpu-third
    
  4. Pods mit drei Netzwerken verbinden Sie können beispielsweise die folgenden Annotationen in Ihrer Pod-Spezifikation verwenden:

    metadata:
      annotations:
        networking.gke.io/default-interface: 'eth0'
        networking.gke.io/interfaces: |
          [
            {"interfaceName":"eth0","network":"default"},
            {"interfaceName":"eth1","network":"tpu-second"},
            {"interfaceName":"eth2","network":"tpu-third"},
          ]
    
  5. Wenden Sie Netzwerk-Sysctls innerhalb des Pods, im Init-Container oder im Anwendungscontainer an. Sie können der Pod-Spezifikation beispielsweise den folgenden Init-Container hinzufügen:

    initContainers:
    - name: "network-optimization-sysctls"
      image: "busybox"
      securityContext:
        privileged: true
      command:
      - bash
      - -c
      - |
        echo 5000 > /proc/sys/net/ipv4/tcp_rto_min_us
        echo 1 > /proc/sys/net/ipv4/tcp_no_metrics_save
        echo 0 > /proc/sys/net/ipv4/tcp_slow_start_after_idle
        echo 131072 > /proc/sys/net/core/optmem_max
        echo "4096 41943040 314572800" > /proc/sys/net/ipv4/tcp_rmem
    
Best Practice:

Verwenden Sie die Schnittstellen eth1 und eth2 anstelle der Schnittstelle eth0, um eine bessere Netzwerkleistung zu erzielen. Dazu fügen Sie der Arbeitslastspezifikation export LIBTPU_INIT_ARGS="$LIBTPU_INIT_ARGS --megascale_grpc_interface_prefixes=eth1,eth2,lo" hinzu.

Logging aktivieren

Logs, die von Containern ausgegeben werden, die auf GKE-Knoten ausgeführt werden, einschließlich TPU-Slice-Knoten, sind im Log-Explorer sichtbar, wenn Sie GKE-System-Logging in Ihrem Cluster aktiviert haben.

Mit dem folgenden Log-Explorer können Sie Ihre Logs aus GKE aufrufen, um die Containerlogs für Ihre Arbeitslast aufzurufen:

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME

Verwenden Sie für TPU-Slice und -Worker den folgenden Filter:

resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME
resource.labels.pod_name:<jobSetName>-<replicateJobName>-<job-index>-<worker-index>

Weitere Informationen finden Sie unter GKE-TPU-Logs ansehen.

Zusätzliche Messwerte aktivieren

Zusätzlich zu den allgemeinen TPU-Messwerten gibt es vier weitere Multislice-spezifische TPU-Laufzeitmesswerte. Diese Messwerte sind in der GKE-Version 1.29.1-gke.1016000 oder höher verfügbar. TPU-Arbeitslast muss die JAX-Version 0.4.24 verwenden

Im Folgenden sind die folgenden Multislice-Messwerte verfügbar:

  • DCN-Übertragungslatenzen (Data Center Network): Verteilung der Netzwerkübertragungslatenzen für Multi-Slice-Traffic.
  • Gesamtlatenzen: Verteilung der Gesamt-End-to-End-Latenz für Multi-Slice-Traffic.
  • Host-zu-Gerät-Übertragungslatenzen: Verteilung der Host-zu-Gerät-Übertragungslatenz für jeden Datenblock für Multi-Slice-Traffic.
  • Übertragungslatenzen von Gerät zu Host: Verteilung der Geräte-zu-Host-Übertragungslatenz für jeden Datenblock für Multi-Slice-Traffic.

Diese Messwerte befinden sich im Kubernetes-Containerschema (k8s_container):

  • kubernetes.io/container/multislice/network/dcn_transfer_latencies
  • kubernetes.io/container/multislice/network/collective_end_to_end_latencies
  • kubernetes.io/container/multislice/accelerator/host_to_device_transfer_latencies
  • kubernetes.io/container/multislice/accelerator/device_to_host_transfer_latencies

TPU-Slice im Vergleich zu Multislice

In der folgenden Tabelle wird die Architektur eines TPU-Slices und eines Multislices unterschieden:

TPU-Slice Multi-Slice
Verbindung Die Arbeitslast wird auf einem einzelnen TPU-Slice ausgeführt. Alle TPU-Chips in einem Segment sind mit ICI verbunden. Die Arbeitslast wird auf mehreren TPU-Slices ausgeführt. Die Kommunikation innerhalb eines Segments erfolgt über ICI. Die Kommunikation zwischen den Segmenten erfolgt über DCN.
Unterstützte Knotenpools TPU-Slice mit einzelnem Host und TPU mit mehreren Hosts Gruppen von TPU-Slices mit mehreren Hosts
Empfohlener Arbeitslasttyp Indexierter Job oder JobSet JobSet

Ressourcen bereinigen

Am einfachsten vermeiden Sie weitere Kosten, wenn Sie das für die Anleitung erstellte Projekt löschen. Google Cloud Alternativ können Sie die einzelnen Ressourcen löschen.

Projekt löschen

    Delete a Google Cloud project:

    gcloud projects delete PROJECT_ID

Einzelne Ressourcen löschen

Löschen Sie den GKE-Cluster:

```sh
gcloud container clusters delete  CLUSTER_NAME \
   --project=PROJECT_ID  \
   --location=CONTROL_PLANE_LOCATION
```

Nächste Schritte