Panoramica di Cloud TPU Multislice

Attenzione: questa pagina si applica all'API Cloud TPU, che non è più in fase di sviluppo attivo. L'API Cloud TPU riceverà solo correzioni di bug e aggiornamenti della sicurezza. Ti consigliamo di utilizzare Google Kubernetes Engine per Multislice. Per saperne di più, consulta Esegui il deployment di TPU multislice in GKE.

Importante: il multislice è supportato solo su Cloud TPU v4 e versioni TPU successive.

Cloud TPU Multislice è una tecnologia di scalabilità delle prestazioni full stack che consente a un job di addestramento di utilizzare più sezioni TPU all'interno di un singolo pod o su sezioni in più pod con parallelismo dei dati standard. Con i chip TPU v4, ciò significa che 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ù fette più piccole, il che consente un tempo di avvio più rapido quando Multislice viene utilizzato con fette più piccole.

Se viene eseguito il deployment in configurazioni multislice, i chip TPU in ogni sezione comunicano tramite inter-chip-interconnect (ICI). I chip TPU in sezioni diverse 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 sull'AI fino a decine di migliaia di chip Cloud TPU con Multislice.

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

Concetti

Tipo di acceleratore

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

Riparazione automatica

Quando uno slice riscontra un evento di manutenzione, un prerilascio o un guasto hardware, Cloud TPU crea un nuovo slice. Se le risorse non sono sufficienti per creare una nuova sezione, la creazione non verrà completata finché non sarà disponibile l'hardware. 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 avvioconfigurato correttamente, lo script di addestramento può essere riavviato automaticamente senza l'intervento dell'utente, caricando e riprendendo 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 a provisioning insieme, contemporaneamente, garantendo che tutte o nessuna delle sezioni vengano sottoposte a provisioning correttamente.

Interchip interconnect (ICI)

Link interni ad alta velocità e bassa latenza che collegano 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 sezione TPU in un ambiente 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 codice QR. Per saperne di più sugli script di avvio di Cloud TPU, consulta Gestisci 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 avere un contratto di prenotazione con Google. Utilizza il flag --reserved quando crei le risorse.

• Spot: ha come target la quota preemptible utilizzando le VM spot. Le tue risorse potrebbero essere prerilasciate 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à prerilasciata. La richiesta di TPU verrà inserita in 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 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 connette 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 nello slice.

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

	ici_data_parallelism	ici_fsdp_parallelism	ici_tensor_parallelism
FSDP a 4 vie	1	4	1
Parallelismo tensore a 4 vie	1	1	4
FSDP a 2 vie + parallelismo dei tensori a 2 vie	1	2	2

Tieni presente che ici_data_parallelism deve essere lasciato su 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 fetta TPU, ad esempio 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 software di 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 slice. Limitato a un massimo di 256 segmenti.
   • STARTUP_SCRIPT: se specifichi uno script di avvio, lo script viene eseguito quando viene eseguito il provisioning o il riavvio della sezione TPU.

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 saperne di più 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 software della TPU specifica la versione del runtime TPU da installare. Per saperne di più, consulta Versioni software TPU.

   7. Fai clic sul pulsante di attivazione/disattivazione Attiva 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 finché la risorsa in coda non si trova nello stato ACTIVE, il che significa che i nodi worker sono nello stato READY. Una volta avviato il provisioning delle risorse in coda, il completamento potrebbe richiedere da uno a cinque 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 risorse in coda.

5. Connettiti alla VM TPU tramite 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:

   Nota: la creazione di un ambiente virtuale per ogni progetto è una best practice per la gestione delle dipendenze Python. Per saperne di più, consulta la documentazione di venv.

   $ 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 sulla tua sezione TPU. L'esecuzione dello script di configurazione richiede alcuni minuti.

   $ bash setup.sh

10. Esegui questo comando per eseguire shardings.py sulla tua slice TPU.

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

    Puoi visualizzare i risultati nei log. Le tue TPU dovrebbero raggiungere circa 260 TFLOP al secondo o un utilizzo di FLOP 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. Puoi 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 la sezione TPU. Devi eseguire questi passaggi di pulizia dall'ambiente in cui hai configurato lo slice (prima esegui exit per uscire dalla sessione SSH). Il completamento dell'eliminazione richiederà da due a cinque minuti. Se utilizzi gcloud CLI, puoi eseguire questo comando in background con il flag facoltativo --async.

    Importante: le risorse in coda devono essere eliminate dopo l'utilizzo. In caso contrario, continueranno a consumare la quota di risorse in coda, che viene monitorata separatamente dalla quota di TPU.

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

Sharding multislice utilizzando il parallelismo DCN

Lo script shardings.py accetta tre parametri che specificano il parallelismo DCN, corrispondente 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 sezioni.

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

	dcn_data_parallelism	dcn_fsdp_parallelism	dcn_tensor_parallelism	N. di sezioni
Parallelismo dei dati bidirezionale	2	1	1	2

dcn_tensor_parallelism deve essere sempre impostato su 1 perché la DCN non è adatta a questo tipo di sharding. Per i workload LLM tipici sui chip v4, dcn_fsdp_parallelism deve essere impostato anche su 1, pertanto dcn_data_parallelism deve essere impostato sul numero di slice, 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. In questo caso 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é il modello linguistico di grandi dimensioni 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 i job a esecuzione prolungata. 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 sezioni. Puoi eseguire multihost_runner.py su una macchina locale o su qualsiasi VM di Compute Engine nello stesso progetto degli slice. L'esecuzione di multihost_runner.py su un worker non è supportata.

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

In questo esempio, esegui shardings.py su due sezioni 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 tua 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à, seleziona l'opzione per 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

   Nota :assicurati di disporre della quota corrispondente prima di selezionare --reserved, --spot o la quota on demand predefinita. Per informazioni sui tipi di quote, consulta le norme sulle quote.

   Il seguente comando mostra come creare una TPU v5e multi-slice. Per utilizzare una versione diversa della TPU, 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. La tecnologia multislice è supportata solo su Cloud TPU v4 e versioni successive di TPU. Per saperne di più 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 delle TPU.

   7. Fai clic sul pulsante di attivazione/disattivazione Abilita 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 Numero di sezioni, inserisci il numero di sezioni che vuoi creare.

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

6. Una volta avviato, il provisioning della risorsa in coda può richiedere fino a cinque minuti per essere completato, 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 risorse in coda.

   Contatta il tuo Google Cloud rappresentante dell'account 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 saperne di più sulla configurazione del parallelismo, consulta Sharding multislice utilizzando il parallelismo DCN e shardings.py.

9. Al termine, esegui la pulizia delle TPU e della risorsa in coda. L'eliminazione richiederà da due a cinque minuti. Se utilizzi gcloud CLI, puoi eseguire questo comando in background con il flag facoltativo --async.

Scalabilità di un workload in 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 per passare a Multislice. Per ottenere prestazioni elevate, la DCN deve essere mappata su assi di parallelismo dei dati, parallelismo dei dati con partizionamento completo o parallelismo della pipeline. Le considerazioni sul rendimento e le strategie di partizionamento sono descritte in maggiore dettaglio in Partizionamento con Multislice per il massimo rendimento.

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

Scelta delle dimensioni delle sezioni per gli ambienti multisezione

Per ottenere un'accelerazione lineare, aggiungi nuove sezioni della stessa dimensione di quella esistente. Ad esempio, se utilizzi una sezione v4-512, Multislice otterrà un rendimento circa doppio 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 carico di lavoro personalizzato in un ambiente multislice:

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

Script di esecuzione della sperimentazione

Lo script multihost_runner.py distribuisce il codice in un ambiente Multislice esistente, esegue il comando su ogni host, copia i log e monitora lo stato di errore di ogni comando. Lo script multihost_runner.py è documentato in MaxText README.

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 configurazione del tuo workload e del tuo hardware.

Script di esecuzione della produzione

Per i job di produzione che richiedono resilienza contro i guasti hardware e altri prerilasci, è meglio eseguire l'integrazione direttamente con l'API Create Queued Resource. Utilizza multihost_job.py come esempio di lavoro che attiva la chiamata API Created Queued Resource con lo script di avvio appropriato per eseguire l'addestramento e riprendere in caso di prerilascio. Lo script multihost_job.py è documentato nel 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 Gestisci un ambiente multi-slice.

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

gcloud

1. Crea una richiesta di risorsa in coda utilizzando il seguente comando:

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --project=${PROJECT} \
       --zone=${ZONE} \
       --node-count=${SLICE_COUNT} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --network=${NETWORK_NAME} \
       --subnetwork=${SUBNETWORK_NAME} \
       --tags=${EXAMPLE_TAG_1},${EXAMPLE_TAG_2} \
       --metadata=startup-script="${STARTUP_SCRIPT}" \
       [--reserved|--spot]

Assicurati di disporre della quota corrispondente prima di selezionare --reserved, --spot o la quota on demand predefinita. Per informazioni sui tipi di quote, consulta Policy sulle quote.

curl

1. Crea un file denominato queued-resource-req.json e copia al suo interno il seguente JSON.

   {
   "guaranteed": { "reserved": true },
   "tpu": {
       "node_spec": [
       {
       "parent": "projects/your-project-number/locations/your-zone",
           "node": {
           "accelerator_type": "accelerator-type",
           "runtime_version": "tpu-vm-runtime-version",
           "network_config": {
               "network": "your-network-name",
               "subnetwork": "your-subnetwork-name",
               "enable_external_ips": true
           },
           "tags" : ["example-tag-1"]
           "metadata": {
               "startup-script": "your-startup-script"
           }
       },
       "multi_node_params": {
           "node_count": slice-count,
           "node_id_prefix": "your-queued-resource-id"
       }
       }
       ]
   }
   }

   Sostituisci i seguenti valori:

   • your-project-number: il tuo numero di progetto Google Cloud
   • your-zone: la zona in cui vuoi creare la risorsa in coda
   • accelerator-type: la versione e le dimensioni di una singola sezione. Multislice è supportato solo su Cloud TPU v4 e versioni successive di TPU.
   • tpu-vm-runtime-version: la versione di runtime della VM TPU che vuoi utilizzare.
   • your-network-name: (facoltativo) una rete a cui verrà collegata la risorsa in coda
   • your-subnetwork-name: (facoltativo) una subnet a cui verrà collegata la risorsa in coda
   • example-tag-1: (facoltativo) una stringa di tag arbitraria
   • your-startup-script: uno script di avvio che verrà eseguito quando la risorsa in coda viene allocata
   • slice-count: il numero di sezioni TPU nel tuo ambiente multislice
   • your-queued-resource-id: L'ID fornito dall'utente per la risorsa in coda

   Per saperne di più, consulta la documentazione dell'API REST Queued Resource per tutte le opzioni disponibili.

   Nota: anche se node_spec è un elenco, è supportata una sola voce node_spec.

   Per utilizzare la capacità Spot, sostituisci:

   "guaranteed": { "reserved": true } con "spot": {}

   Rimuovi la riga per utilizzare la capacità on demand predefinita.

2. Invia la richiesta di creazione della risorsa in coda con il payload JSON:

   $ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" \
   -H "Content-Type: application/json" \
   -d @queuedresourcereq.json \
   https://tpu.googleapis.com/v2alpha1/projects/your-project-id/locations/your-zone/queuedResources\?queued_resource_id\=your-queued-resource-id

   Sostituisci i seguenti valori:

   • your-project-id: l'ID progetto Google Cloud .
   • your-zone: la zona in cui vuoi creare la risorsa in coda
   • your-queued-resource-id: L'ID fornito dall'utente per la risorsa in coda

La risposta dovrebbe essere simile alla seguente:

{
"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
}

Utilizza il valore GUID alla fine del valore della stringa per l'attributo name per ottenere informazioni sulla richiesta di risorse in coda.

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 TPU successive. 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 software della TPU specifica la versione del runtime TPU da installare. Per saperne di più, consulta Versioni software TPU.

7. Fai clic sul pulsante di attivazione/disattivazione Attiva 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 Numero di sezioni, inserisci il numero di sezioni che vuoi creare.

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

Recupera lo stato di una risorsa in coda

gcloud

$ gcloud compute tpus queued-resources describe ${QR_ID} --zone=${ZONE}

Per una risorsa in coda che si trova nello stato ACTIVE, l'output è simile al seguente:

...
state:
    state: ACTIVE
...

curl

$ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2/projects/your-project-id/locations/your-zone/queuedResources/${YOUR_QR_ID}

Per una risorsa in coda che si trova nello stato ACTIVE, l'output è simile al seguente:

{
"name": your-queued-res,
"tpu": {
    "nodeSpec": [
    {
        ... // node 1
    },
    {
        ... // node 2
    },
    ...
    ]
},
...
"state": "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.

Una volta eseguito il provisioning della TPU, puoi anche visualizzare i dettagli della richiesta di risorse in coda andando alla pagina TPU, trovando la tua TPU e facendo clic sul nome della richiesta di risorse in coda corrispondente.

In uno scenario raro, potresti trovare la risorsa in coda nello stato FAILED mentre alcune sezioni sono ACTIVE. In questo caso, elimina le risorse create e riprova tra qualche minuto o contatta l'assistenzaGoogle Cloud .

SSH e installazione delle dipendenze

Esegui il codice JAX sulle sezioni TPU descrive come connetterti alle tue VM TPU utilizzando SSH in una singola sezione. Per connetterti a tutte le VM TPU nel tuo ambiente Multislice tramite SSH e installare le dipendenze, utilizza il seguente comando gcloud:

  $ gcloud compute tpus queued-resources ssh ${QR_ID} \
        --zone=${ZONE} \
        --node=all \
        --worker=all \
        --command="command-to-run" \
        --batch-size=4

Questo comando gcloud invia il comando specificato a tutti i worker e i nodi in QR utilizzando SSH. Il comando viene raggruppato in gruppi di quattro e inviato contemporaneamente. Il gruppo di comandi successivo viene inviato al termine dell'esecuzione del gruppo corrente. Se si verifica un errore con uno dei comandi, l'elaborazione si interrompe e non vengono inviati altri batch. Per saperne di più, consulta il riferimento API delle risorse in coda. Se il numero di sezioni che utilizzi supera il limite di threading del tuo computer locale (chiamato anche limite di raggruppamento in batch), si verificherà un deadlock. Ad esempio, supponiamo che il limite di batch sulla tua macchina locale sia 64. Se provi a eseguire uno script di addestramento su più di 64 slice, ad esempio 100, il comando SSH suddividerà le slice in batch. Eseguirà lo script di addestramento sul primo batch di 64 slice e attenderà il completamento degli script prima di eseguire lo script sul batch rimanente di 36 slice. Tuttavia, il primo gruppo di 64 sezioni non può essere completato finché le restanti 36 sezioni non iniziano a eseguire lo script, causando un deadlock.

Per evitare questo scenario, puoi eseguire lo script di addestramento in background su ogni VM aggiungendo una e commerciale (&) al comando dello script specificato con il flag --command. In questo modo, dopo aver avviato lo script di addestramento sul primo batch di sezioni, il controllo tornerà immediatamente al comando SSH. Il comando SSH può quindi iniziare a eseguire lo script di addestramento sul batch rimanente di 36 sezioni. Dovrai collegare correttamente gli stream stdout e stderr quando esegui i comandi in background. Per aumentare il parallelismo all'interno dello stesso QR, puoi selezionare segmenti specifici utilizzando il parametro --node.

Impostazione rete

Assicurati che le sezioni TPU possano comunicare tra loro eseguendo i seguenti passaggi. Installa JAX su ogni slice. Per saperne di più, vedi Eseguire il codice JAX sugli slice TPU. Verifica che len(jax.devices()) sia uguale al numero di chip nel tuo ambiente multislice. Per farlo, esegui questo comando su ogni sezione:

  $ python3 -c 'import jax; print(jax.devices())'

Se esegui questo codice su quattro sezioni di v4-16, ci sono otto chip per sezione e quattro sezioni, per un totale di 32 chip (dispositivi) che dovrebbero essere restituiti da jax.devices().

Elenca risorse in coda

gcloud

Puoi visualizzare lo stato delle risorse in coda utilizzando il comando queued-resources list:

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

L'output è simile al seguente:

NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
...
que-res-id  us-central1-a  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.

Avvia il job in un ambiente di cui è stato eseguito il provisioning

Puoi eseguire manualmente i workload connettendoti a tutti gli host in ogni slice tramite SSH ed eseguendo il seguente comando su tutti gli host.

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone=${ZONE} \
    --worker=all \
    --command="command-to-run"

Reimpostazione dei codici QR

L'API ResetQueuedResource può essere utilizzata per reimpostare tutte le VM in un QR ACTIVE. La reimpostazione delle VM cancella forzatamente la memoria della macchina e reimposta la VM allo stato iniziale. Tutti i dati archiviati localmente rimarranno intatti e lo script di avvio verrà richiamato dopo un ripristino. L'API ResetQueuedResource può essere utile quando vuoi riavviare tutte le TPU. Ad esempio, quando l'addestramento è bloccato e il ripristino di tutte le VM è più semplice del debug.

I ripristini di tutte le VM vengono eseguiti in parallelo e un'operazione ResetQueuedResource richiede da uno a due minuti. Per richiamare l'API, utilizza il seguente comando:

$ gcloud compute tpus queued-resources reset ${QR_ID} --zone=${ZONE}

Eliminazione delle risorse in coda

Per rilasciare le risorse al termine della sessione di addestramento, elimina la risorsa in coda. Il completamento dell'eliminazione richiede da due a cinque 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} --zone=${ZONE} --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 Elimina.

Recupero automatico dagli errori

In caso di interruzione, Multislice offre la riparazione senza intervento della slice interessata e il ripristino di tutte le slice in seguito. La sezione interessata viene sostituita con una nuova sezione e le restanti sezioni, altrimenti integre, vengono reimpostate. Se non è disponibile capacità per allocare una sezione sostitutiva, l'addestramento viene interrotto.

Per riprendere l'addestramento automaticamente dopo un'interruzione, devi specificare uno script di avvio che controlli e carichi gli ultimi checkpoint salvati. Lo script di avvio viene eseguito automaticamente ogni volta che una sezione viene riassegnata o una VM viene reimpostata. Specifica uno script di avvio nel payload JSON che invii all'API di richiesta di creazione QR.

Il seguente script di avvio (utilizzato in Crea QR) ti consente di eseguire automaticamente il recupero dagli errori e riprendere l'addestramento dai checkpoint archiviati in un bucket Cloud Storage durante l'addestramento MaxText:

{
 "tpu": {
   "node_spec": [
     {
      ...
         "metadata": {
               "startup-script": "#! /bin/bash \n pwd \n runuser -l user1 -c 'cd /home/user1/MaxText && python3 -m MaxText.train MaxText/configs/base.yml run_name=run_test_failure_recovery dcn_data_parallelism=4 ici_fsdp_parallelism=8 steps=10000 save_period=10 base_output_directory='gs://user1-us-central2'' EOF"
         }
     ...
     }
   ]
 }
}

Clona il repository MaxText prima di provare questa operazione.

Profilazione e debug

La profilazione è la stessa negli ambienti a sezione singola e multisezione. Per saperne di più, consulta Profilazione dei programmi JAX.

Ottimizza la formazione

Le sezioni seguenti descrivono come ottimizzare l'addestramento Multislice.

Sharding con Multislice per il massimo rendimento

Per ottenere le massime prestazioni negli ambienti multislice, è necessario considerare come eseguire lo sharding tra i vari slice. In genere esistono tre opzioni (parallelismo dei dati, parallelismo dei dati con partizionamento completo e parallelismo della pipeline). Non consigliamo di partizionare le attivazioni tra le dimensioni del modello (a volte chiamato parallelismo dei tensori) perché richiede troppa larghezza di banda tra le sezioni. Per tutte queste strategie, puoi mantenere la stessa strategia di partizionamento all'interno di una sezione che ha funzionato in passato.

Ti consigliamo di iniziare con il parallelismo dei dati puro. L'utilizzo del parallelismo dei dati completamente partizionati è utile per liberare la memoria utilizzata. Lo svantaggio è che la comunicazione tra le sezioni utilizza la rete DCN e rallenta il workload. Utilizza il parallelismo della pipeline solo quando necessario in base alla dimensione del batch (come analizzato di seguito).

Quando utilizzare il parallelismo dei dati

Il parallelismo dei dati puro funziona bene nei casi in cui hai un workload che funziona bene, ma vuoi migliorarne le prestazioni eseguendo lo scale up su più sezioni.

Per ottenere uno scaling efficace su più slice, il tempo necessario per eseguire all-reduce su DCN deve essere inferiore al tempo necessario per eseguire un passaggio all'indietro. DCN viene utilizzato per la comunicazione tra le sezioni ed è un fattore limitante nella velocità effettiva del workload.

Ogni chip TPU v4 ha un picco di 275 * 10¹² FLOPS al secondo.

Ogni host TPU ha quattro chip e una larghezza di banda di rete massima di 50 Gbps.

Ciò significa che l'intensità aritmetica è 4 * 275 * 10¹² FLOPS / 50 Gbps = 22.000 FLOPS / bit.

Il modello utilizzerà da 32 a 64 bit di larghezza di banda DCN per ogni parametro per passaggio. Se utilizzi due sezioni, il modello utilizzerà 32 bit di larghezza di banda DCN. Se utilizzi più di due sezioni, il compilatore eseguirà un'operazione di riduzione completa e casuale e utilizzerai fino a 64 bit di larghezza di banda DCN per ogni parametro per passaggio. La quantità di FLOPS necessaria per ogni parametro varia a seconda del modello. Nello specifico, per i modelli linguistici basati su Transformer, il numero di FLOPS richiesti per un passaggio in avanti e uno indietro è circa 6 * B * P, dove:

• B è la dimensione del batch in token
• P è il numero di parametri

Il numero di FLOPS per parametro è 6 * B e il numero di FLOPS per parametro durante il backward pass è 4 * B.

Per garantire una scalabilità efficace su più sezioni, assicurati che l'intensità operativa superi l'intensità aritmetica dell'hardware TPU. Per calcolare l'intensità operativa, dividi il numero di FLOPS per parametro durante il backward pass per la larghezza di banda di rete (in bit) per parametro per passaggio: Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

Pertanto, per un modello linguistico basato su Transformer, se utilizzi due sezioni: Operational intensity = 4 * B / 32

Se utilizzi più di due sezioni: Operational intensity = 4 * B/64

Ciò suggerisce una dimensione minima del batch compresa tra 176.000 e 352.000 per i modelli linguistici basati su Transformer. Poiché la rete DCN può perdere brevemente i pacchetti, è meglio mantenere un margine di errore significativo, eseguendo il deployment del parallelismo dei dati solo se la dimensione del batch per pod è di almeno 350.000 (due pod) o 700.000 (molti pod).

Per altre architetture di modelli, dovrai stimare il runtime del passaggio all'indietro per sezione (misurando il tempo con un profiler o contando i FLOPS). Puoi quindi confrontare questo valore con il tempo di esecuzione previsto per l'all-reduce su DCN e ottenere una buona stima per capire se il parallelismo dei dati è adatto alle tue esigenze.

Quando utilizzare il parallelismo dei dati completamente partizionati (FSDP)

Il parallelismo dei dati completamente partizionati (FSDP) combina il parallelismo dei dati (partizionamento dei dati tra i nodi) con il partizionamento dei pesi tra i nodi. Per ogni operazione nei passaggi in avanti e indietro, i pesi vengono raccolti in modo che ogni slice abbia i pesi necessari. Anziché sincronizzare i gradienti utilizzando all-reduce, i gradienti vengono ridotti e distribuiti man mano che vengono prodotti. In questo modo, ogni slice riceve solo i gradienti per i pesi di cui è responsabile.

Analogamente al parallelismo dei dati, FSDP richiederà di scalare la dimensione del batch globale in modo lineare con il numero di slice. FSDP ridurrà la pressione sulla memoria man mano che aumenti il numero di sezioni. Questo perché il numero di pesi e lo stato dell'ottimizzatore per sezione diminuiscono, ma ciò avviene a scapito di un aumento del traffico di rete e di una maggiore possibilità di blocco a causa di un collettivo ritardato.

In pratica, l'FSDP tra slice è la soluzione migliore se aumenti il batch per slice, memorizzi più attivazioni per ridurre al minimo la rimaterializzazione durante il passaggio all'indietro o aumenti il numero di parametri nella rete neurale.

Le operazioni all-gather e all-reduce in FSDP funzionano in modo simile a quelle in DP, quindi puoi determinare se il tuo workload FSDP è limitato dalle prestazioni DCN nello stesso modo descritto nella sezione precedente.

Quando utilizzare il parallelismo della pipeline

Il parallelismo della pipeline diventa rilevante quando si ottengono prestazioni elevate con altre strategie di parallelismo che richiedono una dimensione del batch globale superiore alla dimensione massima del batch preferita. Il parallelismo della pipeline consente alle sezioni che compongono una pipeline di "condividere" un batch. Tuttavia, il parallelismo della pipeline presenta due svantaggi significativi:

1. Si verifica la "bolla della pipeline", in cui i chip sono inattivi perché sono in attesa dei dati.
2. Richiede il micro-batching, che riduce la dimensione effettiva del batch, l'intensità aritmetica e, in definitiva, l'utilizzo dei FLOP del modello.

Il parallelismo della pipeline deve essere utilizzato solo se le altre strategie di parallelismo richiedono una dimensione del batch globale troppo grande. Prima di provare il parallelismo della pipeline, vale la pena fare esperimenti per verificare empiricamente se la convergenza per campione rallenta alla dimensione del batch necessaria per ottenere FSDP ad alte prestazioni. FSDP tende a ottenere un utilizzo più elevato dei FLOP del modello, ma se la convergenza per campione rallenta con l'aumentare della dimensione del batch, il parallelismo della pipeline potrebbe comunque essere la scelta migliore. La maggior parte dei workload può tollerare dimensioni dei batch sufficientemente grandi da non trarre vantaggio dal parallelismo della pipeline, ma il tuo workload potrebbe essere diverso.

Se è necessario il parallelismo della pipeline, ti consigliamo di combinarlo con il parallelismo dei dati o FSDP. In questo modo, potrai ridurre al minimo la profondità della pipeline e, al contempo, aumentare la dimensione del batch per pipeline fino a quando la latenza della DCN non diventa un fattore meno importante per il throughput. In pratica, se hai N sezioni, considera pipeline di profondità 2 e N/2 repliche di parallelismo dei dati, poi pipeline di profondità 4 e N/4 repliche di parallelismo dei dati, continuando nello stesso modo, finché il batch per pipeline non diventa abbastanza grande da consentire di nascondere i collettivi DCN dietro l'aritmetica nel passaggio di backward. In questo modo, il rallentamento introdotto dal parallelismo della pipeline verrà ridotto al minimo, consentendoti al contempo di superare il limite della dimensione del batch globale.

Best practice per Multislice

Le sezioni seguenti descrivono le best practice per l'addestramento con Multislice.

Caricamento dei dati

Durante l'addestramento, carichiamo ripetutamente batch da un set di dati da inserire nel modello. Per evitare di lasciare le TPU senza lavoro, è importante disporre di un caricatore di dati asincrono ed efficiente che suddivida il batch tra gli host. L'attuale caricatore di dati in MaxText fa sì che ogni host carichi un sottoinsieme uguale di esempi. Questa soluzione è adeguata per il testo, ma richiede un nuovo sharding all'interno del modello. Inoltre, MaxText non offre ancora snapshot deterministici che consentano all'iteratore di dati di caricare gli stessi dati prima e dopo la preemption.

Checkpointing

La libreria di checkpointing Orbax fornisce primitive per il checkpointing di JAX PyTrees nell'archiviazione locale o nell'archiviazione Google Cloud . Forniamo un'integrazione di riferimento con il checkpoint sincrono in MaxText in checkpointing.py.

Configurazioni supportate

Le sezioni seguenti descrivono le forme delle sezioni, l'orchestrazione, i framework e il parallelismo supportati per Multislice.

Forme

Tutte le sezioni devono avere la stessa forma (ad esempio, lo stesso AcceleratorType). Le forme delle sezioni eterogenee non sono supportate.

Orchestrazione

L'orchestrazione è supportata con GKE. Per maggiori informazioni, consulta TPU in GKE.

Framework

Multislice supporta solo i workload JAX e PyTorch.

Parallelismo

Consigliamo agli utenti di testare Multislice con il parallelismo dei dati. Per saperne di più sull'implementazione del parallelismo della pipeline con Multislice, contatta il tuo rappresentante dell'accountGoogle Cloud .

Assistenza e feedback

Accogliamo con piacere tutti i feedback. Per condividere un feedback o richiedere assistenza, contattaci utilizzando il modulo di feedback o assistenza di Cloud TPU.