Panoramica di più sezioni di Cloud TPU

Cloud TPU Multislice è una tecnologia full stack che ottimizza le prestazioni che consente a un job di addestramento di utilizzare più sezioni TPU all'interno di un singolo pod in più pod con parallelismo dei dati semplice. Con i chip TPU v4, 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ù sezioni più piccole sono più facilmente disponibili consentendo un tempo di avvio più rapido quando si utilizzano più sezioni sezioni.

Quando viene eseguito il deployment in configurazioni multisezione, i chip TPU in ogni sezione comunicano tramite ICI (Inter-Chip-Interconnect). 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 sezioni. 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 una sezione rileva un evento di manutenzione, un prerilascio o un errore hardware, Cloud TPU creerà una nuova sezione. Nel raro caso in cui risorse insufficienti per creare una nuova sezione, la creazione non verrà completata finché l'hardware non sarà disponibile. Dopo aver creato la nuova sezione, tutte le altre le sezioni nell'ambiente Multislice verranno riavviate, quindi l'addestramento con uno script di avvio configurato correttamente, lo script di addestramento può riavviarsi automaticamente senza l'intervento dell'utente, oltre che durante il caricamento e il ripristino dall'ultimo checkpoint.

Set di dati

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

Networking dei data center (DCN)

Una rete con latenza più alta e velocità effettiva inferiore (se confrontata con ICI) connette le sezioni TPU in una configurazione con più sezioni.

Programmazione per le gang

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 le VM. Un host può eseguire al massimo quattro VM contemporaneamente. Ogni VM ha una TPU dedicata.

Inferenza

Carica un modello di machine learning preaddestrato su un host e fai previsioni in base dati.

Interconnect interchip (ICI)

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

Più sezioni

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

Nodo

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

Pod

Una raccolta di chip TPU collegati da interfacce di rete ICI dedicate. R Il pod consente di distribuire il carico di elaborazione su più TPU.

Risorsa in coda (QR)

Una rappresentazione delle risorse TPU, utilizzata per accodare e gestire una richiesta per un un ambiente TPU a sezione singola o multisezione.

Script di avvio

Uno script di avvio di Compute Engine standard eseguito ogni volta che viene avviata o riavviata una VM. Nel caso di più sezioni, viene specificato nella richiesta di creazione QR. Per ulteriori informazioni per informazioni sugli script di avvio di Cloud TPU, vedi Gestire le risorse TPU.

Sezione TPU

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

VM TPU

Una macchina virtuale che esegue Linux e 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 una macchina 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: ha come target la quota prenotata. Per utilizzare la quota riservata, devi avere un contratto di prenotazione con Google. Utilizza il flag --reserved quando crei le risorse.
• Spot: ha come target la quota prerilasciabile utilizzando le VM spot. Le tue risorse potrebbero essere 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à accodata a un modulo per la quota offerta da Cloud TPU, la disponibilità delle risorse non garantito. Selezionato per impostazione predefinita, non sono necessari flag.

Nota: Cloud TPU limita la dimensione della sezione per tutti i tipi di acceleratori. Per fetta limiti di dimensioni per le richieste on demand, consulta Dimensioni predefinite delle sezioni per i codici 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 un utente TPU v4 esistente e hai una prenotazione, potresti dover eseguire la migrazione la prenotazione 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 con Cloud TPU.

Il codice in shardings.py è progettato per aiutarti a iniziare a sperimentare diversi parallelismo le opzioni di CPU e memoria disponibili. Ad esempio, parallelismo dei dati, parallelismo dei dati con sharding completo (FSDP), e il parallelismo tensore. Il codice scala da una sezione singola a una sezione multipla ambienti cloud-native.

Parallelismo ICI

ICI è l'interconnessione ad alta velocità che collega le TPU in un sezione. Lo sharding ICI corrisponde allo sharding all'interno di una sezione. shardings.py fornisce tre parametri di parallelismo di 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 tensore a quattro vie	1	1	4
FSDP a 2 vie + parallelismo tensore a due vie	1	2	2

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

Questo esempio presuppone che tu abbia familiarità con l'esecuzione di codice su una singola sezione TPU come in Esegui un calcolo su una VM Cloud TPU utilizzando JAX. Questo esempio mostra come eseguire shardings.py su 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 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 con il seguente comando:

   $ 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 flag di comando

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

   runtime-version

   La [versione software di Cloud TPU](/tpu/docs/supported-tpu-configurations#tpu_software_versions).

   node-id

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

   reserved

   Utilizza la quota riservata durante la creazione dei slice.

   spot

   Utilizza la quota di VM Spot durante la creazione dei slice.

   Google Cloud CLI non supporta tutte le opzioni di creazione QR, come i tag. Per maggiori informazioni, consulta la sezione Creare codici QR.

4. Attendi che il codice QR sia nello stato ACTIVE, il che significa che i nodi worker sono nello stato READY. Una volta avviato il provisioning del codice QR, l'operazione potrebbe richiedere da uno a cinque minuti per il completamento, a seconda delle dimensioni del QR. Puoi controllare lo stato di una richiesta di QR utilizzando il seguente comando:

   $ gcloud compute tpus queued-resources \
     list --filter=your-qr-id

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) sulla VM TPU.

7. All'interno della directory del repository MaxText, esegui lo script di configurazione per installare JAX e altre dipendenze dalla sezione 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 suddivisione in 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 il QR e la sezione TPU. Devi eseguire questa pulizia passaggi dall'ambiente in cui hai configurato la sezione (prima esegui exit per uscire dalla sessione SSH). L'eliminazione richiederà da due a cinque minuti e può essere eseguita in background con il flag facoltativo --async.

    Importante: i codici QR devono essere eliminati dopo l'uso. In caso contrario continueranno per consumare la quota QR che viene monitorata separatamente da quella di TPU.

    $ gcloud compute tpus queued-resources
      delete your-qr-id --force (--async)

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.

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

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

dcn_tensor_parallelism deve essere sempre impostato su 1 perché la rete DCN è scadente adatto a questo sharding. Per i carichi di lavoro LLM tipici 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 la dimensione e il batch per slice), aumenti la quantità di parallelismo dei dati.

Esecuzione di shardings.py in un ambiente con più sezioni

Puoi eseguire shardings.py in un ambiente con più sezioni utilizzando multihost_runner.py oppure eseguendo shardings.py su ogni VM TPU. In questo caso 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, ripetutamente riutilizzano le stesse TPU. Poiché lo script multihost_runner.py dipende da per le connessioni SSH di lunga durata, lo sconsigliamo per i job a 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 costituiscono i tuoi slice. Puoi eseguire multihost_runner.py su una o su qualsiasi VM di Compute Engine nello stesso progetto delle tue sezioni. Esecuzione L'istruzione 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 sul tuo runner in una macchina virtuale.

2. Vai alla directory del repository.

3. Crea chiavi SSH per gcloud, consigliamo di lasciare una password vuota (premi inseriscilo due volte dopo aver eseguito questo 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 con più sezioni utilizzando queued-resources create.

   Nota: assicurati di disporre della rispettiva quota prima di selezionare --reserved, --spot o la quota on demand predefinita. Per informazioni sulla quota consulta i Criteri per le 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. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, vedi 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 sezione. 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 e così via in ordine numerico per ogni segmento.

   reserved

   Utilizza la quota prenotata durante la creazione delle sezioni.

   spot

   Utilizza la quota VM spot durante la creazione delle sezioni.

6. Quando inizia il provisioning del QR, il completamento dell'operazione potrebbe richiedere fino a cinque minuti, a seconda delle dimensioni del QR. Attendi l'arrivo della risorsa in coda (QR) lo stato ACTIVE. Puoi controllare lo stato di una richiesta QR utilizzando il seguente comando:

   $ gcloud compute tpus queued-resources list \
   --filter=your-qr-id

   Dovrebbe essere generato un output simile al seguente:

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

   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, pulisci i TPU e il codice QR. L'eliminazione richiede da due a cinque minuti e può essere eseguita in background con il flag facoltativo --async.

Scalabilità di un carico di lavoro a più sezioni

Prima di eseguire il modello in un ambiente con più sezioni, 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 quando si passa a Multislice. Per ottenere prestazioni elevate, la rete DCN deve essere mappata su dati paralleli, con sharding dei dati paralleli o assi paralleli della pipeline. Considerazioni sulle prestazioni e le strategie dello sharding sono illustrate più dettagliatamente nel Sharding con più sezioni per ottenere le massime prestazioni

Per verificare che il tuo codice possa accedere a tutti i dispositivi, puoi affermare che len(jax.devices()) è uguale al numero di chip nel tuo multisezione completamente gestito di Google Cloud. 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 aumentare la velocità lineare, aggiungi nuove sezioni delle stesse dimensioni di quelle esistenti sezione. Ad esempio, se utilizzi una sezione v4-512, Multisezione ottenere circa il doppio delle prestazioni aggiungendo una seconda sezione v4-512 e raddoppiando la dimensione del batch globale. Per ulteriori informazioni, vedi Sharding con più sezioni per ottenere le massime prestazioni

Eseguire il job su più slice

Esistono tre diversi approcci per l'esecuzione di un carico di lavoro personalizzato Ambiente su più sezioni:

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'esecutore 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 in 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 carico di lavoro e alla tua configurazione hardware.

Script di esecutore di produzione

Per i job di produzione che richiedono resilienza contro guasti hardware e altre prerilascio, è preferibile integrare direttamente la risorsa Crea risorsa in coda tramite Google Cloud CLI o tramite l'API Compute Engine. Come esempio pratico, forniamo multihost_job.py, che attiva la chiamata API Created Queued Resource con l'avvio appropriato per eseguire l'addestramento e riprendere con il prerilascio. multihost_job.py sia 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 oppure multihost_job.py per eseguire il tuo carico di lavoro personalizzato la configurazione multisezione. Tuttavia, se preferisci eseguire il provisioning e gestire l'ambiente direttamente con i comandi QR, vedi Gestire un ambiente con più sezioni.

Gestire un ambiente Multislice

Per eseguire manualmente il provisioning e la gestione dei codici QR senza utilizzare gli strumenti fornito nel repository MaxText, leggi le sezioni seguenti.

Crea codici QR

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

Input	Descrizione
your-qr-id	L'ID assegnato dall'utente al codice QR.
PROGETTO	Nome progetto Google Cloud
ZONA	us-central2-b
NETWORK_NAME	Nome delle reti VPC.
SUBNETWORK_NAME	Nome della subnet nelle reti VPC
RUNTIME_VERSION	tpu-ubuntu2204-base
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. Limitato a un massimo di 256 sezioni.
STARTUP_SCRIPT	Se aggiunto alla richiesta di creazione, uno script di avvio può essere eseguito ogni volta che viene eseguito il provisioning o il riavvio di uno slice TPU e se lo slice TPU viene riparato o reimpostato.

Crea una richiesta QR utilizzando gcloud

$ 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 del codice QR.

project

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

zone

La zona Google Cloud in cui creare il QR.

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. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni di 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 la quota riservata durante la creazione dei slice.

spot

Utilizza la quota di 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 quota, consulta i Criteri per le quote.

Crea una richiesta QR utilizzando curl

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

Nota :sostituisci i seguenti valori:

• your-project-number: il numero del tuo progetto Google Cloud
• your-zone - La zona in cui vuoi creare il codice QR
• accelerator-type: la versione e la dimensione di una singola sezione
• tpu-vm-runtime-version - Le versioni del runtime della VM TPU
• your-network-name - (Facoltativo) una rete a cui verrà collegato il codice QR
• your-subnetwork-name - Facoltativo, una subnet a cui verrà allegato il codice QR
• example-tag-1 - Facoltativo, una stringa tag arbitraria
• your-startup-script: uno script di avvio che verrà eseguito quando viene allocato il codice QR
• slice-count: il numero di sezioni TPU nell'ambiente Multislice
• your-qr-id: l'ID fornito dall'utente per il codice QR

Per ulteriori informazioni, consulta la sezione sull'API REST Queued Resource. documentazione 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 linea per utilizzare la capacità on demand predefinita.

Invia la richiesta di creazione del QR 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

Nota :sostituisci i seguenti valori:

• your-project-id: l'ID del tuo progetto Google Cloud
• your-zone - La zona in cui vuoi creare il codice QR
• your-qr-id: l'ID fornito dall'utente per il codice QR

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

Recupera lo stato di un QR

Per ottenere lo stato della richiesta del codice QR, utilizza il seguente comando:

  $ 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/operations/operation-your-qr-guid

Nota: nel comando precedente sostituisci quanto segue:

• your-project-id: l'ID del tuo progetto Google Cloud
• your-zone: la zona in cui creare il QR
• your-qr-guid: il GUID che segue name nell'output della richiesta di creazione del codice QR.

La risposta di questo comando contiene lo stato dell'operazione:

{
  "name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qa-guid>,
  "metadata": {...},
  "done": true,
  "response": {
    "@type": "type.googleapis.com/google.cloud.tpu.v2.QueuedResource",
    ...
    "state": {
      "state": "WAITING_FOR_RESOURCES"
    }
  }
}

Se il codice QR è stato creato correttamente ("done = true"), lo stato nel campo response sarà WAITING_FOR_RESOURCES o FAILED. Se il codice QR è nello stato WAITING_FOR_RESOURCES, è stato inserito in coda e inizierà il provisioning quando saranno disponibili risorse sufficienti. Se il codice QR è nello stato FAILED, il motivo dell'errore sarà nell'output. Per ulteriori informazioni su altri possibili stati, consulta la Guida dell'utente per le risorse in coda.

Terminata l'operazione, usa la descrizione dei codici QR. per monitorare le fasi del QR.

In un raro scenario, potresti trovare il tuo QR nello stato FAILED mentre alcune sezioni sono ACTIVE. In questo caso, elimina le risorse create, e riprova tra qualche minuto o contattaci al team di Cloud TPU per risolvere il problema.

SSH e installazione delle dipendenze

Esegui codice JAX sulle pod di TPU TPU descrive come connetterti alle VM TPU tramite 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 ${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 suddiviso in gruppi di quattro e viene inviato contemporaneamente. Il batch successivo di comandi viene inviato quando il batch corrente completa l'esecuzione. In caso di errore con uno dei comandi, l'elaborazione e non vengono inviati ulteriori batch. Per ulteriori informazioni, consulta il riferimento all'API ResourceQueued. Se il numero di sezioni che utilizzi supera quello dei thread del tuo computer locale (detto anche limite di batching) si imbatterà in un deadlock. Ad esempio, supponiamo che il limite per il batch sulla tua macchina locale sia 64. Se provi a eseguire un su più di 64 sezioni, ad esempio 100, il comando SSH le sezioni 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. Dopo aver avviato lo script di addestramento, sul primo batch di sezioni, il controllo tornerà immediatamente il comando SSH. Il comando SSH può quindi iniziare a eseguire lo script di addestramento il batch rimanente di 36 sezioni. 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 della rete

Assicurati che le sezioni TPU possano comunicare tra loro eseguendo questi passaggi. Installa JAX su ogni slice. Per ulteriori informazioni, vedi Esegui il codice JAX sulle pod di TPU TPU. Dichiarare len(jax.devices()) è uguale al numero di chip nel tuo multisezione completamente gestito di Google Cloud. A tal fine, 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 di codici QR

Puoi visualizzare lo stato dei tuoi QR utilizzando il comando queued-resources list:

$ gcloud compute tpus queued-resources list

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

Descrivi i codici QR

Per visualizzare la configurazione dettagliata e lo stato di un QR, utilizza il "Descrivi l'API QR". Puoi chiamare questa API utilizzando gcloud o curl.

Utilizzo di gcloud:

$ gcloud compute tpus queued-resources describe ${your-qr-id}
...state:
 state: ACTIVE
...

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

state indica lo stato di un codice QR. Per ulteriori informazioni sui possibili stati dei QR, consulta 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 codici QR

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

Le reimpostazioni di tutte le VM vengono eseguite in parallelo e un ResetQueuedResource il completamento dell'operazione richiede uno o due minuti. Per richiamare l'API, utilizza il seguente comando:

$ gcloud compute tpus queued-resources reset your-qr-id

Eliminare i codici QR

Per svincolare le risorse al termine della sessione di formazione, elimina le risorse in coda risorsa con il flag --force. L'eliminazione richiederà da due a cinque minuti e può essere eseguita in background con il flag facoltativo --async.

$ gcloud compute tpus queued-resources \
delete your-qr-id --force (--async)

Ripristino automatico da errori

In caso di interruzione, Multislice offre la riparazione della sezione interessata e il ripristino di tutte le sezioni in seguito. Le persone interessate la sezione viene sostituita con una nuova sezione e le altre sezioni altrimenti integre sono reimpostati. Se non è disponibile alcuna capacità per allocare un nuovo slice, 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 una sezione viene riallocata o viene reimpostata una VM. Specifica una startup nel payload JSON che invii all'API di richiesta QR di creazione.

Il seguente script di avvio (utilizzato nella sezione Creare QR) consente di ripristinare automaticamente gli errori e di riprendere l'addestramento 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"
         }
     ...
     }
   ]
 }
}

Prima di provare questa opzione, clona il repository MaxText fuori.

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 partizionare le attivazioni per tutte le dimensioni del modello (a volte chiamato parallelismo tensore) perché richiede troppa larghezza di banda tra le sezioni. 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 le sezioni utilizza la rete DCN e rallenterà le tue 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 solida scalabilità su più sezioni, la quantità di tempo richiesta per eseguire la riduzione completa su DCN deve essere inferiore al tempo richiesto per eseguire un passaggio all'indietro. 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.

Sono presenti 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 passaggio. Se utilizzi due sezioni, il modello utilizzerà 32 bit di larghezza di banda DCN. Se Se usi più di due sezioni, il compilatore eseguirà uno shuffling completo operativa e utilizzerai fino a 64 bit di larghezza di banda DCN per ciascun parametro 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 necessari per un passaggio in avanti e all'indietro sono 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 passaggio all'indietro è 4 * B.

Per garantire una scalabilità elevata su più sezioni, assicurati che l'ambiente operativo supera l'intensità aritmetica dell'hardware TPU. Per calcolare il intensità operativa, dividi il numero di FLOPS per parametro durante a ritroso della 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 sezioni: Operational intensity = 4 * B / 32

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

Questo suggerisce una dimensione minima del batch compresa tra 176.000 e 352.000 per Transformer basati su modelli linguistici. Poiché la rete DCN può rilasciare rapidamente pacchetti, la soluzione migliore per mantenere un margine di errore significativo, implementando solo il parallelismo dei dati se la dimensione del batch per pod è compresa tra 350.000 (due pod) e 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). Quindi puoi confrontarlo con il tempo di esecuzione previsto per ridurre DCN e capire se il parallelismo dei dati ha senso per te.

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 segmento 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, ciascuna sezione ottiene i gradienti solo per le ponderazioni di cui è responsabile.

Analogamente al parallelismo dei dati, FSDP richiederà di scalare la dimensione globale del batch in modo lineare con il numero di sezioni. FSDP ridurrà la pressione di memoria aumenti il numero di sezioni. 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 pertinente quando si raggiunge prestazioni elevate con altri strategie di parallelismo che richiedono una dimensione del batch globale maggiore di quella dimensione massima preferita per il batch. Il parallelismo della pipeline consente le sezioni costituito da una pipeline per "condividere" un batch. Tuttavia, il parallelismo della pipeline ha due svantaggi significativi:

1. Viene visualizzata la "bolla della pipeline" in cui i chip sono inattivi perché sono in attesa per i dati.
2. Richiede micro-batching che riduce la dimensione effettiva del batch, intensità aritmetica e, infine, modellare l'utilizzo FLOP.

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. Il FSDP tende a raggiungere maggiore utilizzo FLOP del modello, ma se la convergenza per campione rallenta la dimensione del batch aumenta, il parallelismo della pipeline potrebbe comunque essere la scelta migliore. Più alta carichi di lavoro sono in grado di tollerare dimensioni batch sufficientemente grandi per non trarre vantaggio parallelismo della pipeline, ma il tuo carico di lavoro potrebbe essere diverso.

Se è necessario il parallelismo della pipeline, ti consigliamo di combinarlo con i dati il parallelismo o FSDP. Ciò ti consentirà di ridurre al minimo la profondità della pipeline aumentando la dimensione del batch per pipeline fino a quando la latenza DCN non diventa un fattore determinante nella velocità effettiva. Concretamente, se hai N sezioni, considera le pipeline di repliche di parallelismo dei dati e di profondità 2 e N/2, poi pipeline di profondità 4 e N/4 di repliche di parallelismo dei dati e così via, fino a quando le dimensioni del batch per pipeline non diventano abbastanza da poter nascondere i gruppi DCN dietro l'aritmetica 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 più sezioni

Caricamento dei dati

Durante l'addestramento, carichiamo ripetutamente i batch da un set di dati da inserire nel feed un modello di machine learning. È importante disporre di un caricatore di dati asincrono ed efficiente che suddivida il batch tra gli host per evitare di sovraccaricare le TPU. Il caricatore di dati attuale in MaxText ha ogni carico dell'host un sottoinsieme uguale di esempi. Questa soluzione adeguato per il testo, ma richiede un controllo 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 checkpoint Orbax fornisce per il checkpoint di JAX PyTrees nello spazio di archiviazione locale o in Google Cloud Storage. Forniamo un'integrazione di riferimento con checkpoint sincrono in MaxText nel seguente paese: 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 più sezioni con parallelismo dei dati. Per saperne di più sull'implementazione del parallelismo della pipeline con Multislice, contatta 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.