Panoramica di Cloud TPU Multislice

Importante: 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ù slice TPU all'interno di un singolo pod o su slice in più pod con un semplice parallelismo dei dati. Con i chip TPU v4, questo 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, un singolo slice può offrire le migliori prestazioni. Tuttavia, più slice più piccoli sono più facilmente disponibili, consentendo un tempo di avvio più rapido quando Multislice viene utilizzato con slice più piccoli.

Quando vengono implementati in configurazioni multislice, i chip TPU in ogni sezione comunicano tramite inter-chip-interconnect (ICI). I chip TPU in diversi slice comunicano trasferendo i dati alle CPU (host), che a loro volta li trasmettono tramite la rete del data center (DCN).

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 con il calcolo per massimizzare le prestazioni.

Concetti

Tipo di acceleratore

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

Riparazione automatica

Quando uno slice rileva un evento di manutenzione, un prelievo o un guasto hardware, Cloud TPU crea un nuovo slice. Nei rari casi in cui le risorse non sono sufficienti per creare un nuovo slice, la creazione non verrà completata finché l'hardware non sarà disponibile. Dopo aver creato il nuovo livello, tutti gli altri livelli nell'ambiente Multislice verranno riavviati in modo che l'addestramento possa continuare.Con uno script di avvio configurato correttamente, lo script di addestramento può essere riavviato automaticamente senza intervento dell'utente, caricando e riprendendo dall'ultimo checkpoint.

Set di dati

I dati utilizzati da un modello per l'addestramento o l'inferenza.

Data Center Networking (DCN)

Una rete con latenza più elevata e velocità in uscita inferiore (rispetto all'ICI) che collega le sezioni TPU in una configurazione multislice.

Pianificazione in gruppo

Quando tutti i slice TPU vengono provisionati contemporaneamente, viene garantito che il provisioning venga eseguito correttamente su tutti o su nessuno dei slice.

Organizzatore

Un host è un computer fisico che esegue VM. Un host può eseguire al massimo quattro VM contemporaneamente. Ogni VM ha una TPU dedicata.

Inferenza

Carica un modello di machine learning pre-addestrato su un host ed esegui previsioni sui dati.

Interconnect interchip (ICI)

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

Pod

Una raccolta di chip TPU connessi da interfacce di rete ICI dedicate. Un pod ti consente di distribuire il carico di elaborazione su più TPU.

Risorsa in coda (QR)

Una rappresentazione delle risorse TPU, utilizzata per mettere in coda e gestire una richiesta per un ambiente TPU a singolo o multislice.

Script di avvio

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

Sezione TPU

Una sottosezione logica di un pod di TPU costituita da chip TPU. Tutti i chip di una slice comunicano tra loro utilizzando la rete ICI.

VM TPU

Una macchina virtuale con sistema operativo Linux che ha accesso alle TPU sottostanti. Per le TPU v4, ogni VM TPU ha accesso diretto a quattro chip. A volte chiamiamo una VM TPU un worker.

Tensor

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

Tensor Processing Unit (TPU)

Chip di accelerazione del machine learning sviluppato internamente da Google. Sono progettate per offrire un calcolo rapido ed efficiente per attività di machine learning chiave come la moltiplicazione di matrici.

Tipi di capacità Cloud TPU

Le TPU possono essere create da diversi tipi di capacità (consulta le opzioni di utilizzo in Come funzionano i prezzi 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 preassegnata utilizzando le VM spot. Le tue risorse potrebbero essere prelevate per fare spazio alle richieste di un job con priorità più elevata. 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à anticipata. La richiesta TPU verrà inserita in una coda di quota on demand offerta da Cloud TPU, la disponibilità delle risorse non è garantita. Selezionato per impostazione predefinita, non sono necessari flag.

Nota: Cloud TPU limita le dimensioni dei sezioni per tutti i tipi di acceleratore. Per i limiti di dimensione dei seggiolini per le richieste on demand, consulta Dimensioni dei seggiolini predefinite per i QR on demand.

Inizia

Se non hai mai utilizzato le TPU, inizia installando Google Cloud CLI e configurando il tuo ambiente Cloud TPU. Per utilizzare Multislice, le risorse TPU devono essere gestite come risorse in coda.

Se sei già un utente di TPU v4 e hai una prenotazione, potrebbe essere necessario eseguirne la migrazione a un nuovo sistema di prenotazione. Per ulteriori informazioni, contatta il rappresentante dell'account Google Cloud.

Esempio introduttivo

Questo tutorial utilizza il codice del repository GitHub di MaxText. MaxText è un LLM di base ad alte prestazioni, arbitrariamente scalabile, open source e ben testato scritto in Python e Jax. MaxText è stato progettato per essere addestrato in modo 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 suddivisi (FSDP) e parallelismo dei tensori. Il codice si adatta da un singolo slice a ambienti multislice.

Parallelismo ICI

Per ICI si intende l'interconnessione ad alta velocità che collega le TPU in un'unica sezione. Lo sharding ICI corrisponde allo sharding all'interno di uno slice. 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 ciascun 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 nella versione 4-8:

	ici_data_parallelism	ici_fsdp_parallelism	ici_tensor_parallelism
FSDP a 4 vie	1	4	1
Parallelismo tensoriale a 4 vie	1	1	4
FSDP bidirezionale + parallelismo tensore bidirezionale	1	2	2

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 il parallelismo a livello di flusso di dati al parallelismo dei dati.

Questo esempio presuppone che tu abbia dimestichezza con l'esecuzione di codice su un singolo slice TPU, come nell'articolo Eseguire un calcolo su una VM Cloud TPU utilizzando JAX. Questo esempio mostra come eseguire shardings.py su un singolo slice.

1. Configura l'ambiente:

   $ gcloud auth login
   $ gcloud config set project your-project-id
   $ gcloud config set compute/zone your-zone

2. Crea le chiavi SSH per gcloud. Ti consigliamo di lasciare la password vuota (premi Invio due volte dopo aver eseguito il seguente comando). Se ti viene chiesto se 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 YOUR_QR_ID \
       --accelerator-type your-accelerator-type \
       --runtime-version tpu-ubuntu2204-base \
       --node-id qr-id \
       [--reserved |--spot]

   Descrizioni dei flag dei comandi

   YOUR_QR_ID

   Una stringa definita dall'utente che identifica la richiesta del codice QR.

   accelerator-type

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

   runtime-version

   La versione software di Cloud TPU.

   node-id

   L'ID delle risorse TPU che verranno create in risposta alla richiesta di codice QR.

   reserved

   Utilizza una prenotazione durante la creazione dei sezioni.

   spot

   Utilizza le VM Spot durante la creazione dei slice.

   Google Cloud CLI non supporta tutte le opzioni di creazione di codici QR, ad esempio i tag. Per ulteriori informazioni, vedi Creare codici QR.

   Console

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

      Vai alle 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 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 TPU, consulta 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. Per ulteriori informazioni, consulta Immagini VM TPU.

   7. Fai clic sul pulsante di attivazione/disattivazione Attiva l'inserimento in coda.

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

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

4. Attendi che la risorsa in coda sia nello stato ACTIVE, il che significa che i nodi worker sono in stato READY. Una volta avviato il provisioning delle risorse in coda, l'operazione 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 la gcloud CLI o la console Google Cloud:

   gcloud

   $ gcloud compute tpus queued-resources \
       list --filter=YOUR_QR_ID

   Console

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

      Vai alle TPU

   2. Fai clic sulla scheda Risorse in coda.

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

5. Una sezione v4-8 ha una singola VM TPU. Connettiti alla VM TPU tramite SSH:

   $ gcloud compute tpus tpu-vm ssh YOUR_QR_ID

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

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

   $ bash setup.sh

8. Esegui il comando seguente per eseguire shardings.py sul tuo slice TPU.

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

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

9. Non esitare a esplorare altre strategie di frammentazione su ICI, ad esempio potresti provare la seguente combinazione:

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

10. Al termine, elimina la risorsa in coda e il segmento TPU. Devi eseguire questi passaggi di pulizia dall'ambiente in cui hai configurato il segmento (prima, esegui exit per uscire dalla sessione SSH). L'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'uso. In caso contrario, continueranno a consumare la quota di risorse in coda, che viene monitorata distintamente dalla quota TPU.

    gcloud

    $ gcloud compute tpus queued-resources
        delete YOUR_QR_ID --force (--async)

    Console

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

       Vai alle 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 con 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 slice, 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 sempre essere impostato su 1 perché la rete di distribuzione dei contenuti non è adatta a questo tipo di suddivisione. Per i tipici workload LLM sui chip v4, anche dcn_fsdp_parallelism deve essere impostato su 1 e quindi dcn_data_parallelism deve essere impostato sul numero di slice, ma questo dipende dall'applicazione.

Man mano che aumenti il numero di slice (supponendo di mantenere costanti le dimensioni e il batch per slice), aumenti la quantità di 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 che seguono sono molto simili a quelli descritti in Guida introduttiva: esperimenti rapidi su più slice nel 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, che riutilizzano ripetutamente le stesse TPU. Poiché lo script multihost_runner.py dipende da connessioni SSH di lunga durata, non è consigliato per i job di lunga esecuzione. 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 eseguire lo script multihost_runner.py. Utilizziamo il termine worker per indicare le VM TPU che compongono i tuoi slice. Puoi eseguire multihost_runner.py su un computer locale o su qualsiasi VM Compute Engine nello stesso progetto dei tuoi slice. 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, eseguiamo shardings.py su due slice v4-16, per un totale di quattro VM e 16 chip TPU. Puoi modificare l'esempio in modo che venga eseguito su più TPU.

Configura l'ambiente

1. Clona MaxText sulla macchina del tuo runner.

2. Vai alla directory del repository.

3. Crea le chiavi SSH per gcloud, ti consigliamo di lasciare vuota la password (premi ENTER due volte dopo aver eseguito il seguente comando). Se ti viene chiesto se 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 rispettiva quota prima di selezionare --reserved, --spot o la quota on demand predefinita. Per informazioni sui tipi di quote, consulta le norme relative alle quote.

   Il seguente comando mostra come creare una TPU multislice v4. Per utilizzare v5e, specifica un accelerator-type v5e (ad esempio v5litepod-16) e il runtime-version v5e (v2-alpha-tpuv5-lite).

   $ gcloud compute tpus queued-resources \
       create YOUR_QR_ID \
       --accelerator-type=your-accelerator-type \
       --runtime-version=tpu-vm-runtime-version \
       --node-count=node-count \
       --node-prefix=YOUR_QR_ID \
       [--reserved|--spot]

   Descrizioni dei flag dei comandi

   YOUR_QR_ID

   Una stringa definita dall'utente che identifica la richiesta del codice QR.

   accelerator-type

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

   runtime-version

   La versione software di Cloud TPU.

   node-count

   Il numero di sezioni da creare.

   node-prefix

   Il prefisso utilizzato per generare i nomi di ogni slice. A ogni slice viene aggiunto un numero al prefisso. Ad esempio, se imposti node-prefix su mySlice, i segmenti vengono denominati: mySlice-0, mySlice-1, continuando in ordine numerico per ogni segmento.

   reserved

   Utilizza una prenotazione durante la creazione dei sezioni.

   spot

   Utilizza le VM Spot durante la creazione dei slice.

   Console

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

      Vai alle 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 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 TPU, consulta 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. Per ulteriori informazioni, consulta Immagini VM TPU.

   7. Fai clic sul pulsante di attivazione/disattivazione Attiva l'inserimento in coda.

   8. Nel campo Nome della 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 slice, inserisci il numero di slice che vuoi creare.

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

6. Quando inizia il provisioning delle risorse in coda, l'operazione 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 risorsa in coda utilizzando la gcloud CLI o la console Google Cloud:

   gcloud

   $ gcloud compute tpus queued-resources list \
       --filter=YOUR_QR_ID

   L'output dovrebbe essere simile al seguente:

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

   Console

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

      Vai alle TPU

   2. Fai clic sulla scheda Risorse in coda.

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

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

7. Installa le dipendenze.

   $ python3 multihost_runner.py \
       --TPU_PREFIX=YOUR_QR_ID \
       --COMMAND="bash setup.sh"

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

   $ python3 multihost_runner.py \
       --TPU_PREFIX=YOUR_QR_ID \
       --COMMAND="python3 pedagogical_examples/shardings.py \
       --dcn_data_parallelism $SLICE_COUNT \
       --ici_fsdp_parallelism 8 \
       --batch_size 131072 \
       --embedding_dimension 2048"

   Nei file log vedrai circa 230 TFLOP al secondo di prestazioni.

9. Al termine, ripulisci le TPU e la 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.

Scalabilità di un carico di lavoro su 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 instead of jax.experimental.mesh_utils.create_device_mesh quando crei il mesh.

Queste dovrebbero essere le uniche modifiche al codice necessarie per il passaggio a Multislice. Per ottenere prestazioni elevate, il grafico a doppia correlazione deve essere mappato su assi paralleli ai dati, paralleli ai dati completamente suddivisi o paralleli alla pipeline. I fattori di rendimento e le strategie di suddivisione in parti vengono trattati in modo più dettagliato in Suddivisione in parti con Multislice per il rendimento massimo.

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

Scegliere le dimensioni delle sezioni per gli ambienti multislice

Per ottenere un'accelerazione lineare, aggiungi nuovi segmenti delle stesse dimensioni del segmento esistente. Ad esempio, se utilizzi un segmento v4-512, Multislice migliorerà il rendimento di circa il doppio aggiungendo un secondo segmento v4-512 e raddoppiando la dimensione del batch globale. Per ulteriori informazioni, consulta Eseguire lo sharding con Multislice per il massimo rendimento.

Eseguire il job su più slice

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

1. Utilizzando lo script di esecuzione degli esperimenti, multihost_runner.py
2. Utilizzando lo script del runner di produzione, multihost_job.py
3. Utilizzo di un approccio manuale

Script dell'eseguitore della sperimentazione

Lo script multihost_runner.py distribuisce il codice a 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 nel file README di MaxText.

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

Script di esecutore di produzione

Per i job di produzione che richiedono resilienza contro guasti hardware e altre preemption, è meglio eseguire l'integrazione direttamente con l'API Create Queued Resource. Come esempio pratico, forniamo multihost_job.py, che attiva la chiamata dell'API Resource Queued Created con lo script di avvio appropriato per eseguire l'addestramento e riprendere in caso di prelazione. 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 offre un ciclo di iterazione così rapido come multihost_runner.py.

Approccio manuale

Ti consigliamo di utilizzare o adattare multihost_runner.py o multihost_job.py per eseguire il tuo carico di lavoro personalizzato nella configurazione Multislice. Tuttavia, se preferisci eseguire il provisioning e gestire il tuo ambiente direttamente utilizzando 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 repo MaxText, leggi le sezioni seguenti.

Creare risorse in coda

gcloud

1. Imposta le seguenti variabili di ambiente prima di eseguire il provisioning della capacità:

   $ export YOUR_QR_ID=your-queued-resource-id
   $ export PROJECT=your-project-name
   $ export ZONE=us-central2-b
   $ export NETWORK_NAME=your-network-name
   $ export SUBNETWORK_NAME=your-subnetwork-name
   $ export RUNTIME_VERSION=tpu-ubuntu2204-base
   $ export ACCELERATOR_TYPE=v4-16
   $ export SLICE_COUNT=4
   $ export STARTUP_SCRIPT="#!/bin/bash\n ..."
   $ gcloud config set project project-name
   $ gcloud config set compute/zone zone

   Descrizioni delle variabili

   Input	Descrizione
   YOUR_QR_ID	L'ID assegnato dall'utente alla risorsa in coda.
   PROGETTO	Nome del progetto Google Cloud
   ZONA	Specifica la zona in cui creare le risorse.
   NETWORK_NAME	Nome delle reti VPC.
   SUBNETWORK_NAME	Nome della subnet nelle reti VPC
   RUNTIME_VERSION	La versione software di Cloud TPU.
   ACCELERATOR_TYPE	v4-16
   EXAMPLE_TAG_1, EXAMPLE_TAG_2 …	Tag utilizzati per identificare origini o destinazioni valide per i firewall di rete
   SLICE_COUNT	Numero di sezioni. Limite massimo di 256 sezioni.
   STARTUP_SCRIPT	Se specifichi uno script di avvio, lo script viene eseguito quando viene eseguito il provisioning o il riavvio del segmento TPU.

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

   $ gcloud compute tpus queued-resources \
       create ${YOUR_QR_ID} \
       --project your-project-id \
       --zone your-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]

   Descrizioni dei flag dei comandi

   YOUR_QR_ID

   Una stringa definita dall'utente che identifica la richiesta di risorse in coda.

   project

   Il progetto Google Cloud in cui stai creando la richiesta di risorsa in coda.

   zone

   La zona Google Cloud in cui creare la risorsa in coda.

   node-count

   Il numero di sezioni da creare.

   accelerator-type

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

   runtime-version

   La versione software di Cloud TPU.

   network

   Il nome di una rete VPC a cui collegare la risorsa TPU.

   subnetwork

   Il nome di una subnet VPC a cui collegare la risorsa TPU.

   reserved

   Utilizza una prenotazione durante la creazione dei sezioni.

   spot

   Utilizza le VM Spot durante la creazione dei slice.

Assicurati di disporre della rispettiva quota prima di selezionare --reserved, --spot o la quota on demand predefinita. Per informazioni sui tipi di quote, consulta le norme relative alle quote.

curl

1. Crea un file denominato queued-resource-req.json e copia al suo interno il seguente codice 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 numero del tuo progetto Google Cloud
   • your-zone: la zona in cui vuoi creare la risorsa in coda
   • accelerator-type: la versione e le dimensioni di un singolo slice. Multislice è supportato solo su Cloud TPU v4 e versioni successive di TPU.
   • tpu-vm-runtime-version: la versione del 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 sottorete a cui verrà collegata la risorsa in coda
   • example-tag-1 - Facoltativo, una stringa 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 nell'ambiente Multislice
   • YOUR_QR_ID: l'ID fornito dall'utente per la risorsa in coda

   Per ulteriori informazioni, consulta la documentazione dell'API REST Resource Queued per tutte le opzioni disponibili.

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

   Per utilizzare la capacità di 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_QR_ID

   Sostituisci i seguenti valori:

   • your-project-id: l'ID del tuo progetto Google Cloud
   • your-zone: la zona in cui vuoi creare la risorsa in coda
   • YOUR_QR_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 alle 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 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 su versioni TPU successive. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di 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. Per ulteriori informazioni, consulta Immagini VM TPU.

7. Fai clic sul pulsante di attivazione/disattivazione Attiva l'inserimento in coda.

8. Nel campo Nome della 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 slice, inserisci il numero di slice che vuoi creare.

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

Recuperare lo stato di una risorsa in coda

gcloud

$ gcloud compute tpus queued-resources describe ${YOUR_QR_ID}

Per una risorsa in coda 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 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 alle TPU

2. Fai clic sulla scheda Risorse in coda.

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

Dopo aver eseguito il provisioning della TPU, puoi anche visualizzare i dettagli della richiesta di risorsa in coda andando alla pagina TPU, individuando la TPU e facendo clic sul nome della richiesta di risorsa in coda corrispondente.

In un raro scenario, potresti trovare la risorsa in coda nello stato FAILED mentre alcune sezioni sono ACTIVE. In questo caso, elimina le risorse create, riprova dopo alcuni minuti o contatta l'assistenza Google Cloud.

SSH e installa le dipendenze

L'articolo Eseguire il codice JAX nelle sezioni di pod di TPU descrive come connetterti alle VM TPU utilizzando SSH in un'unica 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 ${YOUR_QR_ID} \
        --zone your-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 del QR utilizzando SSH. Il comando viene raggruppato in gruppi di quattro e inviato contemporaneamente. Il successivo batch di comandi viene inviato al termine dell'esecuzione del batch corrente. Se si verifica un errore con uno dei comandi, l'elaborazione si interrompe e non vengono inviati ulteriori batch. Per ulteriori informazioni, consulta il riferimento all'API ResourceInQueue. Se il numero di slice in uso supera il limite di threading (chiamato anche limite di batching) del computer locale, si verificherà un deadlock. Ad esempio, assumere che il limite di raggruppamento sulla tua macchina locale sia 64. Se tenti di eseguire uno script di addestramento su più di 64 slice, ad esempio 100, il comando SSH suddivide i slice in batch. Eseguirà lo script di addestramento sul primo batch di 64 sezioni e attenderà il completamento degli script prima di eseguirlo sul batch rimanente di 36 sezioni. Tuttavia, il primo lotto di 64 slice non può essere completato finché i 36 slice rimanenti non iniziano a eseguire lo script, causando un blocco.

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

Configurazione di rete

Assicurati che i slice TPU possano comunicare tra loro eseguendo i seguenti passaggi. Installa JAX su ogni slice. Per ulteriori informazioni, consulta Eseguire il codice JAX nelle pod di TPU TPU. Assicurati che len(jax.devices()) sia uguale al numero di chip nel tuo ambiente Multislice. Per farlo, in ogni segmento esegui:

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

Se esegui questo codice su quattro slice di v4-16, ci sono otto chip per slice e quattro slice, quindi jax.devices() dovrebbe restituire un totale di 32 chip (dispositivi).

Elenco delle risorse in coda

gcloud

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

$ gcloud compute tpus queued-resources list

L'output è simile al seguente:

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

Console

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

   Vai alle 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 carichi di lavoro connettendoti a tutti gli host di ogni slice tramite SSH e eseguendo il seguente comando su tutti gli host.

$ gcloud compute tpus tpu-vm ssh YOUR_QR_ID \
    --zone=your-zone \
    --worker=all \
    --node=all \
    --command="command-to-run"

Reimpostazione dei QR

L'API ResetQueuedResource può essere utilizzata per ripristinare tutte le VM in un codice QR ACTIVE. Il ripristino delle VM cancella forzatamente la memoria della macchina e reimposta la VM allo stato iniziale. Tutti i dati memorizzati localmente rimarranno invariati e lo script di avvio verrà richiamato dopo un ripristino dei dati di fabbrica. L'API ResetQueuedResource può essere utile quando vuoi riavviare tutte le TPU. Ad esempio, quando l'addestramento è bloccato, è più facile reimpostare tutte le VM rispetto al debug.

I reset 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 YOUR_QR_ID

Eliminazione delle risorse in coda

Per rilasciare le risorse alla fine della sessione di addestramento, elimina la 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.

gcloud

$ gcloud compute tpus queued-resources \
    delete YOUR_QR_ID --force (--async)

Console

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

   Vai alle 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.

Ripristino automatico da errori

In caso di interruzione, Multislice offre la riparazione senza intervento della fetta interessata e il ripristino di tutte le fette in un secondo momento. La slice interessata viene sostituita da una nuova e le altre slice in buono stato vengono reimpostate. Se non è disponibile alcuna capacità per allocare uno slice sostitutivo, l'addestramento si interrompe.

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

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

{
 "tpu": {
   "node_spec": [
     {
      ...
         "metadata": {
               "startup-script": "#! /bin/bash \n pwd \n runuser -l user1 -c 'cd /home/user1/MaxText && python3 MaxText/train.py 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 repo MaxText prima di provare questa funzionalità.

Profilazione e debug

Il profiling è lo stesso negli ambienti a singolo livello e multilivello. Per maggiori informazioni, consulta Profilazione dei programmi JAX.

Ottimizza la formazione

Sharding con Multislice per il massimo rendimento

Per ottenere il massimo rendimento negli ambienti multislice è necessario considerare come eseguire lo sharding nei vari slice. In genere sono disponibili tre opzioni (parallelismo dei dati, parallelismo dei dati completamente suddiviso in parti e parallelismo della pipeline). Sconsigliamo di eseguire lo sharding delle attivazioni nelle dimensioni del modello (a volte chiamato parallelismo tensoriale) perché richiede troppa larghezza di banda inter-slice. Per tutte queste strategie, puoi mantenere la stessa strategia di suddivisione all'interno di uno slice che ha funzionato per te in passato.

Ti consigliamo di iniziare con il parallelismo dei dati puro. L'utilizzo del parallelismo dei dati completamente suddivisi in parti è utile per liberare l'utilizzo della memoria. Lo svantaggio è che la comunicazione tra i vari slice utilizza la rete DCN e rallenta il tuo carico di lavoro. Utilizza il parallelismo della pipeline solo se necessario in base alle dimensioni del batch (come analizzato di seguito).

Quando utilizzare il parallelismo dei dati

Il parallelismo dei dati puri funziona bene nei casi in cui un carico di lavoro è in esecuzione correttamente, ma vuoi migliorarne il rendimento eseguendo il ridimensionamento su più slice.

Per ottenere una scalabilità elevata su più slice, il tempo necessario per eseguire l'operazione all-reduce sulla rete di calcolo distribuita deve essere inferiore al tempo necessario per eseguire un passaggio a ritroso. La rete in data center viene utilizzata per la comunicazione tra i vari slice ed è un fattore limitante del throughput del carico di lavoro.

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

Esistono quattro chip per host TPU e ogni host ha una larghezza di banda di rete massima di 50 Gbps.

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

Il modello utilizzerà da 32 a 64 bit di larghezza di banda DCN per ogni parametro per ogni passaggio. Se utilizzi due slice, il modello utilizzerà 32 bit di larghezza di banda DCN. Se utilizzi più di due slice, il compilatore eseguirà un'operazione di riduzione completa con riordinamento e utilizzerai fino a 64 bit di larghezza di banda DCN per ogni parametro per ogni 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 richiesto per un passaggio in avanti e uno all'indietro è approssimativamente 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 passaggio all'indietro è 4 * B.

Per garantire un'elevata scalabilità su più slice, 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 passaggio all'indietro per la larghezza di banda della rete (in bit) per parametro per passaggio: Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

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

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

Ciò suggerisce un batch minimo compreso 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, implementando il parallelismo dei dati solo se le dimensioni del batch per pod sono almeno 350.000 (due pod) - 700.000 (molti pod).

Per altre architetture del modello, dovrai stimare il tempo di esecuzione del passaggio all'indietro per ogni slice (misurando il tempo utilizzando un profiler o contando i FLOPS). Poi puoi confrontarlo con il tempo di esecuzione previsto per ridurre il sovraccarico su DCN e ottenere una buona stima di se il parallelismo dei dati è adatto alle tue esigenze.

Quando utilizzare il parallelismo dei dati completamente suddiviso (FSDP)

Il parallelismo dei dati completamente suddiviso (FSDP) combina il parallelismo dei dati (suddivisione dei dati tra i nodi) con la suddivisione dei pesi tra i nodi. Per ogni operazione nei passaggi in avanti e all'indietro, i pesi vengono raccolti in modo che ogni frazione abbia i pesi di cui ha bisogno. Invece di sincronizzare i gradienti utilizzando all-reduce, i gradienti vengono sparsi durante la produzione. In questo modo, ogni segmento ottiene solo i gradienti per i pesi di cui è responsabile.

Come per il parallelismo dei dati, la FSDP richiede di scalare la dimensione del batch globale in modo lineare con il numero di slice. L'FSDP riduce la pressione sulla memoria man mano che aumenti il numero di slice. Questo perché il numero di pesi e dello stato dell'ottimizzatore per slice diminuisce, ma a costo di un aumento del traffico di rete e di una maggiore possibilità di blocco a causa di un insieme ritardato.

In pratica, l'FSDP tra i vari slice è ideale se aumenti il batch per slice, memorizzi più attivazioni per ridurre al minimo la rematerializzazione durante il passaggio a ritroso 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 carico di lavoro FSDP è limitato dalle prestazioni del DCN nello stesso modo descritto nella sezione precedente.

Quando utilizzare il parallelismo della pipeline

Il parallelismo della pipeline diventa rilevante per ottenere prestazioni elevate con altre strategie di parallelismo che richiedono una dimensione del batch globale maggiore della dimensione del batch massima preferita. Il parallelismo della pipeline consente ai sezioni che la compongono di "condividere" un batch. Tuttavia, il parallelismo della pipeline presenta due aspetti negativi significativi:

1. Si verifica la "bolla della pipeline" in cui i chip sono inattivi perché sono in attesa di dati.
2. Richiede micro-batching che riduce le dimensioni effettive del batch, l'intensità arithmetica e, in ultima analisi, l'utilizzo di 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 con la dimensione del batch necessaria per ottenere un FSDP ad alte prestazioni. L'FSDP tende a ottenere un utilizzo più elevato dei FLOP del modello, ma se la convergenza per campione rallenta con l'aumento delle dimensioni del batch, il parallelismo della pipeline potrebbe comunque essere la scelta migliore. La maggior parte dei carichi di lavoro può tollerare dimensioni dei batch sufficientemente grandi da non trarre vantaggio dal parallelismo della pipeline, ma il tuo carico di lavoro potrebbe essere diverso.

Se è necessario il parallelismo della pipeline, ti consigliamo di combinarlo con il parallelismo dei dati o con FSDP. In questo modo, puoi ridurre al minimo la profondità della pipeline, aumentando al contempo le dimensioni del batch per pipeline finché la latenza del DCN non diventa meno importante per il throughput. Nello specifico, se hai N slice, valuta le pipeline di profondità 2 e le repliche N/2 del parallelismo dei dati, quindi le pipeline di profondità 4 e le repliche N/4 del parallelismo dei dati e così via, fino a quando il batch per pipeline non diventa sufficientemente grande da nascondere i collettivi DCN dietro l'aritmetica nel passaggio all'indietro. In questo modo, ridurrai al minimo il rallentamento introdotto dal parallelismo della pipeline, consentendoti al contempo di superare il limite di dimensione del batch globale.

Best practice per le immagini multislice

Caricamento dei dati

Durante l'addestramento, carichiamo ripetutamente batch da un set di dati da fornire al modello. È importante disporre di un caricatore di dati asincrono ed efficiente che suddivida il batch tra gli host per evitare di sovraccaricare le TPU. L'attuale caricatore di dati in MaxText prevede che ogni host carichi un sottoinsieme uguale di esempi. Questa soluzione è sufficiente per il testo, ma richiede un nuovo ripartimento all'interno del modello. Inoltre, MaxText non offre ancora snapshot deterministici che consentano all'iteratore dei dati di caricare gli stessi dati prima e dopo la prelazione.

Controllo a punti di controllo

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

Configurazioni supportate

Forme

Tutti i segmenti devono avere la stessa forma (ad esempio la stessa AcceleratorType). Le forme dei segmenti eterogenee non sono supportate.

Orchestrazione

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

Framework

Multislice supporta solo i carichi di lavoro 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 rappresentante dell'account Google Cloud.

Assistenza e feedback

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