Questa pagina descrive come accelerare i workload 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 concetti:
Come funzionano le TPU in Autopilot
Per utilizzare le TPU nei carichi di lavoro Autopilot, richiedi una versione TPU e una topologia supportata per quella versione TPU nel manifest del carico di lavoro. Poi utilizzi
i campi resources.requests
e resources.limits
di Kubernetes per specificare
il numero di chip TPU da utilizzare per il carico di lavoro. 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 sul suo nodo in modo che ogni pod possa accedere a tutte le risorse del nodo con un rischio minimo di interruzione.
Le TPU in Autopilot sono compatibili con le seguenti funzionalità:
Pianifica la configurazione della TPU
Pianifica la configurazione della TPU in base al modello e alla quantità di memoria richiesta. Prima di utilizzare questa guida per eseguire il deployment dei tuoi carichi di lavoro su TPU, completa i passaggi di pianificazione descritti in Pianificare la configurazione delle TPU.
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. Attiva l'API Google Kubernetes Engine
- Se vuoi utilizzare Google Cloud CLI per questa attività,
installa e poi
inizializza gcloud CLI. Se hai già installato gcloud CLI, ottieni 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 le TPU riservate, assicurati di avere una prenotazione della capacità specifica esistente. Per le istruzioni, consulta Utilizzare le risorse di zona prenotate.
Assicurati di disporre della quota TPU
Le sezioni seguenti ti aiutano ad assicurarti di avere 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 di slice TPU in GKE richiede la quota dell'API Compute Engine (compute.googleapis.com), non la quota dell'API Cloud TPU (tpu.googleapis.com). Il nome della quota è diverso nei pod Autopilot standard e nei pod Spot.
Per controllare il limite e l'utilizzo corrente della quota dell'API Compute Engine per le TPU, segui questi passaggi:
Vai alla pagina Quote nella console Google Cloud:
Nella casella
Filtro, procedi nel seguente modo:Seleziona la proprietà Servizio, inserisci API Compute Engine e premi Invio.
Seleziona la proprietà Tipo e scegli Quota.
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 nodicloud.google.com/gke-tpu-accelerator
ètpu-v5-lite-podslice
, inserisciTPU 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 v3 tpu-v3-device
TPU v3 Device chips
Preemptible TPU v3 Device chips
TPU v3 tpu-v3-slice
TPU v3 PodSlice chips
Preemptible TPU v3 PodSlice chips
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
TPU v6e (anteprima) tpu-v6e-slice
TPU v6e Slice chips
Preemptible TPU v6e Lite PodSlice chips
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, inserisciregion:us-west4
se prevedi di creare nodi di slice TPU nella zonaus-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 corrente per la quota corrispondente aumentano in base al numero di chip nella prenotazione TPU. Ad esempio, quando viene creata una prenotazione per 16 chip TPU v5e il cui valore nel selettore dei nodi cloud.google.com/gke-tpu-accelerator
è tpu-v5-lite-podslice
, sia il limite sia l'utilizzo corrente per la quota TPU v5 Lite PodSlice chips
nella regione pertinente aumentano 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 indirizzi IP in uso: ogni nodo Kubernetes utilizza un indirizzo IP. Pertanto, questa quota deve essere impostata almeno sul numero massimo di 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 permax-pods-per-node
deve essere impostata su un valore almeno pari al numero massimo di 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.
- 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 dilibtpu.so
. Per utilizzare le TPU in GKE, assicurati di utilizzare le seguenti versioni:Tipo di TPU libtpu.so
versioneTPU v6e
tpu-v6e-slice
- Versione jax[tpu] consigliata: v0.4.9 o successive
- Versione consigliata di torchxla[tpuvm]: v2.1.0 o successive
TPU v5e
tpu-v5-lite-podslice
tpu-v5-lite-device
- Versione jax[tpu] consigliata: v0.4.9 o successive
- Versione consigliata di torchxla[tpuvm]: v2.1.0 o successive
TPU v5p
tpu-v5p-slice
- Versione consigliata di jax[tpu]: 0.4.19 o successive.
- Versione torchxla[tpuvm] consigliata: è consigliabile utilizzare una compilazione della versione nightly del 23 ottobre 2023.
TPU v4
tpu-v4-podslice
- Jax[TPU] consigliato: v0.4.4 o successive
- torchxla[tpuvm] consigliato: v2.0.0 o successive
TPU v3
tpu-v3-slice
tpu-v3-device
- Jax[TPU] consigliato: v0.4.4 o successive
- torchxla[tpuvm] consigliato: v2.0.0 o successive
- 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 al numero di pod meno uno.TPU_WORKER_HOSTNAMES
: un elenco separato da virgole di nomi host o indirizzi IP delle VM TPU che devono comunicare tra loro all'interno del segmento. 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 perTPU_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
egoogle.com/tpu
di richiesta. GKE aggiunge un servizio senza interfaccia in modo che i record DNS vengano aggiunti per i pod che supportano il servizio.
Una volta completata la preparazione del carico di lavoro, puoi eseguire un job che utilizza le TPU.
Richiedere TPU in un carico di lavoro
Questa sezione mostra come creare un job che richiede TPU in Autopilot. In qualsiasi carico di lavoro che richiede TPU, devi specificare quanto segue:
- Selettori di nodi per la versione e la topologia TPU
- Il numero di chip TPU per un contenitore nel tuo carico di lavoro
Per un elenco delle versioni e delle topologie TPU supportate e del numero corrispondente di chip e nodi TPU in una sezione, consulta Scegliere una configurazione TPU di Autopilot.
Considerazioni per le richieste TPU nei carichi di lavoro
Solo un container in un pod può utilizzare le TPU. 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 una topologia 2x4
, puoi richiedere una delle seguenti opzioni:
4
, che crea due nodi multi-host con 4 chip TPU ciascuno- Chip
8
, 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 workload 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 TPU. Se hai carichi di lavoro che vengono eseguiti su sezioni TPU multi-host, devi anche creare un servizio headless che li selezioni per nome. Questo servizio headless consente ai pod su diversi nodi dello slice multi-host di comunicare tra loro aggiornando la configurazione DNS di Kubernetes in modo che indichi i pod nel carico di lavoro.
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 esempiotpu-v4-podslice
. Deve essere un valore supportato da GKE.TOPOLOGY
: la disposizione dei chip TPU nel sezione, ad esempio2x2x4
. Deve essere una topologia supportata per il tipo di TPU selezionato.NUMBER_OF_CHIPS
: il numero di chip TPU da utilizzare per il contenitore. Deve essere lo stesso valore perlimits
erequests
.
Esegui il deployment del job:
kubectl create -f tpu-autopilot.yaml
Quando crei questo job, GKE esegue automaticamente le seguenti operazioni:
- Esegue il provisioning dei nodi per l'esecuzione dei 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.
- Aggiunge incompatibilità ai pod e tolleranze ai nodi per impedire l'esecuzione di altri carichi di lavoro sugli stessi nodi dei carichi di lavoro TPU.
Esempio: visualizzare il numero totale di chip TPU in una sezione con più 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: 2x2x4
Questa selezione di versione e topologia genera uno slice multi-host.
- 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
- 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.
- 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
- Recupera i log di uno dei pod:
kubectl logs POD_NAME
Sostituisci
POD_NAME
con il nome di uno dei pod creati. Ad esempio,tpu-job-podslice-0-5cd8r
.L'output è simile al seguente:
TPU cores: 16
Esempio: visualizza 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: 2x4
Questa versione e la selezione della topologia generano uno slice con un solo host.
- 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
- 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 a host singolo).
- 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 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 più elevata indica 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 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 corrente della larghezza di banda della memoria dell'acceleratore in uso. Calcolata dividendo la larghezza di banda della memoria impiegata in un periodo di campionamento (60 secondi) per la larghezza di banda massima supportata nello stesso periodo di campionamento.
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, consulta le metriche di Kubernetes e le metriche di sistema GKE.
Logging
I log emessi dai container in esecuzione sui nodi GKE, incluse le VM TPU, vengono raccolti dall'agente di logging GKE, inviati a Logging e sono visibili in Logging.
Consigli per i carichi di lavoro TPU in Autopilot
I seguenti consigli 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à.