Panoramica di Cloud TPU Multislice

Cloud TPU Multislice è una tecnologia di scalabilità delle prestazioni full stack che consente a un job di addestramento di utilizzare più slice TPU all'interno di un singolo pod o su slice in più pod con parallelismo dei dati standard. Con i chip TPU v4, i job di addestramento possono utilizzare più di 4096 chip in una singola esecuzione. Per i job di addestramento che richiedono meno di 4096 chip, una singola sezione può offrire le prestazioni migliori. Tuttavia, sono disponibili più facilmente più segmenti più piccoli, il che consente un tempo di avvio più rapido quando Multislice viene utilizzato con segmenti più piccoli.

Quando vengono implementati in configurazioni multislice, i chip TPU in ogni slice comunicano tramite Inter-Chip Interconnect (ICI). I chip TPU in slice diversi comunicano trasferendo i dati alle CPU (host), che a loro volta trasmettono i dati tramite la rete del data center (DCN). Per saperne di più sulla scalabilità con Multislice, consulta Come scalare l'addestramento dell'AI fino a decine di migliaia di chip Cloud TPU con Multislice.

Gli sviluppatori non devono scrivere codice per implementare la comunicazione DCN tra slice. Il compilatore XLA genera questo codice per te e sovrappone la comunicazione al calcolo per ottenere le massime prestazioni.

Concetti

Tipo di acceleratore

La forma di ogni slice TPU che compone un multislice. Ogni slice in una richiesta multislice è dello stesso tipo di acceleratore. Un tipo di acceleratore è costituito da un tipo di TPU (v4 o versioni successive) seguito dal numero di TensorCore. Ad esempio, v5litepod-128 specifica una TPU v5e con 128 TensorCore.

Riparazione automatica

Quando una slice riscontra un evento di manutenzione, un'interruzione o un guasto hardware, Cloud TPU creerà una nuova slice. Se non ci sono risorse sufficienti per creare una nuova sezione, la creazione non verrà completata finché l'hardware non sarà disponibile. Dopo la creazione della nuova sezione, tutte le altre sezioni nell'ambiente Multislice verranno riavviate in modo che l'addestramento possa continuare. Con uno script di avvio configurato correttamente, lo script di addestramento può essere riavviato automaticamente senza l'intervento dell'utente, caricando e riprendendo l'esecuzione dall'ultimo checkpoint.

Data Center Networking (DCN)

Una rete con latenza più elevata e throughput inferiore (rispetto a ICI) che connette le sezioni TPU in una configurazione multislice.

Pianificazione di gruppo

Quando tutte le sezioni TPU vengono sottoposte al provisioning insieme, contemporaneamente, garantendo che tutte o nessuna delle sezioni venga sottoposta al provisioning correttamente.

Interconnessione interchip (ICI)

Collegamenti interni ad alta velocità e bassa latenza che connettono le TPU all'interno di un pod TPU.

Multislice

Due o più sezioni di chip TPU che possono comunicare tramite DCN.

Nodo

Nel contesto di Multislice, il nodo si riferisce a una singola sezione TPU. A ogni slice TPU in un Multislice viene assegnato un ID nodo.

Script di avvio

Uno script di avvio di Compute Engine standard che viene eseguito ogni volta che una VM viene avviata o riavviata. Per Multislice, viene specificato nella richiesta di creazione del QR. Per saperne di più sugli script di avvio di Cloud TPU, consulta Gestire le risorse TPU.

Tensor

Una struttura di dati utilizzata per rappresentare dati multidimensionali in un modello di machine learning.

Tipi di capacità di Cloud TPU

Le TPU possono essere create da diversi tipi di capacità (vedi Opzioni di utilizzo in Come funziona il prezzo delle TPU):

• Prenotazione: per utilizzare una prenotazione, devi disporre di un contratto di prenotazione con Google. Utilizza il flag --reserved quando crei le risorse.

• Spot: ha come target la quota prerilasciabile utilizzando le VM spot. Le tue risorse potrebbero essere preempted per fare spazio alle richieste di un job con priorità più alta. Utilizza il flag --spot quando crei le risorse.

• On demand: ha come target la quota on demand, che non richiede una prenotazione e non verrà preempted. La richiesta TPU verrà messa in coda a una coda di quote on demand offerta da Cloud TPU, la disponibilità delle risorse non è garantita. Selezionato per impostazione predefinita, non sono necessari flag.

Inizia

1. Configura l'ambiente Cloud TPU.

2. Nella console Google Cloud , attiva Cloud Shell.

   Attiva Cloud Shell

   Nella parte inferiore della console Google Cloud viene avviata una sessione di Cloud Shell e viene visualizzato un prompt della riga di comando. Cloud Shell è un ambiente shell con Google Cloud CLI già installata e con valori già impostati per il progetto corrente. L'inizializzazione della sessione può richiedere alcuni secondi.

Per utilizzare il multislice, le risorse TPU devono essere gestite come risorse in coda.

Esempio introduttivo

Questo tutorial utilizza il codice del repository GitHub MaxText. MaxText è un LLM di base open source ad alte prestazioni, scalabile in modo arbitrario e ben testato scritto in Python e Jax. MaxText è stato progettato per l'addestramento efficiente su Cloud TPU.

Il codice in shardings.py è progettato per aiutarti a iniziare a sperimentare con diverse opzioni di parallelismo. Ad esempio, parallelismo dei dati, parallelismo dei dati completamente partizionati (FSDP) e parallelismo dei tensori. Il codice viene scalato da ambienti a singola sezione a ambienti multisezione.

Parallelismo ICI

ICI si riferisce all'interconnessione ad alta velocità che collega le TPU in una singola slice. Lo sharding ICI corrisponde allo sharding all'interno di una sezione. shardings.py fornisce tre parametri di parallelismo ICI:

• ici_data_parallelism
• ici_fsdp_parallelism
• ici_tensor_parallelism

I valori specificati per questi parametri determinano il numero di shard per ogni metodo di parallelismo.

Questi input devono essere vincolati in modo che ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism sia uguale al numero di chip nella sezione.

La seguente tabella mostra esempi di input utente per il parallelismo ICI per i quattro chip disponibili in v4-8:

Tieni presente che ici_data_parallelism deve essere lasciato come 1 nella maggior parte dei casi perché la rete ICI è abbastanza veloce da preferire quasi sempre FSDP al parallelismo dei dati.

Questo esempio presuppone che tu abbia familiarità con l'esecuzione di codice su una singola sezione TPU, come in Esegui un calcolo su una VM Cloud TPU utilizzando JAX. Questo esempio mostra come eseguire shardings.py su una singola sezione.

1. Configura l'ambiente:

   $ gcloud auth login
   $ export QR_ID=your-queued-resource-id
   $ export TPU_NAME=your-tpu-name
   $ export PROJECT=your-project-name
   $ export ZONE=us-central1-a
   $ export NETWORK_NAME=your-network-name
   $ export SUBNETWORK_NAME=your-subnetwork-name
   $ export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   $ export ACCELERATOR_TYPE=v5litepod-16
   $ export EXAMPLE_TAG_1=your-tag-1
   $ export EXAMPLE_TAG_2=your-tag-2
   $ export SLICE_COUNT=4
   $ export STARTUP_SCRIPT='#!/bin/bash\n'

   Descrizioni delle variabili

   • QR_ID: l'ID assegnato dall'utente della risorsa in coda.
   • TPU_NAME: il nome assegnato dall'utente della TPU.
   • PROJECT: Google Cloud nome del progetto
   • ZONE: specifica la zona in cui creare le risorse.
   • NETWORK_NAME: il nome della rete VPC.
   • SUBNETWORK_NAME: Nome della subnet nella rete VPC
   • RUNTIME_VERSION: la versione del software Cloud TPU.
   • ACCELERATOR_TYPE: il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare.
   • EXAMPLE_TAG_1, EXAMPLE_TAG_2 …: Tag utilizzati per identificare origini o destinazioni valide per i firewall di rete
   • SLICE_COUNT: Numero di settori. Limitato a un massimo di 256 segmenti.
   • STARTUP_SCRIPT: se specifichi uno script di avvio, lo script viene eseguito quando lo slice TPU viene sottoposto a provisioning o riavviato.

2. Crea chiavi SSH per gcloud. Ti consigliamo di lasciare una password vuota (premi Invio due volte dopo aver eseguito il seguente comando). Se ti viene comunicato che il file google_compute_engine esiste già, sostituisci la versione esistente.

   $ ssh-keygen -f ~/.ssh/google_compute_engine

3. Esegui il provisioning delle TPU:

   gcloud

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --node-id=${TPU_NAME} \
       --zone=${ZONE} \
       [--reserved |--spot]

   Google Cloud CLI non supporta tutte le opzioni di creazione di QR, ad esempio i tag. Per saperne di più, consulta Creare QR.

   Console

   1. Nella console Google Cloud , vai alla pagina TPU:

      Vai a TPU

   2. Fai clic su Crea TPU.

   3. Nel campo Nome, inserisci un nome per la TPU.

   4. Nella casella Zona, seleziona la zona in cui vuoi creare la TPU.

   5. Nella casella Tipo di TPU, seleziona un tipo di acceleratore. Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni di TPU.

   6. Nella casella Versione software TPU, seleziona una versione software. Quando crei una VM Cloud TPU, la versione del software TPU specifica la versione del runtime TPU da installare. Per maggiori informazioni, vedi Versioni software TPU.

   7. Fai clic sul pulsante di attivazione/disattivazione Attiva la messa in coda.

   8. Nel campo Nome risorsa in coda, inserisci un nome per la richiesta di risorsa in coda.

   9. Fai clic su Crea per creare la richiesta di risorse in coda.

4. Attendi che la risorsa in coda sia nello stato ACTIVE, il che significa che i nodi worker sono nello stato READY. Una volta avviato il provisioning della risorsa in coda, il completamento potrebbe richiedere da 1 a 5 minuti, a seconda delle dimensioni della risorsa in coda. Puoi controllare lo stato di una richiesta di risorse in coda utilizzando gcloud CLI o la console Google Cloud :

   gcloud

   $ gcloud compute tpus queued-resources \
       list --filter=${QR_ID} --zone=${ZONE}

   Console

   1. Nella console Google Cloud , vai alla pagina TPU:

      Vai a TPU

   2. Fai clic sulla scheda Risorse in coda.

   3. Fai clic sul nome della richiesta di risorsa in coda.

5. Connettiti alla VM TPU utilizzando SSH:

   $ gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE}

6. Clona MaxText (che include shardings.py) nella tua VM TPU:

   $ git clone https://github.com/AI-Hypercomputer/maxtext && cd maxtext

7. Installa Python 3.10:

   $ sudo apt-get update
   $ sudo apt install python3.10
   $ sudo apt install python3.10-venv

8. Crea e attiva un ambiente virtuale:

   $ python3 -m venv your-venv-name
   $ source your-venv-name/bin/activate

9. Nella directory del repository MaxText, esegui lo script di configurazione per installare JAX e altre dipendenze sullo slice TPU. L'esecuzione dello script di configurazione richiede alcuni minuti.

   $ bash setup.sh

10. Esegui questo comando per eseguire shardings.py sullo slice TPU.

    $ python3 -m pedagogical_examples.shardings \
      --ici_fsdp_parallelism 4 \
      --batch_size 131072 \
      --embedding_dimension 2048

    Puoi visualizzare i risultati nei log. Le TPU dovrebbero raggiungere circa 260 TFLOP al secondo o un utilizzo FLOP impressionante superiore al 90%. In questo caso, abbiamo selezionato approssimativamente il batch massimo che rientra nella memoria ad alta larghezza di banda (HBM) della TPU.

11. Non esitare a esplorare altre strategie di sharding su ICI, ad esempio puoi provare la seguente combinazione:

    $ python3 -m pedagogical_examples.shardings \
      --ici_tensor_parallelism 4 \
      --batch_size 131072 \
      --embedding_dimension 2048

12. Al termine, elimina la risorsa in coda e lo slice TPU. Devi eseguire questi passaggi di pulizia dall'ambiente in cui hai configurato lo slice (prima esegui exit per uscire dalla sessione SSH). L'eliminazione richiede da 2 a 5 minuti. Se utilizzi gcloud CLI, puoi eseguire questo comando in background con il flag facoltativo --async.

    gcloud

    $ gcloud compute tpus queued-resources \
        delete ${QR_ID} --force (--async)

    Console

    1. Nella console Google Cloud , vai alla pagina TPU:

       Vai a TPU

    2. Fai clic sulla scheda Risorse in coda.

    3. Seleziona la casella di controllo accanto alla richiesta di risorse in coda.

    4. Fai clic su delete Elimina.

Sharding multislice utilizzando il parallelismo DCN

Lo script shardings.py accetta tre parametri che specificano il parallelismo DCN, corrispondenti al numero di shard di ogni tipo di parallelismo dei dati:

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

I valori di questi parametri devono essere vincolati in modo che dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism sia uguale al numero di segmenti.

Ad esempio, per due fette, utilizza --dcn_data_parallelism = 2.

dcn_tensor_parallelism deve essere sempre impostato su 1 perché il DCN non è adatto a questo tipo di sharding. Per i tipici workload LLM sui chip v4, anche dcn_fsdp_parallelism deve essere impostato su 1 e pertanto dcn_data_parallelism deve essere impostato sul numero di sezioni, ma questo dipende dall'applicazione.

Se aumenti il numero di segmenti (supponendo di mantenere costanti le dimensioni del segmento e il batch per segmento), aumenti il parallelismo dei dati.

Esecuzione di shardings.py in un ambiente Multislice

Puoi eseguire shardings.py in un ambiente multislice utilizzando multihost_runner.py o eseguendo shardings.py su ogni VM TPU. Qui utilizziamo multihost_runner.py. I passaggi seguenti sono molto simili a quelli descritti in Getting Started: Quick Experiments on Multiple slices del repository MaxText, tranne per il fatto che qui eseguiamo shardings.py anziché l'LLM più complesso in train.py.

Lo strumento multihost_runner.py è ottimizzato per esperimenti rapidi, riutilizzando ripetutamente le stesse TPU. Poiché lo script multihost_runner.py dipende da connessioni SSH di lunga durata, non lo consigliamo per job di lunga durata. Se vuoi eseguire un job più lungo (ad esempio ore o giorni), ti consigliamo di utilizzare multihost_job.py.

In questo tutorial, utilizziamo il termine runner per indicare la macchina su cui esegui lo script multihost_runner.py. Utilizziamo il termine worker per indicare le VM TPU che compongono le tue sezioni. Puoi eseguire multihost_runner.py su una macchina locale o su qualsiasi VM di Compute Engine nello stesso progetto delle tue sezioni. L'esecuzione di multihost_runner.py su un worker non è supportata.

multihost_runner.py si connette automaticamente ai worker TPU tramite SSH.

In questo esempio, esegui shardings.py su due slice v5e-16, per un totale di quattro VM e 16 chip TPU. Puoi modificare l'esempio per eseguirlo su più TPU.

Configura l'ambiente

1. Clona MaxText sulla macchina runner:

   $ git clone https://github.com/AI-Hypercomputer/maxtext

2. Vai alla directory del repository.

   $ cd maxtext

3. Crea chiavi SSH per gcloud. Ti consigliamo di lasciare una password vuota (premi INVIO due volte dopo aver eseguito il seguente comando). Se ti viene comunicato che il file google_compute_engine esiste già, scegli di non mantenere la versione esistente.

     $ ssh-keygen -f ~/.ssh/google_compute_engine
     

4. Aggiungi una variabile di ambiente per impostare il conteggio delle sezioni TPU su 2.

     $ export SLICE_COUNT=2
     

5. Crea un ambiente Multislice utilizzando il comando queued-resources create o la console Google Cloud .

   gcloud

   Il seguente comando mostra come creare una TPU multislice v5e. Per utilizzare una versione TPU diversa, specifica un accelerator-type e un runtime-version diversi.

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --node-count=${SLICE_COUNT} \
       --node-prefix=${TPU_NAME} \
       --zone=${ZONE} \
       [--reserved|--spot]

   Console

   1. Nella console Google Cloud , vai alla pagina TPU:

      Vai a TPU

   2. Fai clic su Crea TPU.

   3. Nel campo Nome, inserisci un nome per la TPU.

   4. Nella casella Zona, seleziona la zona in cui vuoi creare la TPU.

   5. Nella casella Tipo di TPU, seleziona un tipo di acceleratore. Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Multislice è supportato solo su Cloud TPU v4 e versioni successive di TPU. Per maggiori informazioni sulle versioni TPU, consulta la sezione Versioni TPU.

   6. Nella casella Versione software TPU, seleziona una versione software. Quando crei una VM Cloud TPU, la versione software TPU specifica la versione del runtime TPU da installare sulle VM TPU. Per maggiori informazioni, vedi Versioni software TPU.

   7. Fai clic sul pulsante di attivazione/disattivazione Attiva la messa in coda.

   8. Nel campo Nome risorsa in coda, inserisci un nome per la richiesta di risorsa in coda.

   9. Fai clic sulla casella di controllo Rendi questa TPU multislice.

   10. Nel campo Conteggio slice, inserisci il numero di slice che vuoi creare.

   11. Fai clic su Crea per creare la richiesta di risorse in coda.

6. Quando inizia il provisioning delle risorse in coda, il completamento può richiedere fino a cinque minuti, a seconda delle dimensioni della risorsa in coda. Attendi che la risorsa in coda sia nello stato ACTIVE. Puoi controllare lo stato di una richiesta di risorse in coda utilizzando gcloud CLI o la console Google Cloud :

   gcloud

   $ gcloud compute tpus queued-resources list \
       --filter=${QR_ID} --zone=${ZONE} --project=${PROJECT}

   In questo modo dovrebbe essere generato un output simile al seguente:

   NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
   ...
   que-res-id  us-central2-b  4           v5litepod-16             ACTIVE
   ...

   Console

   1. Nella console Google Cloud , vai alla pagina TPU:

      Vai a TPU

   2. Fai clic sulla scheda Risorse in coda.

   3. Fai clic sul nome della richiesta di risorsa in coda.

   Contatta il rappresentante del tuo account Google Cloud se lo stato del QR è WAITING_FOR_RESOURCES o PROVISIONING per più di 15 minuti.

7. Installa le dipendenze.

   $ python3 multihost_runner.py \
       --TPU_PREFIX=${TPU_NAME} \
       --ZONE=${ZONE} \
       --COMMAND="bash setup.sh"

8. Esegui shardings.py su ogni worker utilizzando multihost_runner.py.

   $ python3 multihost_runner.py \
       --TPU_PREFIX=${TPU_NAME} \
       --ZONE=${ZONE} \
       --COMMAND="python3 -m pedagogical_examples.shardings \
       --dcn_data_parallelism ${SLICE_COUNT} \
       --ici_fsdp_parallelism 16 \
       --batch_size 131072 \
       --embedding_dimension 2048"

   Nei file di log vedrai un rendimento di circa 230 TFLOP al secondo.

   Per ulteriori informazioni sulla configurazione del parallelismo, consulta Sharding multislice utilizzando il parallelismo DCN e shardings.py.

9. Al termine, libera spazio nelle TPU e nella risorsa in coda. L'eliminazione richiede da due a cinque minuti. Se utilizzi gcloud CLI, puoi eseguire questo comando in background con il flag facoltativo --async.

Scalare un carico di lavoro a Multislice

Prima di eseguire il modello in un ambiente Multislice, apporta le seguenti modifiche al codice:

• Utilizza jax.experimental.mesh_utils.create_hybrid_device_mesh anziché jax.experimental.mesh_utils.create_device_mesh quando crei la mesh.

Queste dovrebbero essere le uniche modifiche al codice necessarie quando si passa a Multislice. Per ottenere prestazioni elevate, DCN deve essere mappato su assi paralleli ai dati, paralleli ai dati completamente partizionati o paralleli alla pipeline. Le considerazioni sul rendimento e le strategie di sharding sono trattate in modo più dettagliato in Sharding con multislice per il massimo rendimento.

Per verificare che il codice possa accedere a tutti i dispositivi, puoi affermare che len(jax.devices()) è uguale al numero di chip nel tuo ambiente Multislice. Ad esempio, se utilizzi quattro fette di v4-16, hai otto chip per fetta * 4 fette, quindi len(jax.devices()) dovrebbe restituire 32.

Scegliere le dimensioni delle sezioni per gli ambienti Multislice

Per ottenere un aumento lineare della velocità, aggiungi nuove fette delle stesse dimensioni di quelle esistenti. Ad esempio, se utilizzi una sezione v4-512, Multislice otterrà prestazioni circa doppie aggiungendo una seconda sezione v4-512 e raddoppiando le dimensioni del batch globale. Per saperne di più, consulta Sharding con multislice per il massimo rendimento.

Esecuzione del job su più sezioni

Esistono tre approcci diversi per eseguire il workload personalizzato in un ambiente multislice:

1. Utilizzando lo script di esecuzione dell'esperimento, multihost_runner.py
2. Utilizzando lo script del runner di produzione, multihost_job.py
3. Utilizzare un approccio manuale

Script di esecuzione dell'esperimento

Lo script multihost_runner.py distribuisce il codice a un ambiente Multislice esistente ed esegue il comando su ogni host, copia i log e tiene traccia dello stato di errore di ogni comando. Lo script multihost_runner.py è documentato nel README di MaxText.

Poiché multihost_runner.py mantiene connessioni SSH persistenti, è adatto solo per esperimenti di dimensioni modeste e di durata relativamente breve. Puoi adattare i passaggi del tutorial multihost_runner.py alla tua configurazione di workload e hardware.

Script di esecuzione della produzione

Per i job di produzione che richiedono resilienza contro guasti hardware e altri interruzioni, è consigliabile l'integrazione diretta con l'API Create Queued Resource. Utilizza multihost_job.py come esempio pratico che attiva la chiamata API Created Queued Resource con lo script di avvio appropriato per eseguire l'addestramento e riprenderlo in caso di preemptive. Lo script multihost_job.py è documentato nel file README di MaxText.

Poiché multihost_job.py deve eseguire il provisioning delle risorse per ogni esecuzione, non fornisce un ciclo di iterazione veloce come multihost_runner.py.

Approccio manuale

Ti consigliamo di utilizzare o adattare multihost_runner.py o multihost_job.py per eseguire il tuo workload personalizzato nella configurazione Multislice. Tuttavia, se preferisci eseguire il provisioning e gestire l'ambiente utilizzando direttamente i comandi QR, consulta Gestire un ambiente multislice.

Gestire un ambiente Multislice

Per eseguire il provisioning e gestire manualmente i QR senza utilizzare gli strumenti forniti nel repository MaxText, leggi le sezioni seguenti.

Crea risorse in coda

{
"name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qr-guid>",
"metadata": {
    "@type": "type.googleapis.com/google.cloud.common.OperationMetadata",
    "createTime": "2023-11-01T00:17:05.742546311Z",
    "target": "projects/<your-project-id>/locations/<your-zone>/queuedResources/<your-qa-id>",
    "verb": "create",
    "cancelRequested": false,
    "apiVersion": "v2alpha1"
},
"done": false
}