Esegui il deployment di TPU Multislices in GKE

Questa pagina mostra come eseguire il deployment dei carichi di lavoro in Google Kubernetes Engine (GKE) utilizzando la configurazione Cloud TPU Multislice per un addestramento su larga scala e conveniente.

Questo tutorial è rivolto a ingegneri di machine learning (ML) e amministratori e operatori di piattaforme che vogliono utilizzare l'orchestrazione dei container Kubernetes per gestire l'addestramento, l'ottimizzazione e l'inferenza di modelli su larga scala utilizzando le TPU. Per saperne di più sui ruoli comuni e sulle attività di esempio a cui viene fatto riferimento nei contenuti di Google Cloud , consulta Ruoli utente e attività comuni di GKE.

Prima di configurare Multislice in GKE, assicurati di conoscere i seguenti concetti:

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

Che cos'è TPU Multislice

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

Ti consigliamo di utilizzare Multislice se il tuo job è troppo grande per essere inserito in una singola slice TPU.

Disponibilità di Multislice in GKE

  • Standard supporta Multislice nella versione 1.27.4-gke.900 e successive.
  • Autopilot supporta Multislice nella versione 1.29.2-gke.1521000 e successive.
  • Multislice supporta i framework JAX e PyTorch. La versione minima supportata di JAX è 2.1.
  • Multislice supporta solo i node pool TPU multi-host. Ad esempio, non puoi utilizzare Multislice con una TPU ct4p-hightpu-4t con una topologia 2x2x1 o una TPU ct5lp-hightpu-4t con una topologia 2x2, perché si tratta di pool di nodi di slice TPU a singolo host.
  • Multislice supporta solo l'addestramento sincrono multicontroller.
  • I carichi di lavoro multi-slice possono essere eseguiti solo su slice TPU che condividono lo stesso tipo, dimensione e topologia di TPU.
  • Multislice non supporta TPU v3.

Prima di iniziare

Prima di iniziare, assicurati di aver eseguito le seguenti operazioni:

  • Attiva l'API Google Kubernetes Engine.
    • Attiva l'API Google Kubernetes Engine
  • Se vuoi utilizzare Google Cloud CLI per questa attività, installala e poi inizializza gcloud CLI. Se hai già installato gcloud CLI, scarica l'ultima versione eseguendo il comando gcloud components update. Le versioni precedenti di gcloud CLI potrebbero non supportare l'esecuzione dei comandi in questo documento.

Scegliere un metodo di implementazione di GKE

Puoi creare un cluster che utilizza le TPU utilizzando Google Cloud CLI o un kit di elaborazione accelerata (XPK).

  • Utilizza gcloud CLI per creare manualmente l'istanza del cluster GKE per la personalizzazione o l'espansione precise degli ambienti GKE di produzione esistenti.
  • Utilizza XPK per creare rapidamente cluster GKE ed eseguire workload per prove concettuali e test. Per saperne di più, consulta il file README di XPK.

Crea manualmente il node pool GKE

Questa sezione mostra come preparare un pool di nodi in modalità Standard per eseguire un carico di lavoro su un Multislice. Se utilizzi la modalità GKE Autopilot, vai alla sezione Esegui un workload multislice. I cluster Autopilot che eseguono la versione 1.29.2-gke.1521000 o successive abilitano le TPU per impostazione predefinita.

Questa sezione illustra i seguenti passaggi:

  1. Crea tre node pool TPU slice multi-host
  2. Verifica lo stato del node pool
  3. Esegui un carico di lavoro multislice

Crea il node pool di slice TPU

Puoi creare più di un node pool TPU multi-host. Ai fini di questa guida, crea tre node pool TPU slice multi-host per eseguire un carico di lavoro Multislice. I passaggi per creare un node pool TPU slice multi-host variano a seconda che utilizzi Ironwood (TPU7x) o una versione precedente della TPU.

Ironwood (TPU7x)

Puoi creare un node pool di una slice TPU multi-host nella versione Ironwood (TPU7x) utilizzando Google Cloud CLI o Terraform:

gcloud

Per creare un pool di nodi TPU slice multi-host con Ironwood (TPU7x), devi prima creare una policy del carico di lavoro.

  1. Crea una policy del workload:

    gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
    --type=HIGH_THROUGHPUT \
    --accelerator-topology=TPU_TOPOLOGY \
    --project=PROJECT_ID \
    --region=REGION

    Sostituisci quanto segue:

    • WORKLOAD_POLICY_NAME: un nome per la policy del carico di lavoro.
    • TPU_TOPOLOGY: la topologia TPU Ironwood (TPU7x). Ad esempio, 2x2x2. Per visualizzare tutte le topologie Ironwood (TPU7x) supportate, consulta la sezione relativa alla topologia.
    • PROJECT_ID: il tuo ID progetto Google Cloud .
    • REGION: la regione della policy del workload. Una policy del workload è una risorsa di regione e può essere riutilizzata in pool di nodi che condividono la stessa topologia.

  2. Crea il pool di nodi con la policy del workload:

    gcloud container node-pools create NODE_POOL_NAME \
  --cluster=CLUSTER_NAME \
  --machine-type=tpu7x-standard-4t \
  --placement-policy=WORKLOAD_POLICY_NAME \
  --location=CONTROL_PLANE_LOCATION \
  --node-locations=NODE_ZONE \
  --project=PROJECT_ID \
  --reservation=RESERVATION_NAME \
  --reservation-affinity=specific

    Sostituisci quanto segue:

    • NODE_POOL_NAME: il nome del nuovo node pool.
    • CLUSTER_NAME: il nome del tuo cluster GKE.
    • WORKLOAD_POLICY_NAME: il nome della policy del workload che hai creato.
    • CONTROL_PLANE_LOCATION: la posizione di Compute Engine del control plane del tuo cluster. Fornisci una regione per i cluster regionali o una zona per i cluster zonali.
    • NODE_ZONE: il nome della zona in base alla versione della TPU che vuoi utilizzare.

      (Facoltativo) Utilizza una zona AI, ad esempio us-central1-ai1a. Le zone AI sono località specializzate ottimizzate per i workload AI/ML all'interno delle regioni Google Cloud .

      Per identificare una località disponibile, consulta Disponibilità di TPU in GKE.
    • PROJECT_ID: il tuo ID progetto Google Cloud .
    • RESERVATION_NAME: il nome della prenotazione da utilizzare.

    In questo comando, il flag --tpu-topology è stato sostituito dal flag --placement-policy.

Terraform

  1. Assicurati di utilizzare la versione 4.84.0 o successive del provider google.

  2. Crea una policy del workload:

    resource "google_compute_resource_policy" {
  name   = "WORKLOAD_POLICY_NAME"
  region = CLUSTER_LOCATION
  workload_policy {
    type = "HIGH_THROUGHPUT"
    accelerator_topology = "TPU_TOPOLOGY"
  }
}

    Sostituisci quanto segue:

    • WORKLOAD_POLICY_NAME: un nome per la policy del workload.
    • CLUSTER_LOCATION: la località di computing per il cluster. Consigliamo di utilizzare un cluster regionale per una maggiore affidabilità del control plane Kubernetes. Puoi anche utilizzare un cluster zonale. Per maggiori informazioni, vedi Selezionare una versione e una topologia TPU.
    • TPU_TOPOLOGY: la topologia TPU Ironwood (TPU7x). Ad esempio, 2x2x2. Per visualizzare tutte le topologie Ironwood (TPU7x) supportate, consulta la sezione Pianificare le TPU.

    Per saperne di più sul riferimento google_compute_resource_policy, consulta Provider Terraform.

  3. Nella configurazione di Terraform, aggiungi il seguente blocco:

    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]
    }
    flex_start = false
    spot = true
  }

  placement_policy {
    policy_name = WORKLOAD_POLICY_NAME
  }
}

    Sostituisci quanto segue:

    • NODE_POOL_RESOURCE_NAME: il nome della risorsa del node pool nel modello Terraform.
    • PROJECT_ID: il tuo ID progetto.
    • CLUSTER_NAME: il nome del cluster esistente a cui aggiungere il pool di nodi.
    • POOL_NAME: il nome del node pool da creare.
    • NODE_ZONES: l'elenco separato da virgole di una o più zone in cui GKE crea il node pool.

      (Facoltativo) Utilizza una zona AI, ad esempio us-central1-ai1a. Le zone AI sono località specializzate ottimizzate per i workload AI/ML all'interno delle regioni Google Cloud .

    • NUM_NODES: il numero di nodi nel node pool. Deve essere zero o il prodotto del numero di chip TPU diviso per quattro, perché nelle sezioni TPU multihost ogni nodo della sezione TPU ha quattro chip. Ad esempio, se TPU_TOPOLOGY è 4x8, allora ci sono 32 chip, il che significa che NUM_NODES deve essere 8. Per scoprire di più sulle topologie TPU, utilizza la tabella in Scegliere la versione della TPU.
    • TPU_TOPOLOGY: indica la topologia fisica selezionata per la sezione TPU. Il formato della topologia dipende dalla versione della TPU che stai utilizzando. Per saperne di più sulle topologie TPU, utilizza la tabella in Scegliere una topologia.

    Se vuoi, puoi utilizzare anche le seguenti variabili:

    • RESERVATION_NAME: se utilizzi una prenotazione TPU, fornisci un elenco di etichette delle risorse di prenotazione da utilizzare durante la creazione del node pool. Per saperne di più su come compilare RESERVATION_LABEL_VALUES nel campo reservation_affinity, consulta Provider Terraform.
    • autoscaling: crea un node pool con la scalabilità automatica abilitata. Quando GKE scala un node pool di sezioni TPU multi-host, lo scala in modo atomico da zero alle dimensioni massime.
      • MAX_NODES: la dimensione massima del node pool. Il valore 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. Ad esempio, se TPU_TOPOLOGY è 2x2x2, il prodotto è 8. Poiché ogni VM in tpu7x-standard-4t ha 4 chip, il numero di nodi è 2.
    • spot: il pool di nodi che utilizzerà le VM spot per i nodi dello slice TPU. Questa impostazione non può essere modificata dopo la creazione del pool di nodi. Per saperne di più, consulta VM spot.
    • flex_start: il pool di nodi che utilizzerà l'opzione di consumo avvio flessibile. Questa impostazione non può essere impostata su true se spot è attivato.

Altre versioni di TPU

Puoi creare un pool di nodi slice TPU multihost nelle versioni v3, v4, v5p, v5e e Trillium (v6e) utilizzando Google Cloud CLI, Terraform o la console Google Cloud .

gcloud

  gcloud container node-pools create POOL_NAME \
      --location=CONTROL_PLANE_LOCATION \
      --cluster=CLUSTER_NAME \
      --node-locations=NODE_ZONE \
      --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]

Sostituisci quanto segue:

  • POOL_NAME: il nome del nuovo pool di nodi.
  • CONTROL_PLANE_LOCATION: il nome della zona in base alla versione della TPU che vuoi utilizzare. Per identificare una località disponibile, consulta Disponibilità di TPU in GKE.
  • CLUSTER_NAME: il nome del cluster.
  • NODE_ZONES: l'elenco separato da virgole di una o più zone in cui GKE crea il node pool.

    (Facoltativo) Utilizza una zona AI, ad esempio us-central1-ai1a. Le zone AI sono località specializzate ottimizzate per i workload AI/ML all'interno delle regioni Google Cloud .

  • MACHINE_TYPE: il tipo di macchina da utilizzare per i nodi. Per scoprire di più sui tipi di macchina disponibili, vedi Scegliere la versione della TPU.

  • TPU_TOPOLOGY: la topologia fisica della sezione TPU. Il formato della topologia dipende dalla versione della TPU. Per saperne di più sulle topologie TPU, utilizza la tabella in Scegliere una topologia.

    Per saperne di più, consulta Topologia.

    Se vuoi, puoi anche utilizzare i seguenti flag:

  • NUM_NODES: il numero di nodi nel node pool. 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 e TPU v5e multi-host, il numero di chip in ogni VM è quattro. Pertanto, se il tuo TPU_TOPOLOGY è 2x4x4 (TPU v4 con quattro chip in ogni VM), allora il NUM_NODES è 32/4, ovvero 8. Se ometti questo flag, il numero di nodi viene calcolato e impostato come predefinito in base alla topologia e al tipo di macchina.

  • RESERVATION_NAME: il nome della prenotazione utilizzata da GKE durante la creazione del node pool. Se ometti questo flag, GKE utilizza i node pool delle sezioni TPU disponibili. Per saperne di più sulle prenotazioni TPU, consulta Prenotazione TPU.

  • --spot: imposta il pool di nodi in modo che utilizzi le VM spot per i nodi dello slice TPU. Questo valore non può essere modificato dopo la creazione del pool di nodi. Per ulteriori informazioni, consulta VM spot.

  • --flex-start: imposta il pool di nodi in modo che utilizzi le VM con avvio flessibile. Le VM con avvio flessibile vengono create utilizzando l'opzione di consumo flex-start, supportata in GKE versione 1.33.0-gke.1712000 o successive.

  • --enable-autoscaling: crea un node pool con la scalabilità automatica abilitata. Quando GKE scala un node pool di sezioni TPU multi-host, lo scala in modo atomico da zero alle dimensioni massime.

    • MAX_NODES: la dimensione massima del node pool. 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.

  • --node-label=cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME, cloud.google.com/gke-workload-type=HIGH_AVAILABILITY: indica a GKE che il pool di nodi della sezione TPU multi-host è una raccolta. Utilizza questo flag se si verificano le seguenti condizioni:

    • Il pool di nodi esegue i carichi di lavoro di inferenza.
    • Il pool di nodi utilizza TPU Trillium.
    • Le VM spot non supportano la pianificazione della raccolta.

    Per maggiori informazioni sulla gestione della pianificazione della raccolta, vedi Gestire la pianificazione della raccolta nelle sezioni TPU multi-host.

  • --placement-type=COMPACT: crea un node pool con il posizionamento compatto abilitato. Questa opzione deve essere utilizzata con il flag --tpu-topology. Per saperne di più, consulta Crea una policy di posizionamento compatto e Topologia TPU.

Terraform

  1. Assicurati di utilizzare la versione 4.84.0 o successive del provider google.

  2. Aggiungi il seguente blocco alla 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]
    }
    flex_start = false
    spot = true
  }

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

    Sostituisci quanto segue:

    • NODE_POOL_RESOURCE_NAME: il nome della risorsa del node pool nel modello Terraform.
    • PROJECT_ID: il tuo ID progetto.
    • CLUSTER_NAME: il nome del cluster esistente a cui aggiungere il pool di nodi.
    • POOL_NAME: il nome del node pool da creare.
    • CLUSTER_LOCATION: la posizione di calcolo per il cluster. Consigliamo di utilizzare un cluster regionale per una maggiore affidabilità del control plane Kubernetes. Puoi anche utilizzare un cluster zonale. Per scoprire di più, consulta la sezione Selezionare una versione e una topologia TPU.
    • NODE_ZONES: l'elenco separato da virgole di una o più zone in cui GKE crea il node pool.

      (Facoltativo) Utilizza una zona AI, ad esempio us-central1-ai1a. Le zone AI sono località specializzate ottimizzate per i workload AI/ML all'interno delle regioni Google Cloud .

    • NUM_NODES: il numero di nodi nel node pool. Deve essere zero o il prodotto del numero di chip TPU diviso per quattro, perché nelle sezioni TPU multihost ogni nodo della sezione TPU ha 4 chip. Ad esempio, se TPU_TOPOLOGY è 4x8, allora ci sono 32 chip, il che significa che NUM_NODES deve essere 8. Per scoprire di più sulle topologie TPU, utilizza la tabella in Scegliere la versione della TPU.
    • TPU_TOPOLOGY: indica la topologia fisica della sezione TPU. Il formato della topologia dipende dalla versione della TPU che stai utilizzando. Per saperne di più sulle topologie TPU, utilizza la tabella in Scegliere una topologia.

    Se vuoi, puoi utilizzare anche le seguenti variabili:

    • RESERVATION_NAME: se utilizzi la prenotazione TPU, questo è l'elenco delle etichette delle risorse di prenotazione da utilizzare durante la creazione del node pool. Per saperne di più su come compilare RESERVATION_LABEL_VALUES nel campo reservation_affinity, consulta Provider Terraform.
    • autoscaling: crea un node pool con la scalabilità automatica abilitata. Quando GKE scala un node pool di sezioni TPU multi-host, lo scala in modo atomico da zero alle dimensioni massime.
      • MAX_NODES: è la dimensione massima del node pool. 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 dello slice TPU. Questo valore non può essere modificato dopo la creazione del pool di nodi. Per saperne di più, consulta VM spot.
    • flex_start: imposta il node pool in modo che utilizzi l'opzione di consumo avvio flessibile. Non può essere impostato su true se spot è abilitato.

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 il nome della zona in base alla versione della TPU che vuoi utilizzare. Per identificare una località disponibile, consulta Disponibilità di TPU in GKE.

  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:

    • CT3: TPU v3, dispositivo host singolo
    • CT3P: TPU v3, sezione di pod multi-host
    • CT4P: TPU v4
    • CT5LP: TPU v5e
    • CT5P: TPU v5p
    • CT6E: TPU Trillium (v6e)

  9. Nel menu a discesa Tipo di macchina, seleziona il nome della macchina da utilizzare per i nodi. Utilizza la tabella Scegli la versione della TPU per scoprire come definire il tipo di macchina e la topologia TPU che creano un node pool della sezione 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 Tipo di disco di avvio sia Standard Persistent Disk o SSD Persistent Disk.

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

  14. Fai clic su Crea.

Verifica lo stato del node pool

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

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

    Sostituisci quanto segue:

    • CLUSTER_NAME: il nome del cluster.
    • PROJECT_ID: il tuo ID progetto.
    • CONTROL_PLANE_LOCATION: la posizione di Compute Engine del control plane del tuo cluster. Fornisci una regione per i cluster regionali o una zona per i cluster zonali.

  2. Utilizza kubectl in Cloud Shell per visualizzare i nodi della slice TPU:

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

    Sostituisci quanto segue:

    • ACCELERATOR_TYPE: il tipo di acceleratore del pool di nodi della slice TPU. Questo valore deve corrispondere al tipo di macchina utilizzato durante la creazione del pool di nodi. Ad esempio, se hai creato un pool di nodi con il tipo di macchina ct5lp-hightpu-4t, utilizza tpu-v5-lite-podslice. Per scoprire quale tipo di acceleratore utilizzare, consulta la tabella in Pianificare la TPU su GKE.
    • TPU_TOPOLOGY: la 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 multislice

In questa sezione esegui un carico di lavoro JAX che mostra il numero globale di chip TPU nella sezione TPU e poi esce.

Per eseguire un workload 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: 2 # This value defines a "Multislice" setup, by creating 2 independent TPU slices.
      template:
        spec:
          parallelism: 4
          completions: 4
          backoffLimit: 0
          template:
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
                cloud.google.com/gke-tpu-topology: 4x4
              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: 4 # This container requires 4 TPU chips

    In questo manifest, seleziona tpu-v5-lite-podslice con una topologia 4x4. Questa configurazione significa quanto segue:

    • La topologia ha 16 chip in totale (4*4=16).
    • tpu-v5-lite-podslice viene implementato in macchine ct5lp-hightpu-4t che hanno quattro chip per VM. Ciò significa che per avere una topologia 4x4, hai bisogno di quattro VM (16/4=4). Pertanto, i campi parallelism e completions devono essere impostati su 4.
    • Ogni container nel pod richiede quattro chip TPU, che corrispondono al numero di chip per VM. Ciò è specificato nel manifest con la riga google.com/tpu: 4.
    • Il campo replicas è impostato su 2 per creare due sezioni.

    Per ulteriori informazioni sui valori possibili per queste variabili, consulta la tabella in Pianificare la TPU su GKE.

    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

    Sostituisci quanto segue:

    • NUM_SLICES: il numero di node pool di sezioni TPU. In questo caso, NUM_SLICES è uguale a 3.
    • ACCELERATOR_TYPE: il tipo di acceleratore del pool di nodi di slice TPU. Questo valore deve corrispondere al tipo di macchina utilizzato durante la creazione pool di nodi. Ad esempio, se hai creato un pool di nodi con il tipo di macchina ct5lp-hightpu-4t, utilizza tpu-v5-lite-podslice. Per scoprire quale tipo di acceleratore utilizzare, consulta la tabella in Pianificare l'utilizzo di TPU su GKE.
    • 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 node pool. 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 la TPU v4 multi-host, il numero di chip TPU in ogni VM è quattro. Per TPU v5e multi-host, il numero di chip TPU in ogni VM è 1, 4 o 8. Pertanto, se il tuo TPU_TOPOLOGY è 2x4x4 (TPU v4 con quattro chip TPU in ogni VM), il NUM_NODES è 32/4, ovvero 8.
    • NUM_CHIPS: il numero di chip per VM.

  2. Applica il manifest:

    kubectl apply -f tpu-multislice.yaml

  3. Verifica che il workload 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

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

  5. (Facoltativo) Rimuovi il carico di lavoro JAX:

    kubectl delete -f tpu-multislice.yaml

Utilizza XPK per creare risorse GKE

  1. Installa XPK completando i passaggi descritti nel file XPK readme.

  2. Crea un cluster GKE in modalità Standard in una regione con TPU disponibili. Il comando seguente crea il cluster GKE con un pool di nodi che include le risorse TPU specificate:

    xpk cluster create \
    --cluster=$CLUSTER_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE \
    --tpu-type=ACCELERATOR_TYPE \
    --num-slices=NUM_SLICES \
    --reservation=$RESERVATION_NAME

    Sostituisci quanto segue:

    • CLUSTER_NAME: un nome per il cluster.
    • PROJECT_ID: il tuo ID progetto Google Cloud .
    • ZONE: il nome della zona in base alla versione della TPU che vuoi utilizzare. Per identificare una località disponibile, consulta Disponibilità di TPU in GKE. Per utilizzare la capacità riservata, assicurati di utilizzare la zona in cui l'hai prenotata.
    • ACCELERATOR_TYPE: la versione e il tipo di TPU. Ad esempio, utilizza tpu7x-standard-4t per Ironwood (TPU7x).
    • NUM_SLICES: il numero di slice TPU da utilizzare per il cluster. Per Multislice, questo valore deve essere maggiore di 1.
    • RESERVATION_NAME: il nome della prenotazione.

    Le ultime due righe dell'output sono simili alle seguenti:

    See your GKE Cluster here:  https://console.cloud.google.com/kubernetes/clusters/details/<location>/CLUSTER_NAME/details?project=PROJECT_ID

  3. Esegui un carico di lavoro che mostri il numero di dispositivi TPU collegati:

    xpk workload create \
    --cluster=$CLUSTER_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE \
    --tpu-type=ACCELERATOR_TYPE \
    --num-slices=NUM_SLICES \
    --docker-image=us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest \
    --workload=WORKLOAD_NAME \
    --command="python -c 'import jax; print(\"TPU cores:\", jax.device_count())'"

    Sostituisci quanto segue:

    • CLUSTER_NAME: un nome per il cluster.
    • PROJECT_ID: il tuo ID progetto Google Cloud .
    • ZONE: la zona del node pool delle macchine scelte.
    • ACCELERATOR_TYPE: il tipo di TPU scelto. Utilizza il completamento automatico di Bash per determinare le opzioni disponibili.
    • NUM_SLICES: il numero di slice TPU configurate nel cluster. Questo valore deve corrispondere a quello utilizzato durante la creazione del cluster.
    • WORKLOAD_NAME: il nome del tuo workload.

Configura le impostazioni aggiuntive

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

Migliorare le prestazioni di rete con hostNetwork

Per migliorare le prestazioni di rete tra le sezioni TPU, ti consigliamo di attivare hostNetworking. Utilizza hostNetwork: true nella specifica del pod per ignorare tutto lo stack di networking di Kubernetes e consentire ai pod Kubernetes di utilizzare direttamente la rete host per la comunicazione da VM a VM.

Per attivare hostNetworking, aggiungi le seguenti due righe alla specifica 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 job di addestramento con ripresa automatica e devi avere gli stessi nomi per ricaricare gli stessi checkpoint.

Se utilizzi TPU Trillium (v6e) o Ironwood (TPU7x) (anteprima) e i tuoi pod utilizzano hostNetworking, installa il seguente DaemonSet per la messa a punto di /proc/sys/net/ipv4/tcp_rmem sul nodo.

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

Migliorare le prestazioni di rete senza hostNetwork su TPU Trillium o Ironwood (TPU7x)

Se utilizzi TPU Trillium o Ironwood (TPU7x) (anteprima) e i tuoi pod non possono utilizzare hostNetworking, attiva il multi-networking con la modalità netdevice per ottenere le migliori prestazioni di rete. Il supporto NIC in modalità netdevice con multi-rete trasferisce la NIC della VM direttamente al pod, bypassando Kubernetes e GKE Dataplane V2.

Il tipo di macchina ct6e-standard-4t è supportato da due NIC fisiche. Kubernetes richiede una vNIC che non può essere passata ai pod. Pertanto, ogni nodo deve avere tre vNIC per consentire ai pod di accedere direttamente a due vNIC per ottenere le migliori prestazioni di entrambe le NIC fisiche.

Per attivare la modalità netdevice per ct6e-standard-4t, completa i seguenti passaggi:

  1. Crea due VPC aggiuntivi che supportano la modalità netdevice
  2. Crea un cluster GKE con funzionalità multi-rete

  3. Configura due netdevice reti. Ad esempio, puoi utilizzare i seguenti oggetti GKENetworkParamSet e Network (SECOND_VPC e THIRD_VPC sono i VPC creati nel passaggio precedente):

    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. Connettere i pod a tre reti. Ad esempio, puoi utilizzare le seguenti annotazioni nella specifica del pod:

    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. Applica i sysctl di rete all'interno del pod, nel container init o nel container dell'applicazione. Ad esempio, puoi aggiungere il seguente init container alle specifiche del pod:

    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:

Utilizza le interfacce eth1 e eth2 per prestazioni di rete migliori, anziché l'interfaccia eth0. Per farlo, aggiungi export LIBTPU_INIT_ARGS="$LIBTPU_INIT_ARGS --megascale_grpc_interface_prefixes=eth1,eth2,lo" alla specifica del workload.

Abilita il logging

I log emessi dai container in esecuzione sui nodi GKE, inclusi i nodi delle sezioni TPU, sono visibili in Esplora log, se hai abilitato la registrazione dei log di sistema GKE nel tuo cluster.

Puoi visualizzare i log di GKE utilizzando Esplora log con il seguente filtro per visualizzare i log dei container per il tuo 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 lo slice 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 delle TPU di GKE.

Abilitare metriche aggiuntive

Oltre alle metriche TPU generali, esistono altre quattro metriche di runtime TPU specifiche per multislice. Queste metriche sono disponibili in GKE versione 1.29.1-gke.1016000 o successive. Il workload TPU deve utilizzare la versione 0.4.24 di JAX

Di seguito sono riportate le metriche multislice disponibili:

  • Latenze di trasferimento DCN (Data Center Network): distribuzione delle latenze di trasferimento di rete per il traffico multislice.
  • Latenze collettive: distribuzione della latenza collettiva end-to-end per il traffico multislice.
  • Latenze di trasferimento da host a dispositivo: distribuzione della latenza di trasferimento da host a dispositivo per ogni blocco di dati per il traffico multislice.
  • Latenze di trasferimento da dispositivo a host: distribuzione della latenza di trasferimento da dispositivo a host per ogni blocco di dati per il traffico multislice.

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 multislice

La tabella seguente distingue l'organizzazione architetturale di uno slice TPU e di un multislice:

Sezione TPU Multislice
Interconnettività Il carico di lavoro viene eseguito su una singola sezione TPU. Tutti i chip TPU in una sezione sono collegati con ICI. Il workload viene eseguito su più slice TPU. La comunicazione all'interno di una sezione avviene tramite ICI. La comunicazione tra le sezioni avviene tramite DCN.
Node pool supportati Sezione TPU single-host e sezione TPU multi-host Gruppi di sezioni TPU multi-host
Tipo di workload consigliato IndexedJob o JobSet JobSet

Esegui la pulizia delle risorse

Il modo più semplice per eliminare la fatturazione è eliminare il progetto Google Cloud che hai creato per il tutorial. In alternativa, puoi eliminare le singole risorse.

Elimina il progetto

    Delete a Google Cloud project:

    gcloud projects delete PROJECT_ID

Elimina singole risorse

Elimina il cluster GKE:

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

Elimina le risorse create con XPK

Se hai utilizzato XPK per creare il cluster, eliminalo per evitare ulteriori addebiti:

xpk cluster delete \
    --cluster=$CLUSTER_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

Sostituisci quanto segue:

  • CLUSTER_NAME: il nome del cluster.
  • PROJECT_ID: il tuo ID progetto Google Cloud .
  • ZONE: la zona in cui è stato creato il cluster.

Passaggi successivi