Questa pagina è stata tradotta dall'API Cloud Translation.

Esegui il deployment di TPU Multislices in GKE

Autopilot Standard

Questa pagina mostra come eseguire il deployment dei carichi di lavoro in Google Kubernetes Engine (GKE) utilizzando la configurazione multislice di Cloud TPU per un addestramento su larga scala e conveniente.

Questo tutorial è rivolto a ingegneri di machine learning (ML) e amministratori e operatori di piattaforme che vogliono utilizzare l'orchestrazione dei container Kubernetes per gestire l'addestramento, l'ottimizzazione e l'inferenza di modelli su larga scala utilizzando le TPU. Per scoprire di più sui ruoli comuni e sulle attività di esempio a cui viene fatto riferimento nei contenuti di Google Cloud , consulta Ruoli utente e attività comuni di GKE.

Prima di configurare Multislice in GKE, assicurati di avere familiarità con i seguenti concetti:

Che cos'è TPU Multislice

TPU Multislice è l'organizzazione architetturale delle VM in una sezione TPU in cui due o più sezioni Cloud TPU comunicano tramite la rete del data center (DCN). Multislice consente l'addestramento full-stack, economico e su larga scala con 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 della TPU. La comunicazione tra i chip TPU all'interno di una sezione avviene tramite interconnessioni tra chip (ICI). La comunicazione tra le sezioni avviene tramite la DCN.

Ti consigliamo di utilizzare Multislice se il tuo job è troppo grande per essere inserito in una singola slice TPU.

Disponibilità di Multislice 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 supportata di JAX è 2.1.
Multislice supporta solo i node pool 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 di slice TPU a singolo host.
Multislice supporta solo l'addestramento sincrono di più controller.
I carichi di lavoro multi-slice possono essere eseguiti solo su slice TPU che condividono lo stesso tipo, dimensione e topologia di TPU.
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à, installala e poi inizializza gcloud CLI. Se hai già installato gcloud CLI, scarica l'ultima versione eseguendo il comando gcloud components update. Le versioni precedenti di gcloud CLI potrebbero non supportare l'esecuzione dei comandi in questo documento.
Nota: per le installazioni esistenti di gcloud CLI, assicurati di impostare la proprietà compute/region. Se utilizzi principalmente cluster zonali, imposta invece compute/zone. Se imposti una località predefinita, puoi evitare errori in gcloud CLI come il seguente: One of [--zone, --region] must be supplied: Please specify location. Potresti dover specificare la posizione in determinati comandi se la posizione del cluster è diversa da quella predefinita che hai impostato.

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

Esegui un carico di lavoro su una Multislice

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

Crea il pool di nodi dello slice TPU

Puoi creare più di un node pool TPU multi-host. Ai fini di questa guida, crea tre node pool TPU multi-host per eseguire un carico di lavoro multislice. I passaggi per creare un pool di nodi TPU slice multi-host variano a seconda che utilizzi Ironwood (TPU7x) o una versione precedente della TPU.

Ironwood (TPU7x)

Puoi creare un node pool slice TPU multi-host nella versione Ironwood (TPU7x) utilizzando Google Cloud CLI o Terraform:

gcloud

Per creare un pool di nodi TPU slice multi-host con Ironwood (TPU7x), devi prima creare una policy del carico di lavoro.

Crea una policy del workload:
```
gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
    --type=HIGH_THROUGHPUT \
    --accelerator-topology=TPU_TOPOLOGY \
    --project=PROJECT_ID \
    --region=REGION
```
Sostituisci quanto segue:
- WORKLOAD_POLICY_NAME: un nome per la policy del workload.
- TPU_TOPOLOGY: la topologia TPU Ironwood (TPU7x). Ad esempio, 2x2x2. Per visualizzare tutte le topologie Ironwood (TPU7x) supportate, consulta la sezione relativa alla topologia.
- PROJECT_ID: il tuo ID progetto Google Cloud .
- REGION: la regione della policy del workload. Una policy del workload è una risorsa di regione e può essere riutilizzata in pool di nodi che condividono la stessa topologia.
Crea il pool di nodi con la policy del workload:
```
gcloud container node-pools create NODE_POOL_NAME \
  --cluster=CLUSTER_NAME \
  --machine-type=tpu7x-standard-4t \
  --placement-policy=WORKLOAD_POLICY_NAME \
  --location=CONTROL_PLANE_LOCATION \
  --node-locations=NODE_ZONE \
  --project=PROJECT_ID \
  --reservation=RESERVATION_NAME \
  --reservation-affinity=specific
```
Sostituisci quanto segue:
- NODE_POOL_NAME: il nome del nuovo pool di nodi.
- CLUSTER_NAME: il nome del tuo cluster GKE.
- WORKLOAD_POLICY_NAME: il nome della policy del workload che hai creato.
- CONTROL_PLANE_LOCATION: la posizione di Compute Engine del control plane del tuo cluster. Fornisci una regione per i cluster regionali o una zona per i cluster zonali.
- NODE_ZONE: il nome della zona in base alla versione della TPU che vuoi utilizzare. Per identificare una località disponibile, consulta Disponibilità di TPU in GKE.
- PROJECT_ID: il tuo ID progetto Google Cloud .
- RESERVATION_NAME: il nome della prenotazione da utilizzare.
In questo comando, il flag --tpu-topology è stato sostituito dal flag --placement-policy.

Terraform

Assicurati di utilizzare la versione 4.84.0 o successive del provider google.
Crea una policy del workload:
```
resource "google_compute_resource_policy" {
  name   = "WORKLOAD_POLICY_NAME"
  region = CLUSTER_LOCATION
  workload_policy {
    type = "HIGH_THROUGHPUT"
    accelerator_topology = "TPU_TOPOLOGY"
  }
}
```
Sostituisci quanto segue:
- WORKLOAD_POLICY_NAME: un nome per la policy del workload.
- CLUSTER_LOCATION: la località di computing per il cluster. Consigliamo di utilizzare un cluster regionale per una maggiore affidabilità del control plane Kubernetes. Puoi anche utilizzare un cluster zonale. Per maggiori informazioni, vedi Selezionare una versione e una topologia TPU.
- TPU_TOPOLOGY: la topologia TPU Ironwood (TPU7x). Ad esempio, 2x2x2. Per visualizzare tutte le topologie Ironwood (TPU7x) supportate, consulta la sezione Pianificare le TPU.
Per saperne di più sul riferimento google_compute_resource_policy, consulta Provider Terraform.
Nella configurazione di Terraform, aggiungi il seguente blocco:
```
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]
    }
    flex_start = false
    spot = true
  }

  placement_policy {
    policy_name = WORKLOAD_POLICY_NAME
  }
}
```
Sostituisci quanto segue:
- NODE_POOL_RESOURCE_NAME: il nome della risorsa del node pool nel modello Terraform.
- PROJECT_ID: il tuo 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.
- 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 node pool. Deve essere zero o il prodotto del numero di chip TPU diviso per quattro, perché nelle sezioni TPU multihost ogni nodo della sezione TPU ha quattro chip. Ad esempio, se TPU_TOPOLOGY è 4x8, allora 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 della TPU.
- TPU_TOPOLOGY: indica la topologia fisica selezionata per la sezione TPU. Il formato della topologia dipende dalla versione della TPU che stai utilizzando. Per saperne di più sulle topologie TPU, utilizza la tabella in Scegliere una topologia.
Se vuoi, puoi utilizzare anche le seguenti variabili:
- RESERVATION_NAME: se utilizzi una prenotazione TPU, fornisci un elenco di etichette delle risorse di prenotazione da utilizzare durante la creazione del pool di nodi. Per scoprire di più su come compilare ilRESERVATION_LABEL_VALUES nel campo reservation_affinity, consulta Provider Terraform.
- autoscaling: crea un pool di nodi con la scalabilità automatica abilitata. Quando GKE scala unpool di nodil di sezioni TPU multi-host, lo scala in modo atomico da zero alle dimensioni massime.
  - MAX_NODES: la dimensione massima del node pool. Il valore 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. Ad esempio, se TPU_TOPOLOGY è 2x2x2, il prodotto è 8. Poiché ogni VM in tpu7x-standard-4t ha 4 chip, il numero di nodi è 2.
- spot: il pool di nodi che utilizzerà le VM spot per i nodi dello slice TPU. Questa impostazione non può essere modificata dopo la creazione del pool di nodi. Per saperne di più, consulta VM spot.
- flex_start: il pool di nodi che utilizzerà l'opzione di consumo avvio flessibile. Questa impostazione non può essere impostata su true se spot è attivato.

Altre versioni di TPU

Puoi creare un pool di nodi slice TPU multihost nelle versioni v3, v4, v5p, v5e e Trillium (v6e) utilizzando Google Cloud CLI, Terraform o la console Google Cloud .

gcloud

  gcloud container node-pools create POOL_NAME \
      --location=CONTROL_PLANE_LOCATION \
      --cluster=CLUSTER_NAME \
      --node-locations=NODE_ZONE \
      --machine-type=MACHINE_TYPE \
      --tpu-topology=TPU_TOPOLOGY \
      [--num-nodes=NUM_NODES] \
      [--spot \]
      [--flex-start \]
      [--enable-autoscaling \
        --max-nodes MAX_NODES]
      [--reservation-affinity=specific \
      --reservation=RESERVATION_NAME] \
      [--node-labels cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME,cloud.google.com/gke-workload-type=HIGH_AVAILABILITY]
      [--placement-type=COMPACT]

Sostituisci quanto segue:

POOL_NAME: il nome del nuovo pool di nodi.
CONTROL_PLANE_LOCATION: la posizione di Compute Engine del piano di controllo del cluster. Fornisci una regione per i cluster regionali o una zona per i cluster zonali.
CLUSTER_NAME: il nome del cluster.
NODE_ZONES: il nome della zona in base alla versione della TPU che vuoi utilizzare. Per identificare una località disponibile, consulta Disponibilità di TPU in GKE.
MACHINE_TYPE: il tipo di macchina da utilizzare per i nodi. Per scoprire di più sui tipi di macchina disponibili, vedi Scegliere la versione della TPU.
TPU_TOPOLOGY: la topologia fisica della sezione TPU. Il formato della topologia dipende dalla versione della TPU. Per maggiori informazioni sulle topologie TPU, utilizza la tabella in Scegliere una topologia.

Per ulteriori informazioni, vedi Topologia.

Se vuoi, puoi anche utilizzare i seguenti flag:
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 il tuo TPU_TOPOLOGY è 2x4x4 (TPU v4 con quattro chip in ogni VM), allora il NUM_NODES è 32/4, ovvero 8. Se ometti questo flag, il numero di nodi viene calcolato e impostato come predefinito in base alla topologia e al tipo di macchina.
RESERVATION_NAME: il nome della prenotazione utilizzata da GKE durante la creazione del pool di nodi. Se ometti questo flag, GKE utilizza i node pool delle sezioni TPU disponibili. Per saperne di più sulle prenotazioni TPU, consulta Prenotazione TPU.
--spot: imposta il pool di nodi in modo che utilizzi le VM spot per i nodi dello slice TPU. Questo valore non può essere modificato dopo la creazione del pool di nodi. Per ulteriori informazioni, consulta VM spot.
--flex-start: imposta il pool di nodi in modo che utilizzi le VM con avvio flessibile. Le VM con avvio flessibile vengono create utilizzando l'opzione di consumo flex-start, supportata in GKE versione 1.33.0-gke.1712000 o successive.
--enable-autoscaling: crea un pool di nodi con la scalabilità automatica abilitata. Quando GKE scala unpool di nodil di sezioni TPU multi-host, lo scala in modo atomico da zero alle dimensioni massime.
- MAX_NODES: la dimensione massima del node pool. 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.
--node-label=cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME, cloud.google.com/gke-workload-type=HIGH_AVAILABILITY: indica a GKE che il pool di nodi della sezione TPU multi-host è una raccolta. Utilizza questo flag se si verificano le seguenti condizioni:
- Il pool di nodi esegue i carichi di lavoro di inferenza.
- Il pool di nodi utilizza TPU Trillium.
- Le VM spot non supportano la pianificazione della raccolta.
Per maggiori informazioni sulla gestione della pianificazione della raccolta, vedi Gestire la pianificazione della raccolta nelle sezioni TPU multi-host.
--placement-type=COMPACT: crea un pool di nodi con il posizionamento compatto abilitato. Questa opzione deve essere utilizzata con il flag --tpu-topology. Per saperne di più, consulta Crea una policy di posizionamento compatto e Topologia TPU.

Terraform

Assicurati di utilizzare la versione 4.84.0 o successive del provider google.
Aggiungi il seguente blocco alla 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]
    }
    flex_start = false
    spot = true
  }

  placement_policy {
    type = "COMPACT"
    tpu_topology = TPU_TOPOLOGY
  }
}
```
Sostituisci quanto segue:
- NODE_POOL_RESOURCE_NAME: il nome della risorsa del node pool nel modello Terraform.
- PROJECT_ID: il tuo 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: la località di computing del cluster. Consigliamo di utilizzare 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 node pool. Deve essere zero o il prodotto del numero di chip TPU diviso per quattro, perché nelle sezioni TPU multihost ogni nodo della sezione TPU ha 4 chip. Ad esempio, se TPU_TOPOLOGY è 4x8, allora 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 della TPU.
- TPU_TOPOLOGY: indica la topologia fisica della sezione TPU. Il formato della topologia dipende dalla versione della TPU che stai utilizzando. Per saperne di più sulle topologie TPU, utilizza la tabella in Scegliere una topologia.
Se vuoi, 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 ilRESERVATION_LABEL_VALUES nel campo reservation_affinity, consulta Provider Terraform.
- autoscaling: crea un pool di nodi con la scalabilità automatica abilitata. Quando GKE scala unpool di nodil di sezioni TPU multi-host, lo scala in modo atomico da zero alle dimensioni massime.
  - MAX_NODES: è la dimensione massima del node pool. 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 della sezione TPU. Questo valore non può essere modificato dopo la creazione del pool di nodi. Per saperne di più, consulta VM spot.
- flex_start: imposta il pool di nodi in modo che utilizzi l'opzione di consumo avvio flessibile. Non può essere impostato su true se spot è abilitato.

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 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:
- CT3: TPU v3, dispositivo host singolo
- CT3P: TPU v3, sezione di pod multi-host
- CT4P: TPU v4
- CT5LP: TPU v5e
- CT5P: TPU v5p
- CT6E: TPU Trillium (v6e)
Nel menu a discesa Tipo di macchina, seleziona il nome della macchina da utilizzare per i nodi. Utilizza la tabella Scegli la versione della 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 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

Recupera le credenziali per poter utilizzare kubectl per accedere al cluster:
```
gcloud container clusters get-credentials CLUSTER_NAME \
    --project=PROJECT_ID \
    --location=CONTROL_PLANE_LOCATION
```
Sostituisci quanto segue:
- CLUSTER_NAME: il nome del cluster.
- PROJECT_ID: il tuo ID progetto.
- CONTROL_PLANE_LOCATION: la posizione di Compute Engine del control plane del tuo cluster. Fornisci una regione per i cluster regionali o una zona per i cluster zonali.
Utilizza kubectl in Cloud Shell per visualizzare i nodi della slice TPU:
```
kubectl get nodes -l cloud.google.com/gke-tpu-accelerator=ACCELERATOR_TYPE \
   -l cloud.google.com/gke-tpu-topology=TPU_TOPOLOGY
```
Sostituisci quanto segue:
- TPU_ACCELERATOR: il tipo di acceleratore TPU che hai utilizzato durante la creazione dei node pool. Ad esempio, tpu-v4-podslice 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 nella sezione TPU e poi esce.

Per eseguire un workload 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: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                ports:
                - containerPort: 8471
                - containerPort: 8080
                - containerPort: 8431
                command:
                - bash
                - -c
                - |
                  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: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
                ports:
                - containerPort: 8471
                - containerPort: 8080
                - containerPort: 8431
                securityContext:
                  privileged: true
                command:
                - bash
                - -c
                - |
                  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 node pool di sezioni TPU. In questo caso, NUM_SLICES è uguale a 3.
ACCELERATOR_TYPE: il tipo di acceleratore TPU che hai utilizzato durante la creazione dei node pool. Ad esempio, tpu-v4-podslice 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 TPU in ogni VM. Per la 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 è 1, 4 o 8. Pertanto, se il tuo TPU_TOPOLOGY è 2x4x4 (TPU v4 con quattro chip TPU in ogni VM), il NUM_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 è 1, 4 o 8. Per saperne di più, consulta Chip TPU sulla VM in una sezione TPU.

In questo manifest:

JobSet è un servizio headless con lo stesso nome di JobSet, in questo caso è multislice-job.
L'annotazione alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool configura l'affinità pod per garantire che tutti i pod siano pianificati sulla stessa sezione.
Il campo image è impostato sull'ultima immagine JAX AI. Per impostare una versione diversa, consulta l'elenco delle immagini AI JAX attuali.
maxRestarts: 4 indica il numero massimo di volte in cui GKE riavvia il JobSet quando un job secondario non riesce. Se i riavvii del JobSet raggiungono il massimo definito, il JobSet viene contrassegnato come non riuscito.
I campi parallelism e completions corrispondono al numero di nodi in ognipool di nodil.
Il valore di backoff è 0 perché Multislice supporta solo l'addestramento sincrono multi-controller. Deve essere impostato su 0. Il job non riesce se un pod non riesce.
I valori nella sezione Affinità assicurano che in un gruppo di Multislice venga eseguito un solo carico di lavoro TPU Multislice.
containerPort: 8080 è la porta per il coordinatore MXLA
containerPort: 8431 è la porta per esportare le metriche di utilizzo della TPU
securityContext: privileged: true indica che i nodi hanno la modalità privilegiata abilitata per accedere alle TPU. I nodi in GKE 1.28 o versioni successive non devono avere la modalità con privilegi abilitata per accedere alle TPU. Per saperne di più, consulta Eseguire container senza modalità privilegiata.

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

Verifica che il workload 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

multislice-job JobSet 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.

(Facoltativo) Rimuovi il carico di lavoro JAX:
```
kubectl delete -f tpu-multislice.yaml
```

Configura le impostazioni aggiuntive

Le sezioni seguenti descrivono le configurazioni aggiuntive che puoi applicare a Multislice.

Migliorare le prestazioni di rete con hostNetwork

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

Per attivare hostNetworking, aggiungi le seguenti due righe alla 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.

Se utilizzi TPU Trillium (v6e) o Ironwood (TPU7x) (anteprima) e i tuoi pod utilizzano hostNetworking, installa il seguente DaemonSet per la messa a punto di /proc/sys/net/ipv4/tcp_rmem sul nodo.

kubectl apply -f https://raw.githubusercontent.com/GoogleCloudPlatform/ai-on-gke/51bf3dcab6ff658cf62cc32867f96860bf58dfdc/scripts/network-setup/v6e-increase-rmem.yaml

Migliorare le prestazioni di rete senza hostNetwork su TPU Trillium o Ironwood (TPU7x)

Se utilizzi TPU Trillium o Ironwood (TPU7x) (anteprima) e i tuoi pod non possono utilizzare hostNetworking, attiva il multi-networking con la modalità netdevice per ottenere le migliori prestazioni di rete. Il supporto NIC in modalità netdevice con multi-network trasferisce la NIC della VM direttamente al pod, bypassando Kubernetes e GKE Dataplane V2.

Il tipo di macchina ct6e-standard-4t è supportato da due NIC fisiche. Kubernetes richiede una vNIC che non può essere passata ai pod. Pertanto, ogni nodo deve avere tre vNIC per consentire ai pod di accedere direttamente a due vNIC per ottenere le migliori prestazioni di entrambe le NIC fisiche.

Per attivare la modalità netdevice per ct6e-standard-4t, completa i seguenti passaggi:

Crea due VPC aggiuntivi che supportano la modalità netdevice
Crea un cluster GKE con funzionalità multi-rete

Configura due netdevice reti. Ad esempio, puoi utilizzare i seguenti oggetti GKENetworkParamSet e Network (SECOND_VPC e THIRD_VPC sono i VPC creati nel passaggio precedente):

apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: tpu-second
spec:
  vpc: SECOND_VPC
  vpcSubnet: SECOND_VPC_SUBNET
  deviceMode: NetDevice
---
apiVersion: networking.gke.io/v1
kind: GKENetworkParamSet
metadata:
  name: tpu-third
spec:
  vpc: THIRD_VPC
  vpcSubnet: SECOND_VPC_SUBNET
  deviceMode: NetDevice
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: tpu-second
spec:
  provider: "GKE"
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: tpu-second
---
apiVersion: networking.gke.io/v1
kind: Network
metadata:
  name: tpu-third
spec:
  provider: "GKE"
  type: "Device"
  parametersRef:
    group: networking.gke.io
    kind: GKENetworkParamSet
    name: tpu-third

Connettere i pod a tre reti. Ad esempio, puoi utilizzare le seguenti annotazioni nella specifica del pod:

metadata:
  annotations:
    networking.gke.io/default-interface: 'eth0'
    networking.gke.io/interfaces: |
      [
        {"interfaceName":"eth0","network":"default"},
        {"interfaceName":"eth1","network":"tpu-second"},
        {"interfaceName":"eth2","network":"tpu-third"},
      ]

Applica i sysctl di rete all'interno del pod, nel container init o nel container dell'applicazione. Ad esempio, puoi aggiungere il seguente init container alle specifiche del pod:

initContainers:
- name: "network-optimization-sysctls"
  image: "busybox"
  securityContext:
    privileged: true
  command:
  - bash
  - -c
  - |
    echo 5000 > /proc/sys/net/ipv4/tcp_rto_min_us
    echo 1 > /proc/sys/net/ipv4/tcp_no_metrics_save
    echo 0 > /proc/sys/net/ipv4/tcp_slow_start_after_idle
    echo 131072 > /proc/sys/net/core/optmem_max
    echo "4096 41943040 314572800" > /proc/sys/net/ipv4/tcp_rmem

Best practice:

Utilizza le interfacce eth1 e eth2 per prestazioni di rete migliori, anziché l'interfaccia eth0. Per farlo, aggiungi export LIBTPU_INIT_ARGS="$LIBTPU_INIT_ARGS --megascale_grpc_interface_prefixes=eth1,eth2,lo" alla specifica del workload.

Abilita il 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 la registrazione dei log di sistema 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 lo slice 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>

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

Abilitare metriche aggiuntive

Oltre alle metriche TPU generali, esistono altre quattro metriche di runtime TPU specifiche per multislice. Queste metriche sono disponibili in GKE versione 1.29.1-gke.1016000 o successive. Il workload TPU deve utilizzare la versione 0.4.24 di JAX

Di seguito sono riportate le metriche multislice disponibili:

Latenze di trasferimento DCN (Data Center Network): distribuzione delle latenze di trasferimento di 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 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 multislice

La tabella seguente distingue l'organizzazione architetturale di uno slice TPU e di un multislice:

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

Esegui la pulizia delle risorse

Il modo più semplice per eliminare la fatturazione è eliminare il progetto Google Cloud che hai creato per il tutorial. In alternativa, puoi eliminare le singole risorse.

Elimina il progetto

Attenzione: l'eliminazione di un progetto ha i seguenti effetti:

L'intero contenuto del progetto viene eliminato. Se hai utilizzato un progetto esistente per le attività descritte in questo documento, eliminerai anche tutto il lavoro che hai svolto nel progetto.
Gli ID progetto personalizzati non sono più disponibili. Quando hai creato questo progetto, potresti aver creato un ID progetto personalizzato che vuoi utilizzare in futuro. Per conservare gli URL che utilizzano l'ID progetto, ad esempio un URL appspot.com, elimina le risorse selezionate all'interno del progetto anziché eliminare l'intero progetto.

Se intendi esplorare più architetture, tutorial o guide rapide, puoi riutilizzare i progetti ed evitare così di superare i limiti di quota.

Delete a Google Cloud project:

gcloud projects delete PROJECT_ID

Elimina singole risorse

Elimina il cluster GKE:

```sh
gcloud container clusters delete  CLUSTER_NAME \
   --project=PROJECT_ID  \
   --location=CONTROL_PLANE_LOCATION
```

Passaggi successivi

Scopri come orchestrare i carichi di lavoro multi-sezione con le sezioni TPU

Esegui il deployment di TPU Multislices in GKE Mantieni tutto organizzato con le raccolte Salva e classifica i contenuti in base alle tue preferenze.

Che cos'è TPU Multislice

Disponibilità di Multislice in GKE

Prima di iniziare

Esegui un carico di lavoro su una Multislice

Prepara un pool di nodi in modalità Standard

Crea il pool di nodi dello slice TPU

Ironwood (TPU7x)

gcloud

Terraform

Altre versioni di TPU

gcloud

Terraform

Console

Verifica lo stato del pool di nodi

Esegui un carico di lavoro multislice

Autopilot

Standard

Configura le impostazioni aggiuntive

Migliorare le prestazioni di rete con hostNetwork

Migliorare le prestazioni di rete senza hostNetwork su TPU Trillium o Ironwood (TPU7x)

Abilita il logging

Abilitare metriche aggiuntive

Sezione TPU e multislice

Esegui la pulizia delle risorse

Elimina il progetto

Elimina singole risorse

Passaggi successivi

Esegui il deployment di TPU Multislices in GKE