esegui il deployment di più sezioni TPU in GKE

Questa pagina mostra come eseguire il deployment dei carichi di lavoro in Google Kubernetes Engine (GKE) utilizzando la configurazione multisezione di Cloud TPU per per l'addestramento su vasta scala.

Prima di configurare la funzionalità Multisezione in In GKE, è bene acquisire familiarità con i seguenti concetti:

  1. Introduzione a Cloud TPU
  2. Architettura di sistema Cloud TPU
  3. Informazioni sulle TPU in GKE
di Gemini Advanced.

Che cos'è TPU Multisezione

TPU Multislice è l'organizzazione architetturale delle VM in una sezione TPU in cui o più sezioni di Cloud TPU comunicano tramite la rete di data center (DCN). Multislice supporta full stack, dai costi contenuti e su larga scala con scalabilità quasi lineare fino a decine di migliaia di chip TPU. In una configurazione multisezione, GKE esegue il deployment di un carico di lavoro con più sezioni su più TPU sezioni. La comunicazione tra i chip TPU all'interno di una sezione avviene su inter chip interconnessioni (ICI). La comunicazione tra le sezioni avviene attraverso la rete DCN.

Ti consigliamo di utilizzare la multisezione se il job è troppo grande per essere singola sezione TPU.

Disponibilità di più sezioni in GKE

  • La versione standard supporta più sezioni nella versione 1.27.4-gke.900 e versioni successive.
  • Autopilot supporta più sezioni nella versione 1.29.2-gke.1521000 e versioni successive.
  • Multislice supporta i framework JAX e PyTorch. Numero minimo di file JAX supportati è la 2.1.
  • Più sezioni supporta solo nodo della sezione TPU multi-host piscine. Ad esempio, non puoi utilizzare Sezione multipla con ct4p-hightpu-4t con una topologia 2x2x1 o una ct5lp-hightpu-4t con una topologia 2x2, perché si tratta di una sezione TPU con host singolo pool di nodi.
  • Multislice supporta solo l'addestramento sincrono multicontroller.
  • I carichi di lavoro con più sezioni possono essere eseguiti solo su sezioni TPU che condividono la stessa TPU tipo, dimensione e topologia.

Prima di iniziare

Prima di iniziare, assicurati di aver eseguito le seguenti attività:

  • Attiva l'API Google Kubernetes Engine.
    • Abilita l'API Google Kubernetes Engine .
  • Se vuoi utilizzare Google Cloud CLI per questa attività, install e poi inizializzare con gcloud CLI. Se hai già installato gcloud CLI, scarica la versione più recente eseguendo gcloud components update.

Esegui un carico di lavoro su una sezione multipla

Questa sezione mostra come eseguire un carico di lavoro su una sezione multipla. Se utilizzi dalla modalità GKE Autopilot, passa Sezione Esegui un carico di lavoro su più sezioni. Autopilot i cluster che eseguono la versione 1.29.2-gke.1521000 o successive abilitano le TPU per impostazione predefinita.

prepara un pool di nodi in modalità Standard

Questa sezione illustra i seguenti passaggi:

  1. Crea tre pool di nodi delle sezioni TPU multi-host
  2. Verifica lo stato del pool di nodi

Crea il pool di nodi della sezione TPU

Puoi creare più di un pool di nodi della sezione TPU multi-host. Ai fini di questo crea tre pool di nodi delle sezioni TPU multi-host per eseguire un carico di lavoro Multislice. Puoi creare un Nodo sezione TPU multi-host utilizzando Google Cloud CLI, Terraform la console Google Cloud.

gcloud

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]

Sostituisci quanto segue:

  • POOL_NAME: il nome del nuovo pool di nodi.

  • LOCATION: il nome della zona basato su la versione di TPU che vuoi utilizzare:

    • Per TPU v4, usa us-central2-b.
    • I tipi di macchine TPU v5e che iniziano con ct5l- non sono mai multi-host.
    • Per i tipi di macchine TPU v5e che iniziano con ct5lp-, usa us-west1-c, us-west4-a, us-west4-b, us-central1-a, us-east1-c, us-east5-b o europe-west4-a.
    • Per i tipi di macchine TPU v5p che iniziano con ct5p-, usa us-east1-d, us-east5-a o us-east5-c.

    Per scoprire di più, consulta Disponibilità delle TPU in GKE.

  • CLUSTER_NAME: il nome del cluster.

  • NODE_ZONE: l'elenco separato da virgole di una o più zone in cui GKE crea il pool di nodi.

  • MACHINE_TYPE: il tipo di macchina da utilizzare per nodi. Per saperne di più sui tipi di macchine disponibili, consulta Mappatura della configurazione di TPU.

  • TPU_TOPOLOGY: la posizione fisica per la sezione TPU. Il formato della topologia dipende dalla TPU dell'oggetto come segue:

    • TPU v4 o v5p: definisci la topologia in 3 tuple ({A}x{B}x{C}), ad esempio 4x4x4.
    • TPU v5e: definisci la topologia in due tuple ({A}x{B}), ad esempio 2x2.

    Per scoprire di più, consulta Topologia.

  • NUM_NODES: il numero di nodi nel pool di nodi. Deve essere zero o il prodotto dei valori definiti in TPU_TOPOLOGY ({A}x{B}x{C}) diviso per il numero di chip in ogni VM. Per TPU v4 multi-host e TPU v5e, il numero di chip in ogni la VM è quattro. Pertanto, se il valore TPU_TOPOLOGY è 2x4x4 (TPU v4 con quattro chip in ogni VM), quindi NUM_NODES è 32/4, che equivale a 8.

Facoltativamente, puoi anche utilizzare i seguenti flag:

  • RESERVATION_NAME: il nome del che GKE utilizza durante la creazione del pool di nodi. Se ometti questo flag, GKE utilizza i pool di nodi delle sezioni TPU disponibili. Per saperne di più sulle prenotazioni TPU, consulta la sezione Prenotazione TPU.
  • --spot: imposta il pool di nodi per l'utilizzo delle VM spot dei nodi della sezione TPU. Questo valore non può essere modificato dopo la creazione del pool di nodi. Per maggiori informazioni le informazioni, vedi VM spot.
  • --enable-autoscaling: crea un pool di nodi con scalabilità automatica abilitata. Quando GKE scala un pool di nodi di sezioni TPU multi-host, fa lo scale up atomicamente del pool di nodi da zero alla dimensione massima.
    • MAX_NODES: la dimensione massima del nodo piscina. Il flag --max-nodes è obbligatorio se viene fornito --enable-autoscaling e deve essere uguale al prodotto dei valori definiti in TPU_TOPOLOGY ({A}x{B}x{C}) diviso per il numero di chip in ogni VM.

Terraform

  1. Assicurati di utilizzare la versione 4.84.0 o successiva del google o il provider di servizi di terze parti.

  2. Aggiungi il blocco seguente alla tua configurazione Terraform:

    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
  }

  placement_policy {
    type = "COMPACT"
    tpu_topology = TPU_TOPOLOGY
  }
}

    Sostituisci quanto segue:

    • NODE_POOL_RESOURCE_NAME: il nome del nodo risorsa del pool nel modello Terraform.
    • PROJECT_ID: l'ID progetto.
    • CLUSTER_NAME: il nome del cluster esistente a cui aggiungere il pool di nodi.
    • POOL_NAME: il nome del pool di nodi da creare.
    • CLUSTER_LOCATION: località di calcolo per in un cluster Kubernetes. Consigliamo di avere un cluster a livello di regione per una maggiore affidabilità il piano di controllo Kubernetes. Puoi anche usare un cluster a livello di zona. Per saperne di più, vedi Seleziona una versione e una topologia di TPU.
    • NODE_ZONES: l'elenco separato da virgole di una o più zone in cui GKE crea il pool di nodi.
    • NUM_NODES: il numero di nodi nel nodo piscina. Deve essere zero o il prodotto del numero di chip TPU diviso per quattro, poiché nelle sezioni TPU multi-host ogni nodo TPU ha 4 chip. Ad esempio, se TPU_TOPOLOGY è 4x8, Ora ci sono 32 chip, il che significa NUM_NODES deve essere 8. Per saperne di più sulle topologie TPU, utilizza la tabella in Mappatura della configurazione di TPU.
    • TPU_TOPOLOGY: indica lo stato desiderato topologia fisica per la sezione TPU. Il formato della topologia dipende dalla versione della TPU in uso:
      • Per TPU v4: definisci la topologia in 3 tuple ({A}x{B}x{C}), per esempio 4x4x4.
      • Per TPU v5e: definisci la topologia in 2 tuple ({A}x{B}), per esempio 2x2.

    Facoltativamente, puoi anche utilizzare le seguenti variabili:

    • RESERVATION_NAME: se utilizzi TPU prenotazione, questo è l'elenco delle etichette delle risorse di prenotazione da utilizzare quando la creazione del pool di nodi. Per scoprire di più su come compilare RESERVATION_LABEL_VALUES nel Campo reservation_affinity, vedi Fornitore Terraform.
    • autoscaling: crea un pool di nodi con scalabilità automatica abilitata. Quando GKE scala un pool di nodi di sezioni TPU multi-host, fa lo scale up atomicamente del pool di nodi da zero alla dimensione massima.
      • MAX_NODES: è la dimensione massima del nodo piscina. Deve essere uguale al prodotto dei valori definiti in TPU_TOPOLOGY ({A}x{B}x{C}) diviso per il numero di chip in ogni VM.
    • spot: consente al pool di nodi di utilizzare le VM spot per i nodi della sezione TPU. Questo valore non può essere modificato dopo la creazione del pool di nodi. Per ulteriori informazioni, vedi VM spot.

Console

Per creare un pool di nodi con TPU:

  1. Vai alla pagina Google Kubernetes Engine nella console Google Cloud.

    Vai a Google Kubernetes Engine

  2. Nell'elenco dei cluster, fai clic sul nome del cluster da modificare.

  3. Fai clic su Aggiungi pool di nodi.

  4. Nella sezione Dettagli del pool di nodi, seleziona la casella Specifica le località dei nodi.

  5. Seleziona la zona in base alla versione della TPU che vuoi utilizzare:

    • Per TPU v4, usa us-central2-b.
    • I tipi di macchine TPU v5e che iniziano con ct5l- non sono mai multi-host.
    • Per i tipi di macchine TPU v5e che iniziano con ct5lp-, usa us-west1-c, us-west4-a, us-west4-b, us-central1-a, us-east1-c, us-east5-b o europe-west4-a.
    • Per i tipi di macchine TPU v5p che iniziano con ct5p-, usa us-east1-d, us-east5-a o us-east5-c.

  6. Nel riquadro di navigazione, fai clic su Nodi.

  7. Nella sezione Configurazione macchina, seleziona TPU.

  8. Nel menu a discesa Serie, seleziona una delle seguenti opzioni:

    • CT4P: per TPU v4.
    • CT5LP: per TPU v5e.

  9. Nel menu a discesa Tipo di macchina, seleziona il nome della macchina da utilizzare per nodi. Utilizza la Tabella Mappatura della configurazione di TPU per scoprire come definire il tipo di macchina e la topologia TPU che creano un pool di nodi di sezioni TPU multi-host.

  10. Nel menu a discesa Topologia TPU, seleziona la topologia fisica per la sezione TPU.

  11. Nella finestra di dialogo Modifiche necessarie, fai clic su Apporta modifiche.

  12. Assicurati che il Tipo di disco di avvio sia Disco permanente standard o Disco permanente SSD.

  13. (Facoltativo) Seleziona la casella di controllo Abilita i nodi sulle VM Spot per utilizzare Spot VM per i nodi pool di nodi.

  14. Fai clic su Crea.

Verifica lo stato del pool di nodi

  1. Richiedi le credenziali per poter utilizzare kubectl per accedere al cluster:

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

    Sostituisci quanto segue:

    • CLUSTER_NAME: il nome del cluster.
    • PROJECT_ID: l'ID progetto.

  2. Usa kubectl, in Cloud Shell, per visualizzare i nodi delle sezioni TPU:

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

    Sostituisci quanto segue:

    • TPU_ACCELERATOR: il tipo di acceleratore TPU utilizzato quando hai creato i pool di nodi. Ad esempio, tpu-v4-podslice, tpu-v5-lite-device o tpu-v5-lite-podslice.
    • TPU_TOPOLOGY: il topologia fisica per la sezione TPU.

    L'output è simile al seguente:

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

Esegui un carico di lavoro su più sezioni

In questa sezione eseguirai un carico di lavoro JAX che mostra il numero globale di chip nella sezione TPU e poi si chiude.

Per eseguire un carico di lavoro JAX:

  1. Crea il seguente manifest tpu-multislice.yaml:

    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

    Sostituisci quanto segue:

    • NUM_SLICES: il numero di pool di nodi della sezione TPU. In questo Maiuscole/minuscole, NUM_SLICES è uguale a 3.
    • ACCELERATOR_TYPE: il tipo di acceleratore TPU utilizzato durante la creazione dei pool di nodi. Ad esempio, tpu-v4-podslice, tpu-v5-lite-device o tpu-v5-lite-podslice.
    • TPU_TOPOLOGY: la topologia fisica per la sezione TPU. Ad esempio 4x4x4 o 2x2, a seconda della versione della TPU.
    • NUM_NODES: il numero di nodi nel pool di nodi. Deve essere zero o il prodotto dei valori definiti in TPU_TOPOLOGY ({A}x{B}x{C}) diviso per il numero di chip TPU in ogni VM. Per TPU multi-host v4, il numero di TPU in ogni VM sono quattro. Per TPU v5e multi-host, il numero di TPU in ogni VM è uno, quattro o otto. Pertanto, se TPU_TOPOLOGY è 2x4x4 (TPU v4 con quattro chip TPU in ogni VM), NUM_NODES è 32/4, che equivale a 8.
    • NUM_CHIPS: per TPU v4 multi-host, il numero di TPU in ogni VM sono quattro. Per TPU v5e multi-host, il numero di TPU in ogni VM è uno, quattro o otto. Per saperne di più, consulta la sezione Chip TPU sulla VM in una sezione TPU.

    In questo manifest:

    • Il JobSet è un servizio headless con lo stesso nome del nome JobSet, in in questo caso è multislice-job.
    • maxRestarts: 4 indica il numero massimo di volte GKE riavvia il JobSet in caso di errore di un job secondario. Se Il riavvio del set di job raggiunge il valore massimo definito, quindi il set di job viene contrassegnato come non riuscito.
    • I campi parallelism e completions equivalgono al numero di nodi in per ogni pool di nodi.
    • backoff è 0 perché l'opzione multisezione supporta solo la funzione sincrona l'addestramento multi-controller. Deve essere impostato su 0. L'esecuzione del job non viene completata correttamente quando un pod non riesce.
    • I valori nella sezione di affinità assicurano che sia presente una sola TPU Carico di lavoro su più sezioni in esecuzione in un gruppo di sezioni multiple.
    • containerPort: 8080 è il trasferimento per il coordinatore MXLA
    • containerPort: 8431 è la porta per esportare le metriche di utilizzo della TPU
    • securityContext: privileged: true indica che i nodi sono con privilegi abilitati per l'accesso alle TPU. Nodi in GKE La versione 1.28 o successive non richiede l'abilitazione della modalità con privilegi per accedere TPU. Per saperne di più, vedi Esegui i container senza modalità con privilegi.

  2. Applica il manifest:

    kubectl apply -f tpu-multislice.yaml

  3. Verifica che il carico di lavoro sia ammesso:

    kubectl get jobsets

    L'output è simile al seguente:

    NAME            RESTARTS   COMPLETED   AGE
multislice-job                         3s

  4. Monitora lo stato dei pod di cui è stato eseguito il provisioning:

    kubectl get pods

    L'output è simile al seguente:

     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

Il set di job multislice-job pianifica, crea ed esegue i pod fino al completamento. I nomi dei pod sono nel formato <jobsetName>-<jobName>-<jobReplicaIndex>-<randomSuffix>. jobsetName determina il JobSet a cui appartiene il pod.

Configurazioni aggiuntive

Le seguenti sezioni descrivono le configurazioni aggiuntive che puoi applicare il tuo Multislice.

Abilita hostNetwork sui pod GKE Standard

Per migliorare le prestazioni di rete tra le sezioni TPU, ti consigliamo di attivare hostNetworking. Utilizza hostNetwork: true nelle specifiche del pod per saltare tutte le lo stack di networking di Kubernetes e consenti ai tuoi pod Kubernetes di utilizzare la rete host per la comunicazione da VM a VM.

Per attivare hostNetworking, rimuovi le due righe seguenti dalle specifiche del pod:

hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet

Per continuare a utilizzare podHostnames per il rilevamento dei nodi worker con hostNetwork, imposta dnsPolicy: ClusterFirstWithHostNet. Questo è importante quando esegui il ripristino automatico di addestramento e devi avere gli stessi nomi per ricaricare dei checkpoint.

Logging

I log emessi dai container in esecuzione sui nodi GKE, inclusi i nodi delle sezioni TPU, sono visibili in Esplora log, se dispongono del logging di sistema GKE abilitato in in un cluster Kubernetes.

Puoi visualizzare i tuoi log da GKE utilizzando il Esplora log con il seguente filtro per visualizzare i log dei container per carico di lavoro:

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

Utilizza il seguente filtro per la sezione TPU e i worker:

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>

Per saperne di più, consulta Visualizzare i log di GKE TPU.

Osservabilità e metriche

Oltre alle metriche TPU generali, sono disponibili altre quattro metriche di runtime TPU specifiche per più sezioni. Queste metriche sono disponibili in GKE versione 1.29.1-gke.1016000 o successive. Il carico di lavoro TPU deve utilizzare la versione JAX 0.4.24

Di seguito sono riportate le metriche disponibili in più sezioni:

  • Latenze di trasferimento della rete DCN (Data Center Network): distribuzione delle latenze di trasferimento di rete per il traffico su più sezioni.
  • Latenze collettive: distribuzione della latenza collettiva end-to-end per il traffico su più sezioni.
  • Latenze di trasferimento da host a dispositivo: distribuzione della latenza di trasferimento da host a dispositivo per ogni blocco di dati per il traffico su più sezioni.
  • Latenze di trasferimento da dispositivo a host: distribuzione della latenza di trasferimento dal dispositivo all'host per ogni blocco di dati per il traffico multisezione.

Queste metriche si trovano nello schema del container Kubernetes (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

Sezione TPU e più sezioni

La tabella seguente differenzia l'organizzazione architetturale di una sezione TPU e una multisezione:

Sezione TPU Multislice
Interconnessione Il carico di lavoro viene eseguito su una singola sezione TPU. Tutti i chip TPU in una sezione sono collegati a ICI. Il carico di lavoro viene eseguito su più sezioni TPU. La comunicazione all'interno di una sezione avviene tramite ICI. La comunicazione tra le sezioni avviene tramite DCN.
Pool di nodi supportati Sezione TPU con host singolo e Sezione TPU multi-host Gruppi di sezioni TPU multi-host
Tipo di carico di lavoro consigliato IndexedJob o JobSet JobSet

Passaggi successivi