Questa pagina mostra come eseguire il deployment dei carichi di lavoro in Google Kubernetes Engine (GKE) utilizzando la configurazione multisezione di Cloud TPU per un addestramento conveniente e su larga scala.
Prima di configurare più sezioni in GKE, devi conoscere i seguenti concetti:
Che cos'è TPU Multisezione
TPU Multislice è l'organizzazione architetturale delle VM in una sezione TPU in cui due o più sezioni Cloud TPU comunicano sulla rete di data center (DCN). Multislice consente l'addestramento full stack ed economico su larga scala con scalabilità quasi lineare fino a decine di migliaia di chip TPU. In una configurazione con più sezioni, GKE esegue il deployment di un carico di lavoro multisettore su più sezioni TPU. La comunicazione tra i chip TPU all'interno di una sezione avviene tramite le interconnessioni tra chip (ICI). La comunicazione tra le sezioni avviene attraverso la rete DCN.
Ti consigliamo di utilizzare l'opzione Multisezione se il job è troppo grande per poter essere inserito in una singola sezione TPU.
Disponibilità di più sezioni 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 di JAX supportata è la 2.1.
- Più sezioni supporta solo i pool di nodi delle sezioni TPU multi-host. Ad esempio, non puoi utilizzare
Multislice con un
ct4p-hightpu-4t
con una topologia2x2x1
o unct5lp-hightpu-4t
con una topologia2x2
, perché si tratta di pool di nodi di una sezione TPU a host singolo. - Multislice supporta solo l'addestramento sincrono multicontroller.
- I carichi di lavoro con più sezioni possono essere eseguiti solo su sezioni TPU che condividono lo stesso tipo, dimensione e topologia di TPU.
Prima di iniziare
Prima di iniziare, assicurati di aver eseguito le seguenti attività:
- Abilita l'API Google Kubernetes Engine. Abilita l'API Google Kubernetes Engine
- Se vuoi utilizzare Google Cloud CLI per questa attività, installa e quindi initialize gcloud CLI. Se hai già installato gcloud CLI, ottieni la versione più recente eseguendo
gcloud components update
.
- Crea un cluster Standard o Autopilot che esegue una versione che supporta Multislice. Per le versioni supportate, consulta Disponibilità di più sezioni in GKE.
- Assicurati che il progetto abbia una quota sufficiente per Cloud TPU in GKE.
- Installa JobSet v0.2.3 o versioni successive.
Esegui un carico di lavoro su una sezione multipla
Questa sezione mostra come eseguire un carico di lavoro su una sezione multipla. Se utilizzi la modalità GKE Autopilot, passa alla sezione Esegui un carico di lavoro con più sezioni. I cluster Autopilot che eseguono la versione 1.29.2-gke.1521000 o successive abilitano le TPU per impostazione predefinita.
prepara un pool di nodi in modalità Standard
Questa sezione illustra i seguenti passaggi:
- Crea tre pool di nodi delle sezioni TPU multi-host
- Verifica lo stato del pool di nodi
Crea il pool di nodi della sezione TPU
Puoi creare più di un pool di nodi della sezione TPU multi-host. Ai fini di questa guida, crea tre pool di nodi di sezioni TPU multi-host per eseguire un carico di lavoro Multislice. Puoi creare un pool di nodi con una sezione 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 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-
, utilizzaus-west1-c
,us-west4-a
,us-west4-b
,us-central1-a
,us-east1-c
,us-east5-b
oeurope-west4-a
. - Per i tipi di macchine TPU v5p che iniziano con
ct5p-
, usaus-east1-d
,us-east5-a
ous-east5-c
.
Per scoprire di più, consulta Disponibilità delle TPU in GKE.
- Per TPU v4, usa
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 pagina Mappatura della configurazione di TPU.TPU_TOPOLOGY
: la topologia fisica per la sezione TPU. Il formato della topologia dipende dalla versione di TPU, come segue:- TPU v4 o v5p: definisci la topologia in 3 tuple (
{A}x{B}x{C}
), ad esempio4x4x4
. - TPU v5e: definisci la topologia in due tuple (
{A}x{B}
), ad esempio2x2
.
Per scoprire di più, consulta Topologia.
- TPU v4 o v5p: definisci la topologia in 3 tuple (
NUM_NODES
: il numero di nodi nel pool di nodi. Deve essere zero o il prodotto dei valori definiti inTPU_TOPOLOGY
({A}x{B}x{C}
) diviso per il numero di chip in ogni VM. Per TPU v4 multi-host e TPU v5e, il numero di chip in ogni VM è quattro. Pertanto, se il valore diTPU_TOPOLOGY
è2x4x4
(TPU v4 con quattro chip in ogni VM), ilNUM_NODES
è 32/4, ovvero 8.
Facoltativamente, puoi anche utilizzare i seguenti flag:
RESERVATION_NAME
: il nome della prenotazione utilizzata da GKE durante la creazione del pool di nodi. Se ometti questo flag, GKE utilizza i pool di nodi della sezione TPU disponibili. Per scoprire di più sulle prenotazioni TPU, consulta la pagina relativa alla prenotazione TPU.--spot
: imposta il pool di nodi in modo da utilizzare le VM spot per i nodi della sezione TPU. Questo valore non può essere modificato dopo la creazione del pool di nodi. Per ulteriori informazioni, consulta VM spot.--enable-autoscaling
: crea un pool di nodi con scalabilità automatica abilitata. Quando GKE scala un pool di nodi di sezioni TPU multi-host, fa lo scale up atomicamente del pool di nodi da zero alla dimensione massima.MAX_NODES
: la dimensione massima del pool di nodi. Il flag--max-nodes
è obbligatorio se viene fornito il valore--enable-autoscaling
e deve essere uguale al prodotto dei valori definiti inTPU_TOPOLOGY
({A}x{B}x{C}
) diviso per il numero di chip in ogni VM.
Terraform
- Assicurati di utilizzare la versione 4.84.0 o successive del provider
google
. Aggiungi il blocco seguente alla tua configurazione Terraform:
resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" { provider = google project = PROJECT_ID cluster = CLUSTER_NAME name = POOL_NAME location = CLUSTER_LOCATION node_locations = [NODE_ZONES] initial_node_count = NUM_NODES autoscaling { max_node_count = MAX_NODES location_policy = "ANY" } node_config { machine_type = MACHINE_TYPE reservation_affinity { consume_reservation_type = "SPECIFIC_RESERVATION" key = "compute.googleapis.com/reservation-name" values = [RESERVATION_LABEL_VALUES] } spot = true } placement_policy { type = "COMPACT" tpu_topology = TPU_TOPOLOGY } }
Sostituisci quanto segue:
NODE_POOL_RESOURCE_NAME
: il nome della risorsa del 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
: località di calcolo del cluster. Consigliamo di avere un cluster a livello di regione per una maggiore affidabilità del piano di controllo Kubernetes. Puoi anche usare un cluster a livello di zona. Per scoprire di più, 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.NUM_NODES
: il numero di nodi nel pool. Deve essere zero o il prodotto del numero dei chip TPU diviso per quattro, perché nelle sezioni TPU multi-host ogni nodo della sezione TPU ha 4 chip. Ad esempio, seTPU_TOPOLOGY
è4x8
, ci sono 32 chip, il che significa cheNUM_NODES
deve essere 8. Per ulteriori informazioni sulle topologie TPU, utilizza la tabella in Mappatura della configurazione TPU.TPU_TOPOLOGY
: indica la topologia fisica voluta 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}
), ad esempio4x4x4
. - Per TPU v5e: definisci la topologia in 2 tuple (
{A}x{B}
), ad esempio2x2
.
- Per TPU v4: definisci la topologia in 3 tuple (
Facoltativamente, puoi anche utilizzare le seguenti variabili:
RESERVATION_NAME
: se utilizzi la prenotazione TPU, questo è l'elenco di etichette delle risorse di prenotazione da utilizzare durante la creazione del pool di nodi. Per scoprire di più su come compilareRESERVATION_LABEL_VALUES
nel camporeservation_affinity
, consulta Provider Terraform.autoscaling
: crea un pool di nodi con scalabilità automatica abilitata. Quando GKE scala un pool di nodi di sezioni TPU multi-host, fa lo scale up atomicamente del pool di nodi da zero alla dimensione massima.MAX_NODES
: è la dimensione massima del pool di nodi. Deve essere uguale al prodotto dei valori definiti inTPU_TOPOLOGY
({A}x{B}x{C}
) diviso per il numero di chip in ogni VM.
spot
: consente al pool di nodi di utilizzare le VM spot per i nodi della sezione TPU. Questo valore non può essere modificato dopo la creazione del pool di nodi. Per maggiori informazioni, consulta VM spot.
Console
Per creare un pool di nodi con TPU:
Vai alla pagina Google Kubernetes Engine nella console Google Cloud.
Nell'elenco dei cluster, fai clic sul nome del cluster da modificare.
Fai clic su add_box Aggiungi pool di nodi.
Nella sezione Dettagli del pool di nodi, seleziona la casella Specifica le località dei nodi.
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-
, utilizzaus-west1-c
,us-west4-a
,us-west4-b
,us-central1-a
,us-east1-c
,us-east5-b
oeurope-west4-a
. - Per i tipi di macchine TPU v5p che iniziano con
ct5p-
, usaus-east1-d
,us-east5-a
ous-east5-c
.
- Per TPU v4, usa
Nel riquadro di navigazione, fai clic su Nodi.
Nella sezione Configurazione macchina, seleziona TPU.
Nel menu a discesa Serie, seleziona una delle seguenti opzioni:
- CT4P: per TPU v4.
- CT5LP: per TPU v5e.
Nel menu a discesa Tipo di macchina, seleziona il nome della macchina da utilizzare per i nodi. Utilizza la tabella Mappatura della configurazione TPU per scoprire come definire il tipo di macchina e la topologia TPU che creano un pool di nodi della sezione TPU multi-host.
Nel menu a discesa Topologia TPU, seleziona la topologia fisica per la sezione TPU.
Nella finestra di dialogo Modifiche necessarie, fai clic su Apporta modifiche.
Assicurati che il Tipo di disco di avvio sia Disco permanente standard o Disco permanente SSD.
(Facoltativo) Seleziona la casella di controllo Abilita i nodi sulle VM Spot per utilizzare le VM spot per i nodi nel pool di nodi.
Fai clic su Crea.
Verifica lo stato del pool di nodi
Richiedi le credenziali per poter utilizzare
kubectl
per accedere al cluster:gcloud container clusters get-credentials CLUSTER_NAME \ --project=PROJECT_ID
Sostituisci quanto segue:
CLUSTER_NAME
: il nome del cluster.PROJECT_ID
: l'ID progetto.
Usa
kubectl
, in Cloud Shell, per visualizzare i nodi delle sezioni TPU:kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=TPU_ACCELERATOR \ -l cloud.google.com/gke-tpu-topology=TPU_TOPOLOGY
Sostituisci quanto segue:
TPU_ACCELERATOR
: il tipo di acceleratore TPU utilizzato durante la creazione dei pool di nodi. Ad esempio,tpu-v4-podslice
,tpu-v5-lite-device
otpu-v5-lite-podslice
.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 su più sezioni
In questa sezione eseguirai un carico di lavoro JAX che mostra il numero globale di chip TPU nella sezione TPU, poi esce.
Per eseguire un carico di lavoro JAX:
Crea il seguente manifest
tpu-multislice.yaml
:Autopilot
apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: multislice-job annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: NUM_SLICES template: spec: parallelism: NUM_NODES completions: NUM_NODES backoffLimit: 0 template: spec: nodeSelector: cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 - containerPort: 8431 command: - bash - -c - | pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html python -c 'import jax; print("Global device count:", jax.device_count())' sleep 60 resources: limits: google.com/tpu: NUM_CHIPS
Standard
apiVersion: jobset.x-k8s.io/v1alpha2 kind: JobSet metadata: name: multislice-job annotations: alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool spec: failurePolicy: maxRestarts: 4 replicatedJobs: - name: slice replicas: NUM_SLICES template: spec: parallelism: NUM_NODES completions: NUM_NODES backoffLimit: 0 template: spec: hostNetwork: true dnsPolicy: ClusterFirstWithHostNet nodeSelector: cloud.google.com/gke-tpu-accelerator: ACCELERATOR_TYPE cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY containers: - name: jax-tpu image: python:3.8 ports: - containerPort: 8471 - containerPort: 8080 - containerPort: 8431 securityContext: privileged: true command: - bash - -c - | pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html python -c 'import jax; print("Global device count:", jax.device_count())' sleep 60 resources: limits: google.com/tpu: NUM_CHIPS
Sostituisci quanto segue:
NUM_SLICES
: il numero di pool di nodi della sezione TPU. In questo caso,NUM_SLICES
è uguale a3
.ACCELERATOR_TYPE
: il tipo di acceleratore TPU utilizzato durante la creazione dei pool di nodi. Ad esempio,tpu-v4-podslice
,tpu-v5-lite-device
otpu-v5-lite-podslice
.TPU_TOPOLOGY
: la topologia fisica per la sezione TPU. Ad esempio4x4x4
o2x2
, a seconda della versione della TPU.NUM_NODES
: il numero di nodi nel pool di nodi. Deve essere zero o il prodotto dei valori definiti inTPU_TOPOLOGY
({A}x{B}x{C}
) diviso per il numero di chip TPU in ogni VM. Per TPU multi-host v4, il numero di chip TPU in ogni VM è quattro. Per TPU v5e multi-host, il numero di chip TPU in ogni VM è uno, quattro o otto. Di conseguenza, se il tuoTPU_TOPOLOGY
è2x4x4
(TPU v4 con quattro chip TPU in ogni VM), ilNUM_NODES
è 32/4, ovvero 8.NUM_CHIPS
: per 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 è uno, quattro o otto. Per saperne di più, consulta la sezione Chip TPU sulla VM in una sezione TPU.
In questo manifest:
- JobSet è un servizio headless con lo stesso nome del nome JobSet, in questo caso
multislice-job
. maxRestarts: 4
indica il numero massimo di volte in cui GKE riavvia il JobSet in caso di errore di un job figlio. Se il JobSet si riavvia raggiunge il valore massimo definito, il JobSet viene contrassegnato come non riuscito.- I campi
parallelism
ecompletions
equivalgono al numero di nodi in ciascun pool di nodi. backoff
è 0 perché l'opzione Multislice supporta solo l'addestramento sincrono multi-controller. Deve essere impostato su 0. Se un pod non funziona correttamente, il job non riesce.- I valori nella sezione di affinità assicurano che sia presente un solo carico di lavoro TPU Multislice in esecuzione in un gruppo di sezioni multiple.
containerPort: 8080
è il trasferimento per il coordinatore MXLAcontainerPort: 8431
è la porta per esportare le metriche di utilizzo della TPUsecurityContext: privileged: true
indica che i nodi hanno la modalità con privilegi abilitata per accedere alle TPU. Per i nodi in GKE 1.28 o versioni successive non è necessario che la modalità con privilegi sia abilitata per accedere alle TPU. Per scoprire di più, consulta Eseguire i container senza la modalità con privilegi.
Applica il manifest:
kubectl apply -f tpu-multislice.yaml
Verifica che il carico di lavoro sia ammesso:
kubectl get jobsets
L'output è simile al seguente:
NAME RESTARTS COMPLETED AGE multislice-job 3s
Monitora lo stato dei pod di cui è stato eseguito il provisioning:
kubectl get pods
L'output è simile al seguente:
NAME READY STATUS RESTARTS AGE multislice-job-slice-0-0-wzq9t 0/1 Completed 0 2m31s multislice-job-slice-0-1-zf4dp 0/1 Completed 0 2m30s multislice-job-slice-1-0-hbfn5 0/1 Completed 0 2m31s multislice-job-slice-1-1-45fgl 0/1 Completed 0 2m30s multislice-job-slice-2-0-wjbp4 0/1 Completed 0 2m30s multislice-job-slice-2-1-lwnvs 0/1 Completed 0 2m30s
Il set di job multislice-job
pianifica, crea ed esegue i pod fino al completamento. I nomi dei pod sono nel formato <jobsetName>-<jobName>-<jobReplicaIndex>-<randomSuffix>
. Il prefisso jobsetName
determina il JobSet a cui appartiene il pod.
Configurazioni aggiuntive
Le seguenti sezioni descrivono le configurazioni aggiuntive che puoi applicare al tuo Multislice.
Abilita hostNetwork sui pod GKE Standard
Per migliorare le prestazioni di rete tra le sezioni TPU, ti consigliamo di attivare
hostNetworking
. Utilizza hostNetwork: true
nella specifica del pod per saltare tutto lo stack di networking di Kubernetes e consentire ai tuoi pod Kubernetes di utilizzare la rete host direttamente per la comunicazione da VM a VM.
Per attivare hostNetworking
, rimuovi le due righe seguenti dalle specifiche del pod:
hostNetwork: true
dnsPolicy: ClusterFirstWithHostNet
Per continuare a utilizzare podHostnames
per il rilevamento dei nodi worker con hostNetwork
, imposta dnsPolicy: ClusterFirstWithHostNet
. Questo è importante quando esegui job di addestramento con ripristino automatico e devi avere gli stessi nomi per ricaricare gli stessi checkpoint.
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 il logging di sistema di 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 la sezione TPU e i worker:
resource.type="k8s_container"
resource.labels.cluster_name=CLUSTER_NAME
labels."k8s-pod/jobset_sigs_k8s_io/jobset-name"=JOBSET_NAME
resource.labels.pod_name:<jobSetName>-<replicateJobName>-<job-index>-<worker-index>
Osservabilità e metriche
Oltre alle metriche TPU generali, sono disponibili altre quattro metriche di runtime TPU specifiche per più sezioni. Queste metriche sono disponibili in GKE versione 1.29.1-gke.1016000 o successive. Il carico di lavoro TPU deve utilizzare la versione JAX 0.4.24
Di seguito sono riportate le metriche disponibili in più sezioni:
- Latenze di trasferimento della rete DCN (Data Center Network): distribuzione delle latenze di trasferimento di rete per il traffico su più sezioni.
- Latenze collettive: distribuzione della latenza collettiva end-to-end per il traffico su più sezioni.
- Latenze di trasferimento da host a dispositivo: distribuzione della latenza di trasferimento da host a dispositivo per ogni blocco di dati per il traffico su più sezioni.
- Latenze di trasferimento da dispositivo a host: distribuzione della latenza di trasferimento dal dispositivo all'host per ogni blocco di dati per il traffico multisezione.
Queste metriche si trovano nello schema del container Kubernetes (k8s_container
):
kubernetes.io/container/multislice/network/dcn_transfer_latencies
kubernetes.io/container/multislice/network/collective_end_to_end_latencies
kubernetes.io/container/multislice/accelerator/host_to_device_transfer_latencies
kubernetes.io/container/multislice/accelerator/device_to_host_transfer_latencies
Sezione TPU e più sezioni
La tabella seguente differenzia l'organizzazione architetturale di una sezione TPU e di una sezione multisezione:
Sezione TPU | Più sezioni | |
---|---|---|
Interconnessione | Il carico di lavoro viene eseguito su una singola sezione TPU. Tutti i chip TPU in una sezione sono collegati a ICI. | Il carico di lavoro viene eseguito su più sezioni TPU. La comunicazione all'interno di una sezione avviene tramite ICI. La comunicazione tra le sezioni avviene tramite DCN. |
Pool di nodi supportati | Sezione TPU con host singolo e Sezione TPU multi-host | Gruppi di sezioni TPU multi-host |
Tipo di carico di lavoro consigliato | IndexedJob o JobSet | JobSet |