Questa pagina è stata tradotta dall'API Cloud Translation.

Esegui il deployment di TPU Multislice in GKE

Autopilot Standard

Questa pagina mostra come eseguire il deployment dei carichi di lavoro in Google Kubernetes Engine (GKE) utilizzando la configurazione Cloud TPU Multislice per l'addestramento su larga scala e a costi contenuti.

Prima di configurare Multislice in GKE, assicurati di conoscere i seguenti concetti:

Che cos'è TPU Multislice

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

Ti consigliamo di utilizzare Multislice se il tuo job è troppo grande per adattarsi a un singolo slice TPU.

Disponibilità di più slice 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 è 2.1.
Multislice supporta solo pool di nodi di sezioni TPU multi-host. Ad esempio, non puoi utilizzare Multislice con un ct4p-hightpu-4t con una topologia 2x2x1 o un ct5lp-hightpu-4t con una topologia 2x2, perché si tratta di pool di nodi con sezioni TPU a singolo host.
Multislice supporta solo l'addestramento multicontroller sincrono.
I carichi di lavoro multislice possono essere eseguiti solo su sezioni TPU che condividono lo stesso tipo, le stesse dimensioni e la stessa topologia.
Multislice non supporta TPU v3.

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.
Nota: per le installazioni esistenti di gcloud CLI, assicurati di impostare le proprietà compute/region e compute/zone. Se imposti le posizioni predefinite, puoi evitare errori in gcloud CLI come il seguente: One of [--zone, --region] must be supplied: Please specify location.

Crea un cluster standard o un cluster Autopilot che esegue una versione che supporta Multislice. Per le versioni supportate, consulta Disponibilità di Multislice in GKE.
Assicurati che il tuo progetto disponga di una quota sufficiente per Cloud TPU in GKE.
Installa JobSet versione 0.2.3 o successive.

Esegui un carico di lavoro su un Multislice

Questa sezione mostra come eseguire un carico di lavoro su un Multislice. Se utilizzi la modalità GKE Autopilot, vai alla sezione Eseguire un carico di lavoro multislice. I cluster Autopilot che eseguono la versione 1.29.2-gke.1521000 o successive attivano le TPU per impostazione predefinita.

Preparare un pool di nodi in modalità standard

Questa sezione illustra i seguenti passaggi:

Crea tre pool di nodi di slice TPU multi-host
Verifica lo stato del pool di nodi

Crea il pool di nodi del segmento TPU

Puoi creare più pool di nodi di sezioni TPU multi-host. Ai fini di questa guida, crea tre pool di nodi di slice TPU multi-host per eseguire un carico di lavoro multislice. 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 in base alla versione TPU che vuoi utilizzare. Per identificare una località disponibile, consulta la sezione Disponibilità di 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 Scegliere la versione TPU.
TPU_TOPOLOGY: la topologia fisica per la sezione TPU. Il formato della topologia dipende dalla versione della TPU. Per scoprire di più sulle topologie TPU, utilizza la tabella in Scegliere una topologia.

Per saperne di più, consulta Topologia.
NUM_NODES: il numero di nodi nel pool di nodi. Deve essere pari a zero o al 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 tuo TPU_TOPOLOGY è 2x4x4 (TPU v4 con quattro chip in ogni VM), il NUM_NODES è 32/4, ovvero 8.

Se vuoi, 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 node pool di sezioni TPU disponibili. Per scoprire di più sulle prenotazioni delle TPU, consulta la sezione Prenotazione TPU.
--spot: imposta il pool di nodi in modo da utilizzare 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 la sezione VM spot.
--enable-autoscaling: crea un pool di nodi con la scalabilità automatica abilitata. Quando GKE esegue il ridimensionamento di un pool di nodi di sezioni TPU multi-host, lo esegue atomicamente da zero alla dimensione massima.
- 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

Assicurati di utilizzare la versione 4.84.0 o successive del fornitore google.
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 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: posizione di calcolo per il cluster. Ti 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 pool di nodi. 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, significa che ci sono 32 chip, il che significa che NUM_NODES deve essere 8. Per scoprire di più sulle topologie TPU, utilizza la tabella in Scegliere la versione TPU.
- TPU_TOPOLOGY: indica la topologia fisica preferita per la sezione TPU. Il formato della topologia dipende dalla versione della TPU in uso. Per scoprire di più sulle topologie TPU, utilizza la tabella in Scegliere una topologia.
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 esegue il ridimensionamento di un pool di nodi di sezioni TPU multi-host, lo esegue atomicamente 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 della sezione TPU. 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:

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 il nome della zona in base alla versione della TPU che vuoi utilizzare. Per identificare una località disponibile, consulta la sezione Disponibilità di TPU in GKE.
Nel riquadro di navigazione, fai clic su Nodi.
Nella sezione Configurazione macchina, seleziona TPU.
Nel menu a discesa Serie, seleziona una delle seguenti opzioni:
- CT3P: per TPU v3.
- 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 Scegliere la versione TPU per scoprire come definire il tipo di macchina e la topologia TPU che creano un pool di nodi di sezioni TPU multi-host.
Nel menu a discesa Topologia TPU, seleziona la topologia fisica per il seme TPU.
Nella finestra di dialogo Modifiche necessarie, fai clic su Apporta modifiche.
Assicurati che Tipo di disco di avvio sia Disco permanente standard o Disco permanente SSD.
Se vuoi, 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

Recupera 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.
Utilizza kubectl in Cloud Shell per visualizzare i nodi dei tuoi slice 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 quando hai creato i 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 multislice

In questa sezione esegui un carico di lavoro JAX che mostra il numero globale di chip TPU nel segmento 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 di sezioni TPU. In questo caso, NUM_SLICES è uguale a 3.
ACCELERATOR_TYPE: il tipo di acceleratore TPU utilizzato per creare i 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 pari a zero o al prodotto dei valori definiti in TPU_TOPOLOGY ({A}x{B}x{C}) diviso per il numero di chip TPU in ogni VM. Per le TPU v4 multi-host, il numero di chip TPU in ogni VM è quattro. Per le TPU v5e multi-host, il numero di chip TPU in ogni VM è uno, quattro o otto. Pertanto, se il tuo TPU_TOPOLOGY è 2x4x4 (TPU v4 con quattro chip TPU in ogni VM), NUM_NODES è 32/4, ovvero 8.
NUM_CHIPS: per TPU v4 multi-host, il numero di chip TPU in ogni VM è quattro. Per le 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:

Il JobSet è un servizio headless con lo stesso nome del JobSet, in questo caso multislice-job.
L'annotazione alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool configura l'affinità dei pod per garantire che tutti i pod vengano pianificati nello stesso slice.
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 numero di riavvii del JobSet raggiunge il valore massimo definito, il JobSet viene contrassegnato come non riuscito.
I campi parallelism e completions corrispondono al numero di nodi in ciascun pool di nodi.
backoff è 0 perché Multislice supporta solo l'addestramento con più controller in sincrono. Deve essere impostato su 0. Interrompi il job quando un pod non va a buon fine.
I valori nella sezione di affinità assicurano che sia in esecuzione un solo carico di lavoro TPU Multislice in un gruppo di Multislice.
containerPort: 8080 è la porta per il coordinatore MXLA
containerPort: 8431 è la porta per esportare le metriche di utilizzo delle TPU
Il simbolo securityContext: privileged: true indica che i nodi hanno attivato la modalità privilegiata per accedere alle TPU. I nodi in GKE versione 1.28 o successive non devono avere la modalità privilegiata abilitata per accedere alle TPU. Per saperne di più, consulta Eseguire i container senza modalità privilegiata.

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

Verifica che il carico di lavoro sia accettato:

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 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 nei pod GKE Standard

Per migliorare le prestazioni della rete tra i slice TPU, ti consigliamo di attivare hostNetworking. Utilizza hostNetwork: true nella specifica del pod per saltare tutto lo stack di rete Kubernetes e consentire ai pod Kubernetes di utilizzare direttamente la rete dell'host per la comunicazione da VM a VM.

Per attivare hostNetworking, rimuovi le due righe seguenti dalla specifica 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 ripresa automatica 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 della sezione TPU, sono visibili in Esplora log, se hai attivato i log di sistema GKE nel tuo cluster.

Puoi visualizzare i log di GKE utilizzando Esplora log con il seguente filtro per visualizzare i log del contenitore 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 i segmenti e i worker TPU:

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>

Per scoprire di più, consulta Visualizzare i log di TPU GKE.

Osservabilità e metriche

Oltre alle metriche TPU generali, sono disponibili altre quattro metriche di runtime TPU specifiche per multislice. 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 multislice disponibili:

Latenze di trasferimento della rete del data center (DCN): distribuzione delle latenze di trasferimento della rete per il traffico multislice.
Latenze collettive: distribuzione della latenza collettiva end-to-end per il traffico multislice.
Latenze di trasferimento da host a dispositivo: distribuzione della latenza di trasferimento da host a dispositivo per ogni blocco di dati per il traffico multislice.
Latenze di trasferimento da dispositivo a host: distribuzione della latenza di trasferimento da dispositivo a host per ogni blocco di dati per il traffico multislice.

Queste metriche si trovano nello schema del contenitore 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 Multislice

La tabella seguente distingue l'organizzazione architettonica di uno slice TPU e di una TPU multislice:

	Sezione TPU	Multislice
Interconnessione	Il carico di lavoro viene eseguito su una singola sezione TPU. Tutti i chip TPU in una sezione sono collegati con ICI.	Il carico di lavoro viene eseguito su più sezioni TPU. La comunicazione all'interno di uno slice avviene tramite ICI. La comunicazione tra i vari slice avviene tramite la rete DCN.
Node pool supportati	Sezione TPU mono-host e sezione TPU multi-host	Gruppi di sezioni TPU multi-host
Tipo di workload consigliato	IndexedJob o JobSet	JobSet

Passaggi successivi

Scopri come orchestrare i carichi di lavoro multislice con le sezioni TPU