Questa pagina è stata tradotta dall'API Cloud Translation.

Esegui il deployment di sezioni TPU multiple in GKE

Autopilot Standard

Questa pagina introduce la configurazione TPU multisezione in Google Kubernetes Engine (GKE). Prima di configurare la configurazione multisezione in GKE, dovresti conoscere i seguenti concetti:

Che cos'è TPU Multislice

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

Ti consigliamo di utilizzare la sezione Multislice se il job è troppo grande per essere adatto a una singola sezione TPU.

Disponibilità multisezione in GKE

La versione Standard supporta il 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 JAX minima supportata è la 2.1.
La multisezione supporta solo pool di nodi di sezioni TPU multi-host. Ad esempio, non puoi utilizzare la multisezione con un elemento ct4p-hightpu-4t con topologia 2x2x1 o ct5lp-hightpu-4t con topologia 2x2, perché si tratta di pool di nodi di sezione TPU con host singolo.
Il multislice supporta solo l'addestramento sincrono del multicontroller.
I carichi di lavoro multisezione 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 initialize gcloud CLI. Se hai già installato gcloud CLI, scarica la versione più recente eseguendo gcloud components update.
Nota: per le installazioni esistenti dell'interfaccia a riga di comando gcloud, assicurati di impostare le proprietà compute/region e compute/zone. Impostando le località predefinite, puoi evitare errori nellgcloud CLI come segue: One of [--zone, --region] must be supplied: Please specify location.

Crea un cluster Standard o un cluster Autopilot che esegue una versione che supporta il multislice. Per le versioni supportate, consulta Disponibilità della multisezione in GKE.
Assicurati che il tuo 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 multisezione

Questa sezione mostra come eseguire un carico di lavoro su una sezione multisezione. Se utilizzi la modalità GKE Autopilot, passa alla sezione Esegui un carico di lavoro multisezione. 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 della sezione TPU multi-host
Verifica lo stato del pool di nodi

Crea il pool di nodi TPU

Puoi creare più di un pool di nodi TPU multi-host. Ai fini di questa guida, crea tre pool di nodi TPU multi-host per eseguire un carico di lavoro multislice. Puoi creare un pool di nodi di 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 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-, utilizza 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 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 macchina disponibili, consulta la pagina Mappatura della configurazione 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 tre tuple ({A}x{B}x{C}), ad esempio 4x4x4.
- TPU v5e: definisci la topologia in due tuple ({A}x{B}), ad esempio 2x2.
Per saperne 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 è quattro. Pertanto, se TPU_TOPOLOGY è 2x4x4 (TPU v4 con quattro chip in ogni VM), allora NUM_NODES è 32/4, che equivale a 8.

Facoltativamente, puoi utilizzare anche i seguenti flag:

RESERVATION_NAME: il nome della prenotazione utilizzato da GKE durante la creazione del pool di nodi. Se ometti questo flag, GKE utilizza i pool di nodi TPU disponibili. Per scoprire di più sulle prenotazioni TPU, consulta Prenotazione TPU.
--spot: imposta il pool di nodi in modo che utilizzi le VM spot per i nodi TPU. Questo valore non può essere modificato dopo la creazione del pool di nodi. Per ulteriori informazioni, consulta la pagina VM spot.
--enable-autoscaling: crea un pool di nodi con scalabilità automatica abilitata. Quando GKE scala un pool di nodi della sezione TPU multi-host, fa lo scale up del pool di nodi atomicamente da zero fino 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 in TPU_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 successiva del provider google.
Aggiungi il blocco seguente 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: nome della risorsa del pool di nodi nel modello Terraform.
- PROJECT_ID: l'ID del tuo 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. Per una maggiore affidabilità del piano di controllo Kubernetes, consigliamo di avere un cluster a livello di regione. Puoi anche utilizzare un cluster di zona. Per saperne 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 di nodi. Deve essere zero o il prodotto del numero di chip TPU diviso per quattro, perché nelle sezioni TPU con più host ogni nodo TPU ha 4 chip. Ad esempio, se TPU_TOPOLOGY è 4x8, ci sono 32 chip, il che significa che NUM_NODES deve essere 8. Per scoprire di più sulle topologie TPU, utilizza la tabella in Mappatura della configurazione TPU.
- TPU_TOPOLOGY: indica la topologia fisica desiderata per la sezione TPU. Il formato della topologia dipende dalla versione della TPU che stai utilizzando:
  - Per TPU v4: definisci la topologia in tre tuple ({A}x{B}x{C}), ad esempio 4x4x4.
  - Per TPU v5e: definisci la topologia in due tuple ({A}x{B}), ad esempio 2x2.
Facoltativamente, puoi utilizzare anche le seguenti variabili:
- RESERVATION_NAME: se utilizzi la prenotazione TPU, questo è l'elenco delle etichette delle risorse di prenotazione da utilizzare durante la creazione del pool di nodi. Per scoprire di più su come compilare RESERVATION_LABEL_VALUES nel campo reservation_affinity, consulta Provider Terraform.
- autoscaling: crea un pool di nodi con scalabilità automatica abilitata. Quando GKE scala un pool di nodi della sezione TPU multi-host, fa lo scale up del pool di nodi atomicamente da zero fino 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 le VM spot per i nodi TPU. Questo valore non può essere modificato dopo la creazione del pool di nodi. Per ulteriori informazioni, consulta la pagina VM spot.

Console

Per creare un pool di nodi con TPU:

Vai alla pagina Google Kubernetes Engine nella console Google Cloud.

Vai a Google Kubernetes Engine
Nell'elenco dei cluster, fai clic sul nome del cluster da modificare.
Fai clic su 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, 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-, utilizza 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.
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 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 l'opzione Tipo di disco di avvio sia Disco permanente standard o Disco permanente SSD.
(Facoltativo) Seleziona la casella di controllo Abilita 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 del tuo progetto.
Usa kubectl, in Cloud Shell, per visualizzare i nodi 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 o tpu-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 multisezione

In questa sezione eseguirai un carico di lavoro JAX che mostra il numero globale di chip TPU nella sezione TPU e 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 TPU. In questo caso, NUM_SLICES è uguale a 3.
ACCELERATOR_TYPE: il tipo di acceleratore TPU utilizzato durante la creazione dei pool di nodi. Ad esempio, tpu-v4-podslice, tpu-v5-lite-device o tpu-v5-lite-podslice.
TPU_TOPOLOGY: la topologia fisica per la sezione TPU. Ad esempio 4x4x4 o 2x2, a seconda della versione della TPU.
NUM_NODES: il numero di nodi nel 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 multi-host, il numero di chip in ogni VM è quattro. Per TPU v5e multi-host, il numero di chip in ogni VM è uno, quattro o otto. Pertanto, se TPU_TOPOLOGY è 2x4x4 (TPU v4 con quattro chip in ogni VM), NUM_NODES è 32/4, che equivale a 8.
NUM_CHIPS: per TPU v4 multi-host, il numero di chip in ogni VM è quattro. Per TPU v5e multi-host, il numero di chip in ogni VM è uno, quattro o otto. Per saperne di più, consulta Chip TPU sulla VM TPU.

In questo file manifest:

JobSet è un servizio headless con lo stesso nome del JobSet, in questo caso multislice-job.
maxRestarts: 4 indica il numero massimo di volte in cui GKE riavvia il JobSet quando un job secondario non va a buon fine. Se il JobSet si riavvia raggiunge il valore massimo definito, viene contrassegnato come non riuscito.
I campi parallelism e completions corrispondono al numero di nodi in ogni pool di nodi.
Il valore backoff è 0 perché Multislice supporta solo l'addestramento multi-controller sincrono. Deve essere impostato su 0. Non eseguire il job in caso di errore di un pod.
I valori nella sezione di affinità assicurano che in un gruppo di multislice sia in esecuzione un solo carico di lavoro TPU Multislice.
containerPort: 8080 è il porto per il coordinatore MXLA
containerPort: 8431 è la porta su cui esportare le metriche di utilizzo della TPU
securityContext: privileged: true indica che i nodi hanno la modalità privilegi abilitata per l'accesso alle TPU. Per i nodi in GKE versione 1.28 o successive non è necessario abilitare la modalità con privilegi per accedere alle TPU. Per scoprire di più, consulta Eseguire container senza modalità con privilegi.

Applica il manifest:
```
kubectl apply -f tpu-multislice.yaml
```

Conferma 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 JobSet multislice-job pianifica, crea e esegue i pod fino al completamento. Il formato dei nomi dei pod è <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 alla tua sezione multisezione.

Abilita hostNetwork sui tuoi pod GKE Standard

Per migliorare le prestazioni di rete tra le sezioni TPU, ti consigliamo di attivare hostNetworking. Usa hostNetwork: true nella specifica dei pod per saltare tutti gli 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 seguenti due righe 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, incluse le VM TPU, sono visibili in Esplora log, se hai abilitato il logging di sistema GKE nel 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 multisezione. 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 multisezione disponibili:

Latenze di trasferimento DCN (Rete di data center): distribuzione delle latenze di trasferimento di rete per il traffico multisezione.
Latenze collettive: distribuzione della latenza collettiva end-to-end per il traffico multisezione.
Latenze di trasferimento da host a dispositivo: distribuzione della latenza di trasferimento da host a dispositivo per ogni blocco di dati per il traffico multisezione.
Latenze di trasferimento da dispositivo a host: distribuzione della latenza di trasferimento tra dispositivo e 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 multisezione

La seguente tabella differenzia l'organizzazione dell'architettura di una sezione TPU e di una sezione multisezione:

	Sezione TPU	Multisezione
Interconnessione	Il carico di lavoro viene eseguito su una singola sezione TPU. Tutti i chip TPU in una sezione sono connessi con 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 singolo host e sezione TPU multi-host	Gruppi di sezioni TPU multi-host
Tipo di carico di lavoro consigliato	IndexedJob o JobSet	JobSet

Passaggi successivi

Scopri come orchestrare carichi di lavoro multisezione con sezioni TPU