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.

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 Knotenpools mit mehreren Hosts. Sie können Multi-Slice beispielsweise nicht mit einem ct4p-hightpu-4t mit einer 2x2x1 Topologie oder ein ct5lp-hightpu-4t mit einer 2x2 Topologie nutzen, 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.

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.

• Erstellen Sie einen Standard- oder einen Autopilot-Cluster, der eine Version ausführt, die Multislice unterstützt. Unterstützte Versionen finden Sie unter Multi-Slice-Verfügbarkeit in GKE.
• Prüfen Sie, ob Ihr Projekt ein ausreichendes Kontingent für Cloud TPU in GKE hat.
• Installieren Sie JobSet v0.2.3 oder höher.

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 container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONE \
    --machine-type=MACHINE_TYPE \
    --tpu-topology=TPU_TOPOLOGY \
    --num-nodes=NUM_NODES \
    [--spot \]
    [--enable-autoscaling \
      --max-nodes MAX_NODES]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME]

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
   

   Ersetzen Sie Folgendes:

   • CLUSTER_NAME: Der Name des Clusters.
   • PROJECT_ID: Ihre Projekt-ID.

2. Verwenden Sie kubectl in Cloud Shell, um Ihre TPU-Slice-Knoten anzuzeigen:

   kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=TPU_ACCELERATOR \
      -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, tpu-v5-lite-device 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: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   - containerPort: 8431
                   command:
                   - bash
                   - -c
                   - |
                     pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                     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: python:3.8
                   ports:
                   - containerPort: 8471
                   - containerPort: 8080
                   - containerPort: 8431
                   securityContext:
                     privileged: true
                   command:
                   - bash
                   - -c
                   - |
                     pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                     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, tpu-v5-lite-device 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 Anmerkung alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool wird die Pod-Affinität konfiguriert, damit alle Pods auf demselben Slice geplant werden.
   • 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 in 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.

Zusätzliche Konfigurationen

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

hostNetwork auf Ihren GKE-Standard-Pods aktivieren

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.

Entfernen Sie die folgenden beiden Zeilen aus der Pod-Spezifikation, 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.

Logging

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.

Beobachtbarkeit und Messwerte

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:

Nächste Schritte

• Multislice-Arbeitslasten mit TPU-Slices orchestrieren