Esegui il deployment di TPU Multislice in GKE

Questa pagina mostra come eseguire il deployment dei carichi di lavoro in Google Kubernetes Engine (GKE) utilizzando la configurazione Cloud TPU Multislice 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 si fa riferimento nei contenuti di Google Cloud , consulta Ruoli utente e attività comuni di GKE Enterprise.

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

1. Introduzione a Cloud TPU
2. Architettura di sistema di Cloud TPU
3. Informazioni sulle TPU in GKE

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 sezioni TPU a singolo host.
• Multislice supporta solo l'addestramento sincrono multicontroller.
• I carichi di lavoro multislice 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 inizializzala. Se hai già installato gcloud CLI, scarica l'ultima versione eseguendo gcloud components update.

• Crea un cluster Standard o un cluster Autopilot che esegue una versione che supporta Multislice. Per le versioni supportate, consulta Disponibilità di più slice 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:

1. Crea tre node pool TPU slice multi-host
2. 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. Puoi creare un pool di nodi slice TPU multi-host utilizzando Google Cloud CLI, Terraform o la console Google Cloud .

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONES \
    --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]

Verifica lo stato del pool di nodi

1. 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.

2. Utilizza kubectl in Cloud Shell per visualizzare i nodi dello 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 carico di lavoro JAX:

1. 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.10
                   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.10
                   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 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 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à pod per garantire che tutti i pod siano pianificati sulla stessa sezione.
   • maxRestarts: 4 indica il numero massimo di volte in cui GKE riavvia il JobSet quando un job secondario non riesce. Se i riavvii di JobSet raggiungono il massimo definito, JobSet viene contrassegnato come non riuscito.
   • I campi parallelism e completions corrispondono al numero di nodi in ogni pool di nodi.
   • backoff è 0 perché Multislice supporta solo l'addestramento sincrono di più 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.

2. Applica il manifest:

   kubectl apply -f tpu-multislice.yaml
   

3. Verifica che il workload sia ammesso:

   kubectl get jobsets
   

   L'output è simile al seguente:

   NAME            RESTARTS   COMPLETED   AGE
   multislice-job                         3s
   

4. 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.

5. (Facoltativo) Rimuovi il workload JAX:

   kubectl delete -f tpu-multislice.yaml
   

Configura le impostazioni aggiuntive

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

Migliorare le prestazioni di rete con hostNetwork

Per migliorare le prestazioni di rete tra gli slice TPU, ti consigliamo di attivare hostNetworking. Utilizza hostNetwork: true nella specifica del pod per ignorare tutto lo stack di rete 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) 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

Se utilizzi TPU Trillium 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 più reti 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:

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

3. 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
   

4. 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"},
         ]
   

5. 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
   

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 architettonica di uno slice TPU e di un Multislice:

Esegui la pulizia delle risorse

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

Elimina il progetto

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