Esegui il deployment dei carichi di lavoro TPU su GKE Autopilot

Questa pagina descrive come accelerare i carichi di lavoro di machine learning (ML) eseguendo il deployment utilizzando gli acceleratori Cloud TPU (TPU) nei cluster Autopilot di Google Kubernetes Engine (GKE). Prima di leggere questa pagina, assicurati di conoscere i seguenti argomenti di base:

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

Come funzionano le TPU in Autopilot

Per utilizzare le TPU nei carichi di lavoro Autopilot, richiedi una versione di TPU per quella versione TPU nel manifest del carico di lavoro. Quindi utilizzi i campi resources.requests di Kubernetes e resources.limits da specificare il numero di chip TPU che il carico di lavoro può utilizzare. Quando esegui il deployment del carico di lavoro, GKE esegue il provisioning dei nodi con la configurazione TPU richiesta e pianifica i pod sui nodi. GKE posiziona ogni carico di lavoro il proprio nodo, in modo che ogni pod possa accedere a tutte le risorse del nodo il rischio di interruzioni.

Le TPU in Autopilot sono compatibili con le seguenti funzionalità:

  1. Spot Pod
  2. Prenotazioni di capacità specifica
  3. Pod con tempo di esecuzione esteso

Pianifica la configurazione della TPU

Prima di richiedere le TPU, decidi la configurazione che preferisci in base ai requisiti di CPU e memoria del tuo carico di lavoro. Devi decidere quanto segue:

  • Versione TPU: la versione specifica di Cloud TPU, ad esempio v5e.
  • Topologia per la versione di TPU selezionata: la disposizione e il numero di TPU.

La versione TPU e la topologia selezionate determinano se GKE provisiona i nodi come sezioni a host singolo o sezioni a più host. Nelle sezioni con un solo host, ogni nodo è indipendente dagli altri nodi della sezione TPU. Nelle sezioni multi-host, GKE crea un gruppo di nodi con VM interconnesse in una sezione TPU. Le sezioni multi-host sono atomiche, il che significa che GKE esegue il ridimensionamento dell'intero gruppo di nodi interconnessi come una singola unità.

Per informazioni sulle versioni TPU disponibili, sulle topologie corrispondenti, sulla capacità della CPU e della memoria e sul tipo di sezione risultante, consulta Scegliere una configurazione TPU Autopilot.

Prezzi

Per informazioni sui prezzi, consulta Prezzi di Autopilot.

Prima di iniziare

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

  • 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.
  • Assicurati di avere un cluster Autopilot che esegue GKE versione 1.29.2-gke.1521000 o successive.
  • Per utilizzare TPU prenotate, assicurati di avere una capacità specifica esistente prenotazione. Per le istruzioni, consulta Consumo di risorse di zona prenotate.

Assicurati di disporre della quota TPU

Le sezioni seguenti ti aiutano a assicurarti di disporre di una quota sufficiente quando utilizzi le TPU in GKE.

Per creare nodi di slice TPU, devi disporre di una quota TPU disponibile, a meno che tu non stia utilizzando una prenotazione della capacità esistente. Se utilizzi TPU riservate, salta questa sezione.

La creazione di nodi sezione TPU in GKE richiede una quota API Compute Engine (compute.googleapis.com), non quota API Cloud TPU (tpu.googleapis.com). La della quota sia diverso nei pod Autopilot regolari Spot pod.

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 e al valore TPU nel selettore dei nodi cloud.google.com/gke-tpu-accelerator. Ad esempio, se prevedi di creare nodi TPU v5e on demand il cui valore nel selettore dei nodi cloud.google.com/gke-tpu-accelerator è tpu-v5-lite-podslice, inserisci TPU v5 Lite PodSlice chips.

      Versione TPU cloud.google.com/gke-tpu-accelerator Nome della quota per le istanze on demand Nome della quota per le istanze Spot2
      TPU v4 tpu-v4-podslice TPU v4 PodSlice chips Preemptible TPU v4 PodSlice chips
      TPU v5e tpu-v5-lite-device TPU v5 Lite Device chips Preemptible TPU v5 Lite Device chips
      TPU v5e tpu-v5-lite-podslice TPU v5 Lite PodSlice chips Preemptible TPU v5 Lite PodSlice chips
      TPU v5p tpu-v5p-slice 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 di cui hai bisogno e devi richiedere un aumento della quota TPU.

Quando viene creata una prenotazione TPU, sia il limite sia i valori di utilizzo attuale l'aumento della quota corrispondente del numero di chip nella TPU prenotazione. Ad esempio, quando viene creata una prenotazione per 16 chip TPU v5e la cui il valore nel selettore di nodi cloud.google.com/gke-tpu-accelerator è tpu-v5-lite-podslice, sia Limite che Utilizzo attuale per la quota TPU v5 Lite PodSlice chips nei pertinenti in termini di regione di 16.

Quote per risorse GKE aggiuntive

Potresti dover aumentare le seguenti quote relative a GKE nelle regioni in cui GKE crea le tue risorse.

  • Quota SSD (GB) di Persistent Disk: il disco di avvio di ogni nodo Kubernetes richiede 100 GB per impostazione predefinita. Pertanto, questa quota deve essere impostata su un valore almeno uguale al prodotto del numero massimo di nodi GKE che prevedi di creare e 100 GB (nodi * 100 GB).
  • Quota di indirizzi IP in uso: ogni nodo Kubernetes utilizza un indirizzo IP. Di conseguenza, questa quota deve essere impostata almeno su un valore pari al numero massimo di i nodi GKE che prevedi di creare.
  • Assicurati che max-pods-per-node sia in linea con l'intervallo della subnet: ogni nodo Kubernetes utilizza intervalli IP secondari per i pod. Ad esempio, max-pods-per-node di 32 richiede 64 indirizzi IP, che si traducono in una subnet /26 per nodo. Tieni presente che questo intervallo non deve essere condiviso con nessun altro cluster. Per evitare di esaurire l'intervallo di indirizzi IP, utilizza il flag --max-pods-per-node per limitare il numero di pod consentiti da pianificare su un nodo. La quota per max-pods-per-node deve essere impostata almeno pari al numero massimo di i nodi GKE che prevedi di creare.

Per richiedere un aumento della quota, vedi Richiedere una quota superiore.

Prepara l'applicazione TPU

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

  1. I framework come JAX, PyTorch e TensorFlow accedono alle VM TPU utilizzando 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. 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 consigliata di jax[tpu]: 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 container 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 al numero 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 inietta automaticamente queste variabili di ambiente utilizzando un webhook con mutazioni quando viene creato un job con le proprietà completionMode: Indexed, subdomain, parallelism > 1 e google.com/tpu di richiesta. GKE aggiunge un servizio headless vengono aggiunti i record DNS Pod che supportano il servizio.

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

Richiedi TPU in un carico di lavoro

Questa sezione mostra come creare un job che richiede TPU Autopilot. In qualsiasi carico di lavoro che richiede TPU, devi specificare seguenti:

  • Selettori di nodi per la versione e la topologia TPU
  • Il numero di chip TPU per un container nel carico di lavoro

Per un elenco di versioni di TPU supportate, topologie e il numero corrispondente di chip e nodi TPU in una sezione, vedi Scegli una configurazione TPU Autopilot.

Considerazioni sulle richieste di TPU nei carichi di lavoro

Le TPU possono essere utilizzate da un solo container in un pod. Il numero di chip TPU richiesti da un contenitore deve essere uguale al numero di chip TPU collegati a un nodo della sezione. Ad esempio, se richiedi TPU v5e (tpu-v5-lite-podslice) con un 2x4 topologia, puoi richiedere una delle seguenti opzioni:

  • 4, che crea due nodi multi-host con quattro chip TPU ciascuno
  • 8 chip, che crea un nodo host singolo con 8 chip TPU

Come best practice per massimizzare l'efficienza dei costi, utilizza sempre tutte le TPU nel segmento richiesto. Se richiedi uno slice multi-host di due nodi con 4 chip TPU ciascuno, devi eseguire il deployment di un carico di lavoro che funzioni su entrambi i nodi e consumi tutti e 8 i chip TPU nello slice.

Crea un carico di lavoro che richiede TPU

I passaggi seguenti creano un job che richiede le TPU. Se hai carichi di lavoro su sezioni TPU multi-host, devi anche creare un servizio headless che seleziona il carico di lavoro per nome. Questo servizio headless consente ai pod su diverse ai nodi nella sezione multi-host di comunicare tra loro aggiornando Configurazione DNS di Kubernetes in modo che punti ai pod nel carico di lavoro.

  1. Salva il seguente manifest come tpu-autopilot.yaml:

    apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: TPU_TYPE
        cloud.google.com/gke-tpu-topology: TOPOLOGY
      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.
        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: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: 500Gi
            google.com/tpu: NUMBER_OF_CHIPS

    Sostituisci quanto segue:

    • TPU_TYPE: il tipo di TPU da utilizzare, ad esempio tpu-v4-podslice. Deve essere un valore supportato da GKE.
    • TOPOLOGY: la disposizione dei chip TPU nel sezione, ad esempio 2x2x4. Deve essere una topologia supportata per il tipo di TPU selezionato.
    • NUMBER_OF_CHIPS: il numero di chip TPU per il container da utilizzare. Deve essere lo stesso valore per limits e requests.

  2. Esegui il deployment del job:

    kubectl create -f tpu-autopilot.yaml

Quando crei questo job, GKE esegue automaticamente le seguenti operazioni:

  1. Esegue il provisioning dei nodi per eseguire i pod. A seconda del tipo di TPU, della topologia e delle richieste di risorse specificate, questi nodi sono sezioni con un solo host o sezioni con più host.
  2. Aggiunge incompatibilità ai pod e tolleranze ai nodi per prevenire qualsiasi l'esecuzione di altri carichi di lavoro sugli stessi nodi dei carichi di lavoro TPU.

Esempio: visualizzare i chip TPU totali in una sezione multi-host

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: 2 x 2 x 4

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.
        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 i pod. Ad esempio, tpu-job-podslice-0-5cd8r.

    L'output è simile al seguente: 

    TPU cores: 16

Esempio: visualizzare i chip TPU in un singolo nodo

Il seguente carico di lavoro è un pod statico che mostra il numero di chip TPU collegati a un nodo specifico. Per creare un nodo a host singolo, il carico di lavoro ha 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.
    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 dei nodi con otto sezioni TPU a host singolo che utilizzano TPU v5e. Ogni nodo TPU ha otto chip TPU (sezione TPU con host singolo).

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

    L'output è simile al seguente: 

    Total TPU chips: 8

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 hai attivato le metriche di sistema 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 servizio: percentuale di tempo nell'ultimo periodo di campionamento (60 secondi) durante che i TensorCore elaboravano attivamente 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.
  • Memoria totale: memoria dell'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 versione 1.28.1-gke.1066000 o successive, VM in una sezione TPU esporta le metriche di utilizzo delle TPU come GKE metriche di sistema. 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 TensorCore è la divisione delle operazioni TensorCore che sono state eseguita nell'ultimo periodo campione (60 secondi) per il numero supportato di le operazioni di TensorCore 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 nello schema del nodo Kubernetes (k8s_node) e del container Kubernetes (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, vedi Metriche Kubernetes e le metriche di sistema di GKE.

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.

Consigli per i carichi di lavoro TPU in Autopilot

I seguenti suggerimenti potrebbero migliorare l'efficienza dei carichi di lavoro TPU:

  • Utilizza pod con tempo di esecuzione esteso per un periodo di tolleranza fino a sette giorni prima che GKE interrompa i pod per riduzioni o upgrade dei nodi. Puoi utilizzare le finestre di manutenzione e le esclusioni con i pod con tempo di esecuzione esteso per ritardare ulteriormente gli upgrade automatici dei nodi.
  • Utilizza le prenotazioni della capacità per assicurarti che i tuoi carichi di lavoro ricevano le TPU richieste senza essere inseriti in una coda per la disponibilità.