Esegui il deployment dei carichi di lavoro TPU in GKE Standard

Questa pagina mostra come richiedere ed eseguire il deployment l'addestramento di modelli basati su intelligenza (AI) e machine learning (ML), ottimizzazione e inferenza dei carichi di lavoro tramite Cloud TPU acceleratori (TPU) nei cluster Google Kubernetes Engine (GKE) Standard.

Prima di configurare ed eseguire il deployment di carichi di lavoro TPU in GKE, Prendi familiarità con i seguenti concetti:

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

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, ottieni la versione più recente eseguendo gcloud components update.

Disponibilità di TPU in GKE

Utilizza GKE per creare e gestire i pool di nodi con TPU. Puoi utilizzare questi acceleratori appositamente progettati per eseguire l'addestramento, il tuning e l'inferenza dei modelli di IA su larga scala.

Consulta un elenco delle versioni di TPU supportate in GKE.

Pianifica la configurazione della TPU

Pianifica la configurazione della TPU in base al modello di machine learning e alla quantità di memoria richiesta. Di seguito sono riportati i passaggi pertinenti per la pianificazione della configurazione della TPU:

  1. Seleziona una versione e una topologia TPU.
  2. Seleziona il tipo di pool di nodi da utilizzare.

Assicurati una quota sufficiente per le VM on demand o spot

Se stai creando un pool di nodi di sezioni TPU con VM on demand o a spot, devi avere una quota TPU sufficiente disponibile nella regione che vuoi utilizzare.

La creazione di un pool di nodi di sezioni TPU che utilizza una prenotazione TPU non richiede alcuna quota TPU.1 Puoi saltare questa sezione in tutta sicurezza per le TPU riservate.

Creazione di un pool di nodi della sezione TPU on demand o spot in GKE richiede una quota per l'API Compute Engine. La quota dell'API Compute Engine (compute.googleapis.com) non è uguale alla quota dell'API Cloud TPU (tpu.googleapis.com), necessaria quando si creano TPU con l'API Cloud TPU.

Per controllare il limite e l'utilizzo corrente della quota dell'API Compute Engine per le TPU, segui questi passaggi:

  1. Vai alla pagina Quote nella console Google Cloud:

    Vai a Quote

  2. Nella casella Filtro di , esegui le seguenti:

    1. Seleziona la proprietà Servizio, inserisci API Compute Engine e premi Invio.

    2. Seleziona la proprietà Tipo e scegli Quota.

    3. Seleziona la proprietà Nome e inserisci il nome della quota in base alla versione TPU e al tipo di macchina. Ad esempio, se prevedi di creare modelli I nodi TPU v5e il cui tipo di macchina inizia con ct5lp-, inserisci TPU v5 Lite PodSlice chips.

      Versione TPU Il tipo di macchina inizia con Nome della quota per le istanze on demand Nome della quota per le istanze Spot2
      TPU v4 ct4p- TPU v4 PodSlice chips Preemptible TPU v4 PodSlice chips
      TPU v5e ct5l- TPU v5 Lite Device chips Preemptible TPU v5 Lite Device chips
      TPU v5e ct5lp- TPU v5 Lite PodSlice chips Preemptible TPU v5 Lite PodSlice chips
      TPU v5p ct5p- TPU v5p chips Preemptible TPU v5p chips

    4. Seleziona la proprietà Dimensioni (ad es. località) e inserisci region: seguito dal nome della regione in cui prevedi di creare TPU in GKE. Ad esempio, inserisci region:us-west4 se prevedi di creare nodi di slice TPU nella zona us-west4-a. La quota TPU è regionale, pertanto tutte le zone all'interno della stessa regione consumano la stessa quota TPU.

Se nessuna quota corrisponde al filtro inserito, significa che al progetto non è stata assegnata alcuna quota specificata per la regione desiderata e devi richiedere un aumento della quota TPU.

  1. Quando crei un pool di nodi della sezione TPU, utilizza --reservation e --reservation-affinity=specific flag a per creare istanza dedicata. Le prenotazioni TPU sono disponibili quando acquisti un impegno. ↩

  2. Quando crei un pool di nodi della sezione TPU, utilizza --spot per creare un'istanza Spot. ↩

Assicurati la disponibilità della prenotazione

La creazione di un pool di nodi della sezione TPU riservato, che utilizza una prenotazione, non per qualsiasi quota di TPU. Tuttavia, la prenotazione deve avere chip TPU disponibili o non utilizzati sufficienti al momento della creazione del pool di nodi.

Per vedere quali prenotazioni esistono all'interno di un progetto, visualizza un elenco prenotazioni.

Per visualizzare il numero di chip TPU disponibili all'interno di una prenotazione TPU, visualizza i dettagli di una prenotazione.

Crea un cluster

Crea un cluster GKE in modalità standard in una regione con TPU disponibili. I nostri suggerimenti di utilizzare cluster a livello di regione, che offrono un'alta disponibilità dal piano di controllo Kubernetes. Puoi utilizzare Google Cloud CLI o la console Google Cloud.

gcloud container clusters create CLUSTER_NAME \
  --location LOCATION \
  --cluster-version VERSION

Sostituisci quanto segue:

  • CLUSTER_NAME: il nome del nuovo cluster.
  • LOCATION: la regione in cui è disponibile la tua capacità TPU.
  • VERSION: la versione GKE, che deve supportare il tipo di macchina che vuoi utilizzare. Tieni presente che la versione GKE predefinita potrebbe non essere disponibile per la TPU di destinazione. Per sapere quali sono le versioni GKE minime disponibili per tipo di macchina TPU, consulta la disponibilità delle TPU in GKE

Crea un pool di nodi

Sezione TPU con host singolo

Puoi creare un pool di nodi di slice TPU a un solo host utilizzando Google Cloud CLI, Terraform o la console Google Cloud.

gcloud

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONES \
    --machine-type=MACHINE_TYPE

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.
    • Per i tipi di macchine TPU v5e che iniziano con ct5l-, utilizza us-central1-a o europe-west4-b.
    • 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 TPU v5p, 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 i nodi. Per ulteriori informazioni sui tipi di macchine compatibili con le TPU, utilizza la tabella in Mappatura della configurazione TPU.

Se vuoi, puoi anche utilizzare i seguenti flag:

  • --num-nodes=NUM_NODES: il numero iniziale di nodi nel pool di nodi in ogni zona. Se ometti questo flag, il valore predefinito è 3. Se la scalabilità automatica è abilitata per il pool di nodi che utilizza il --enable-autoscaling, ti consigliamo di impostare NUM_NODES su 0, poiché il gestore della scalabilità automatica provisiona nodi aggiuntivi non appena i carichi di lavoro li richiedono.
  • --reservation=RESERVATION_NAME: il nome del che GKE utilizza durante la creazione del pool di nodi. Se ometti questo flag, GKE utilizza le TPU disponibili. Per scoprire di più sulle prenotazioni TPU, consulta Prenotazione TPU.
  • --enable-autoscaling: crea un pool di nodi con la scalabilità automatica abilitata. Richiede i seguenti flag aggiuntivi:
    • --total-min-nodes=TOTAL_MIN_NODES: valore minimo di tutti i nodi nel pool di nodi.
    • --total-max-nodes=TOTAL_MAX_NODES: massimo di tutti i nodi nel pool di nodi.
    • --location-policy=ANY: assegna la priorità all'utilizzo delle prenotazioni inutilizzate e ridurre il rischio di prerilascio delle VM spot.
  • --spot: imposta il pool di nodi in cui utilizzare VM spot per i nodi nel pool di nodi. Questo valore non può essere modificato dopo la creazione del pool di nodi.

Per un elenco completo di tutti i flag che puoi specificare, consulta la documentazione di riferimento di gcloud container clusters create.

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 seguente blocco alla configurazione di 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]

  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
  }
}

Sostituisci quanto segue:

  • NODE_POOL_RESOURCE_NAME: il nome del pool di nodi risorsa nel modello Terraform.
  • PROJECT_ID: l'ID del tuo progetto.
  • CLUSTER_NAME: il nome del cluster esistente.
  • POOL_NAME: il nome del pool di nodi da creare.
  • CLUSTER_LOCATION: le zona di computing in un cluster Kubernetes. Specifica la regione in cui è disponibile la versione TPU. Per apprendere consulta Selezionare 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.
  • MACHINE_TYPE: il tipo di macchina TPU da utilizzare. Per visualizzare i tipi di macchine compatibili con le TPU, utilizza la tabella in Mappatura della configurazione TPU.

Facoltativamente, puoi anche utilizzare le seguenti variabili:

  • autoscaling: crea un pool di nodi con la scalabilità automatica abilitata. Per sezione TPU con host singolo, GKE scala tra Valori di TOTAL_MIN_NODES e TOTAL_MAX_NODES.
    • TOTAL_MIN_NODES: numero minimo di tutti i nodi nel pool di nodi. Questo campo è facoltativo, a meno che non sia impostata anche la scalabilità automatica specificato.
    • TOTAL_MAX_NODES: numero massimo di tutti i nodi nel pool di nodi. Questo campo è facoltativo, a meno che non sia impostata anche la scalabilità automatica specificato.
  • RESERVATION_NAME: se utilizzi Prenotazione TPU, 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 il valore RESERVATION_LABEL_VALUES nel campo reservation_affinity, consulta Provider Terraform.
  • spot: imposta il pool di nodi in modo che utilizzi le VM spot per la TPU nodi. Questo valore non può essere modificato dopo la creazione del pool di nodi. Per ulteriori informazioni, consulta 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.
    • Per i tipi di macchine TPU v5e che iniziano con ct5l-, utilizza us-central1-a o europe-west4-b.
    • 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 TPU v5p, 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: TPU v4
    • CT5LP: TPU v5e
    • CT5P: TPU v5p

  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 della sezione TPU a host singolo.

  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 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 nel pool di nodi.

  14. Fai clic su Crea.

Sezione TPU multi-host

Puoi creare un pool di nodi di slice TPU multi-host utilizzando Google Cloud CLI, Terraform o 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, utilizza 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-, utilizza us-east1-d, us-east5-a o us-east5-c.

    Per saperne di più, consulta la sezione 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 i nodi. Per scoprire di più sui tipi di macchine disponibili, consulta la sezione Mappatura della configurazione 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 e TPU v5e multi-host, il numero di chip in ogni VM è pari a 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 della prenotazione utilizzata da GKE per creare il pool di nodi. Se ometti questo flag, GKE utilizza i pool di nodi delle sezioni TPU disponibili. Per scoprire di più sulle prenotazioni TPU, consulta la sezione Prenotazione TPU.
  • --spot: imposta il pool di nodi in modo da utilizzare VM spot per i nodi dei sezioni TPU. Questo valore non può essere modificato dopo la creazione del pool di nodi. Per ulteriori informazioni, consulta VM spot.
  • --enable-autoscaling: crea un pool di nodi con scalabilità automatica abilitata. Quando GKE esegue il ridimensionamento di un pool di nodi con sezione TPU multi-host, lo esegue atomicamente da zero alle dimensioni massime.
    • MAX_NODES: la dimensione massima del pool di nodi. 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 successive del fornitore google.

  2. Aggiungi il seguente blocco alla configurazione di 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 della risorsa pool di nodi 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: posizione di calcolo per il cluster. Consigliamo di avere 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 pool di nodi.
    • NUM_NODES: il numero di nodi nel nodo piscina. Deve essere pari a zero o al prodotto del numero di chip TPU diviso per quattro, perché nelle sezioni TPU multi-host ogni nodo della sezione 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 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.

    Se vuoi, puoi anche utilizzare le seguenti variabili:

    • RESERVATION_NAME: se utilizzi la prenotazione TPU, questo è l'elenco delle etichette delle risorse di prenotazione da utilizzare per creare il pool di nodi. Per scoprire di più su come compilareRESERVATION_LABEL_VALUES nel camporeservation_affinity, consultaProvider Terraform.
    • autoscaling: crea un pool di nodi con la 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 pool di nodi. 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 VM spot per i nodi del segmento 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. Se vuoi, 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.

Stato provisioning

Se GKE non può creare il pool di nodi della sezione TPU a causa di Capacità TPU disponibile insufficiente. GKE restituisce un messaggio di errore. che indica che non è possibile creare i nodi della sezione TPU per mancanza di capacità.

Se stai creando un pool di nodi di una sezione TPU con un solo host, il messaggio di errore sarà simile al seguente:

2 nodes cannot be created due to lack of capacity. The missing nodes will be
created asynchronously once capacity is available. You can either wait for the
nodes to be up, or delete the node pool and try re-creating it again later.

Se stai creando un pool di nodi di sezioni TPU multi-host, il messaggio di errore sarà simile al seguente:

The nodes (managed by ...) cannot be created now due to lack of capacity. They
will be created asynchronously once capacity is available. You can either wait
for the nodes to be up, or delete the node pool and try re-creating it again
later.

La richiesta di provisioning di TPU potrebbe rimanere in coda per molto tempo e rimanere nello stato "Provisioning" mentre è in coda.

Una volta che la capacità è disponibile, GKE crea i nodi rimanenti che non sono stati creati.

Se hai bisogno di capacità prima, ti consigliamo di provare le VM spot, anche se tieni presente che le VM spot consumano una quota diversa rispetto alle istanze on demand.

Puoi eliminare la richiesta TPU in coda eliminando il nodo della sezione TPU pool.

Esegui il carico di lavoro sui nodi delle sezioni TPU

Preparazione del carico di lavoro

I carichi di lavoro TPU hanno i seguenti requisiti di preparazione.

  1. Framework come JAX, PyTorch e TensorFlow accedono alle VM TPU utilizzando la libtpu libreria condivisa. libtpu include il compilatore XLA, il software di runtime TPU e il driver TPU. Ogni release di PyTorch e JAX richiede una determinata versione di libtpu.so. A per utilizzare le TPU in GKE, assicurati di utilizzare le seguenti versioni:
    Tipo di TPU libtpu.so versione
    TPU v5e
    tpu-v5-lite-podslice
    tpu-v5-lite-device
    TPU v5p
    tpu-v5p-slice
    • Versione jax[tpu] consigliata: 0.4.19 o successive.
    • Versione torchxla[tpuvm] consigliata: l'utilizzo di una versione consigliata per la notte il 23 ottobre 2023.
    TPU v4
    tpu-v4-podslice
  2. Imposta le seguenti variabili di ambiente per il contenitore che richiede le risorse TPU:
    • TPU_WORKER_ID: un numero intero univoco per ogni pod. Questo ID indica un ID worker univoco nel segmento TPU. I valori supportati per questo campo vanno da zero a di pod meno uno.
    • TPU_WORKER_HOSTNAMES: un elenco separato da virgole di nomi host o IP delle VM TPU che devono comunicare tra loro all'interno della sezione. Deve essere presente un nome host o un indirizzo IP per ogni VM TPU nel slice. L'elenco di indirizzi IP o nomi host è ordinato e ha indice zero per TPU_WORKER_ID.

      • GKE inserisce automaticamente queste variabili di ambiente utilizzando un modello webhook quando viene creato un job con completionMode: Indexed, subdomain, parallelism > 1 e richiesta google.com/tpu proprietà. GKE aggiunge un servizio headless vengono aggiunti i record DNS Pod che supportano il servizio.

      Quando esegui il deployment di risorse multi-host TPU con Kuberay, GKE fornisce un webhook disponibile come parte dei modelli Terraform per l'esecuzione di Ray su GKE. Istruzioni per eseguire Ray su GKE con TPU disponibile nel Guida dell'utente di TPU. Il webhook con mutazioni inietta queste variabili di ambiente nei cluster Ray che richiedono proprietà google.com/tpu e un selettore di nodi cloud.google.com/gke-tpu-topology multi-host.

    • Nel manifest del carico di lavoro, aggiungi i selettori dei nodi Kubernetes per assicurarti che GKE pianifichi il carico di lavoro TPU sul tipo di macchina TPU. Topologia TPU che hai definito:

        nodeSelector:
    cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR
    cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY

      Sostituisci quanto segue:

      • TPU_ACCELERATOR: il nome dell'acceleratore TPU:
        • Per TPU v4, utilizza tpu-v4-podslice
        • Per i tipi di macchine TPU v5e che iniziano con ct5l-, usa tpu-v5-lite-device
        • Per i tipi di macchine TPU v5e che iniziano con ct5lp-, usa tpu-v5-lite-podslice
        • Per TPU v5p, utilizza tpu-v5p-slice
      • TPU_TOPOLOGY: la topologia fisica per la sezione TPU. Il formato della topologia dipende dalla TPU dell'oggetto come segue:
        • TPU v4: definisci la topologia in tuple di 3 elementi ({A}x{B}x{C}), ad esempio 4x4x4.
        • TPU v5e: definisci la topologia in tuple di 2 elementi ({A}x{B}), ad esempio 2x2.
        • TPU v5p: definisci la topologia in tuple di 3 elementi ({A}x{B}x{C}), ad esempio 4x4x4.

Dopo aver completato la preparazione del carico di lavoro, puoi eseguire un job che utilizza le TPU.

Le sezioni seguenti mostrano esempi di come eseguire un job che esegue calcoli semplici con le TPU.

Esempio 1: esegui un carico di lavoro che mostra il numero di chip TPU disponibili in un pool di nodi di una sezione TPU

Il seguente carico di lavoro restituisce il numero di chip TPU su tutti i nodi di uno slice TPU multi-host. Per creare uno slice multi-host, il carico di lavoro ha i seguenti parametri:

  • Versione TPU: TPU v4
  • Topologia: 2x2x4

Questa selezione di versione e topologia genera un segmento multi-host.

  1. Salva il seguente manifest come available-chips-multihost.yaml: 
    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice
        cloud.google.com/gke-tpu-topology: 2x2x4
      containers:
      - name: tpu-job
        image: python:3.10
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        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("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: 4
  2. Esegui il deployment del manifest: 
    kubectl create -f available-chips-multihost.yaml

    GKE esegue una sezione TPU v4 con quattro VM (sezione TPU multi-host). La sezione ha 16 chip TPU interconnessi.

  3. Verifica che il job abbia creato quattro pod: 
    kubectl get pods

    L'output è simile al seguente:

    NAME                       READY   STATUS      RESTARTS   AGE
tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s
  4. Recupera i log di uno dei pod: 
    kubectl logs POD_NAME

    Sostituisci POD_NAME con il nome di uno degli utenti creati di pod. Ad esempio, tpu-job-podslice-0-5cd8r.

    L'output è simile al seguente: 

    TPU cores: 16

Esempio 2: eseguire un carico di lavoro che visualizza il numero di chip TPU disponibili nella sezione TPU

Il seguente carico di lavoro è un pod statico che mostra il numero di chip TPU collegati a un nodo specifico. Per creare un nodo host singolo, il carico di lavoro prevede i seguenti parametri:

  • Versione TPU: TPU v5e
  • Topologia: 2 x 4

Questa versione e selezione della topologia generano una sezione con host singolo.

  1. Salva il seguente manifest come available-chips-singlehost.yaml: 
    apiVersion: v1
kind: Pod
metadata:
  name: tpu-job-jax-v5
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice
    cloud.google.com/gke-tpu-topology: 2x4
  containers:
  - name: tpu-job
    image: python:3.10
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    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("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8
      limits:
        google.com/tpu: 8
  2. Esegui il deployment del manifest: 
    kubectl create -f available-chips-singlehost.yaml

    GKE esegue il provisioning di nodi con otto sezioni TPU a host singolo che utilizzano TPU v5e. Ogni nodo TPU ha otto chip TPU (sezione TPU con un solo host).

  3. Recupera i log del pod: 
    kubectl logs tpu-job-jax-v5

    L'output è simile al seguente: 

    Total TPU chips: 8

Esegui l'upgrade dei pool di nodi utilizzando gli acceleratori (GPU e TPU)

GKE upgrade automatico Cluster standard, inclusi i pool di nodi. Puoi anche manualmente esegui l'upgrade del nodo se vuoi che i tuoi nodi usino una versione successiva prima. Per controllare il funzionamento degli upgrade per il tuo cluster, utilizza i canali di rilascio, i periodi di manutenzione ed esclusioni e la sequenziazione del rollout.

Puoi anche configurare un upgrade dei nodi per del tuo pool di nodi, come surge upgrade o blu-verde upgrade. Configurando queste strategie, puoi assicurarti che venga eseguito l'upgrade dei pool di nodi in modo da raggiungere l'equilibrio ottimale tra velocità e interruzione completamente gestito di Google Cloud. Per i pool di nodi delle sezioni TPU multi-host, anziché utilizzare la strategia di upgrade dei nodi configurata, GKE ricrea l'intero pool di nodi in un unico passaggio. Per scoprire di più, consulta la definizione di atomicità in Terminologia relativa alle TPU in GKE.

L'utilizzo di una strategia di upgrade dei nodi richiede temporaneamente a GKE di eseguire il provisioning di risorse aggiuntive, a seconda della configurazione. Se Google Cloud ha una capacità limitata per le risorse del tuo pool di nodi, ad esempio visualizzi errori di disponibilità delle risorse quando cerchi di creare altri nodi con GPU o TPU, consulta Eseguire l'upgrade in un ambiente con risorse limitate.

Esegui la pulizia

Per evitare che al tuo account Google Cloud vengano addebitati costi relativi alle risorse utilizzate in questa guida, ti consigliamo di eliminare i pool di nodi di slice TPU che non hanno più carichi di lavoro pianificati. Se i carichi di lavoro in esecuzione devono essere terminato, utilizza kubectl drain per pulire i carichi di lavoro prima di eliminare il nodo.

  1. Elimina un pool di nodi della sezione TPU:

    gcloud container node-pools delete POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME

    Sostituisci quanto segue:

    • POOL_NAME: il nome del pool di nodi.
    • CLUSTER_NAME: il nome del cluster.
    • LOCATION: la località di computing del cluster.

Configurazioni aggiuntive

Le seguenti sezioni descrivono le configurazioni aggiuntive che puoi applicare ai carichi di lavoro TPU.

Multislice

Puoi aggregare le sezioni più piccole in una multisezione per gestire carichi di lavoro di addestramento più grandi. Per ulteriori informazioni, consulta TPU multislice in GKE.

Esegui la migrazione della prenotazione TPU

Se hai già prenotazioni TPU, devi prima eseguire la migrazione della prenotazione TPU a un nuovo sistema di prenotazione basato su Compute Engine. Puoi anche creare un sistema di prenotazione basato su Compute Engine in cui non è necessaria alcuna migrazione. Per scoprire come eseguire la migrazione delle prenotazioni TPU, consulta Prenotazione TPU.

Logging

Log emessi dai container in esecuzione su nodi GKE, tra cui TPU Le VM, vengono raccolte dell'agente Logging GKE, inviato a Logging visibile in Logging.

Utilizzare il provisioning automatico dei nodi GKE

Puoi configurare GKE in modo da creare ed eliminare automaticamente il nodo per soddisfare le esigenze in base alle esigenze di risorse dei carichi di lavoro TPU. Per ulteriori informazioni, consulta la pagina sulla configurazione di Cloud TPU.

Riparazione automatica dei nodi della sezione TPU

Se un nodo di sezione TPU in un pool di nodi di sezione TPU multi-host non è in stato di esecuzione, l'intero pool di nodi viene ricreato. Invece, in un pool di nodi di una sezione TPU con un solo host, viene riparato automaticamente solo il nodo TPU non funzionante.

Le condizioni che comportano nodi della sezione TPU non integro includono seguenti:

  • Qualsiasi nodo della sezione TPU con condizioni comuni dei nodi.
  • Qualsiasi nodo della sezione TPU con un numero di TPU non allocabili maggiore di zero.
  • Qualsiasi istanza VM in uno slice TPU che è stata arrestata (a causa della preemption) o terminata.
  • Manutenzione dei nodi: se un nodo della sezione TPU all'interno di un pool di nodi della sezione TPU multi-host si arresta per la manutenzione dell'host, GKE ricrea l'intero pool di nodi della sezione TPU.

Puoi visualizzare lo stato della riparazione (incluso il motivo dell'errore) nei cronologia delle operazioni, Se l'errore è causato da una quota insufficiente, contatta il tuo del rappresentante dell'account Google Cloud per aumentare la quota corrispondente.

Configurare l'interruzione controllata dei nodi di sezioni TPU

Nei cluster GKE con il piano di controllo in esecuzione nella versione 1.29.1-gke.1425000 o successive, i nodi della sezione TPU supportano gli indicatori SIGTERM che avvisano il nodo di un imminente scollegamento. La notifica di arresto imminente è configurabile fino a cinque minuti nei nodi TPU.

Per configurare GKE in modo da terminare i carichi di lavoro in modo controllato entro questo periodo di notifica, segui i passaggi Gestisci l'interruzione dei nodi GKE per GPU e TPU.

Esegui container senza modalità con privilegi

Se il nodo della sezione TPU esegue versioni precedenti alla 1.28, leggi la sezione seguente:

Un contenitore in esecuzione su una VM in una sezione TPU deve accedere a limiti più elevati per la memoria bloccata in modo che il driver possa comunicare con i chip TPU tramite accesso diretto alla memoria (DMA). Per abilitare questa funzionalità, devi configurare una ulimit Se vuoi ridurre l'ambito delle autorizzazioni nel contenitore, completa i seguenti passaggi:

  1. Modifica securityContext in modo da includere i seguenti campi:

    securityContext:
  capabilities:
    add: ["SYS_RESOURCE"]

  2. Aumenta ulimit eseguendo il seguente comando all'interno del contenitore prima di configurare i carichi di lavoro per l'utilizzo delle risorse TPU:

    ulimit -l 68719476736

Nota: per TPU v5e, è disponibile l'esecuzione di container senza modalità con privilegi nei cluster nella versione 1.27.4-gke.900 e successive.

Osservabilità e metriche

Dashboard

Nella pagina Cluster Kubernetes della console Google Cloud, la scheda Osservabilità mostra le metriche di osservabilità di TPU. Per ulteriori informazioni, consulta le metriche di osservabilità di GKE.

La dashboard TPU viene compilata solo se metriche di sistema abilitate nel tuo cluster GKE.

Metriche di runtime

In GKE 1.27.4-gke.900 o versioni successive, i carichi di lavoro TPU che utilizzano JAX versione 0.4.14 o successive e specificano containerPort: 8431 esportano le metriche di utilizzo delle TPU come GKE metriche di sistema. In Cloud Monitoring sono disponibili le seguenti metriche per monitorare le prestazioni di runtime del tuo carico di lavoro TPU:

  • Ciclo di lavoro: percentuale di tempo nell'ultimo periodo di campionamento (60 secondi) durante il quale i TensorCore hanno eseguito attivamente l'elaborazione su un chip TPU. Una percentuale maggiore significa un migliore utilizzo della TPU.
  • Memoria utilizzata: quantità di memoria dell'acceleratore allocata in byte. Campionamento eseguito ogni 60 secondi.
  • Totale memoria: memoria acceleratore totale in byte. Campionamento eseguito ogni 60 secondi.

Queste metriche si trovano nello schema del nodo Kubernetes (k8s_node) e del container Kubernetes (k8s_container).

Container Kubernetes:

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

Nodo Kubernetes:

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

Metriche host

In GKE 1.28.1-gke.1066000 o versioni successive, le VM in uno slice TPU esportano le metriche di utilizzo delle TPU come metriche di sistema di GKE. In Cloud Monitoring sono disponibili le seguenti metriche per monitorare le prestazioni dell'host TPU:

  • Utilizzo TensorCore: la percentuale attuale di TensorCore utilizzata. Il valore di TensorCore è uguale alla somma delle unità di moltiplicazione a matrice (MXU) più l'unità vettoriale. Il valore di utilizzo di TensorCore è la divisione delle operazioni TensorCore che sono state eseguite nell'ultimo periodo di campionamento (60 secondi) per il numero di operazioni TensorCore supportate nello stesso periodo. Un valore più elevato indica un utilizzo migliore.
  • Utilizzo larghezza di banda memoria: percentuale attuale dell'acceleratore per la larghezza di banda della memoria in uso. Viene calcolato dividendo la larghezza di banda della memoria utilizzata in un periodo campione (60 secondi) dalla larghezza di banda massima supportata su nello stesso periodo campione.

Queste metriche si trovano nel nodo Kubernetes (k8s_node) e in Kubernetes del container (k8s_container).

Container Kubernetes:

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

Nodo Kubernetes:

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

Per ulteriori informazioni, consulta le metriche di Kubernetes e le metriche di sistema GKE.

Problemi noti

  • Il gestore della scalabilità automatica del cluster potrebbe calcolare erroneamente la capacità per i nuovi nodi di slice TPU prima che questi nodi segnalino le TPU disponibili. Il gestore della scalabilità automatica dei cluster potrebbe quindi eseguire un ulteriore scale up e, di conseguenza, creare più nodi del necessario. Gruppo il gestore della scalabilità automatica fa lo scale down di nodi aggiuntivi, se non necessari, dopo la normale operazione di fare lo scale down.
  • Il gestore della scalabilità automatica dei cluster annulla il ridimensionamento in aumento dei pool di nodi dei sezioni TPU che rimangono in stato di attesa per più di 10 ore. Il gestore della scalabilità automatica del cluster riprova a fare lo scale up operazioni in seguito. Questo comportamento potrebbe ridurre la disponibilità delle TPU per i clienti che non utilizzano le prenotazioni.
  • I carichi di lavoro non TPU che hanno una tolleranza per l'incompatibilità della TPU possono impedire lo fare lo scale down pool di nodi se vengono ricreati durante lo svuotamento del pool di nodi della sezione TPU.
  • La metrica Utilizzo larghezza di banda memoria non è disponibile per le TPU v5e.

Passaggi successivi