Panoramica multisezione di Cloud TPU

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

Se viene eseguito il deployment in configurazioni multisezione, i chip TPU in ogni sezione comunicano tramite Inter-chip-Interconnect (ICI). I chip TPU in sezioni diverse comunicano trasferendo i dati alle CPU (host), che a loro volta li trasmettono attraverso la rete dei data center (DCN).

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

Concetti

Tipo di acceleratore

La forma di ogni sezione TPU che comprende una sezione multisezione. Ogni sezione in una richiesta multisezione è 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 guasto hardware, Cloud TPU crea una nuova sezione. Nei rari casi in cui le risorse a disposizione siano insufficienti per creare una nuova sezione, la creazione non verrà completata finché l'hardware non sarà disponibile. Dopo la creazione della nuova sezione, tutte le altre sezioni nell'ambiente Multislice verranno riavviate in modo che l'addestramento possa continuare.Con uno script di avvio configurato correttamente, lo script di addestramento può riavviarsi automaticamente senza intervento dell'utente, caricando e riprendendo dal checkpoint più recente.

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ù elevata e velocità effettiva inferiore (rispetto a ICI) che collega le sezioni TPU in una configurazione multisezione.

Programmazione bande

Quando viene eseguito contemporaneamente il provisioning di tutte le sezioni TPU, in modo da garantire che venga eseguito correttamente il provisioning di tutte o nessuna delle sezioni.

Organizzatore

Un host è un computer fisico che esegue delle VM. Un host può eseguire al massimo 4 VM alla volta. Ogni VM ha una TPU dedicata.

Inferenza

Carica un modello di machine learning preaddestrato su un host ed esegui previsioni sui dati.

Interchip Interconnect (ICI)

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

Multisezione

Due o più sezioni di chip TPU in grado di comunicare tramite DCN.

Nodo

Nel contesto multisezione, il nodo fa riferimento a una singola sezione TPU. A ogni sezione TPU in una sezione multisezione viene assegnato un ID nodo.

Pod

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

Risorsa in coda (QR)

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

Script di avvio

Uno script di avvio di Compute Engine standard che viene eseguito a ogni avvio o riavvio di una VM. Per Multislice, è 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 composto da chip TPU. Tutti i chip in una sezione comunicano tra loro tramite la rete ICI.

VM TPU

Una macchina virtuale con 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 i dati multidimensionali in un modello di machine learning.

Tensor Processing Unit (TPU)

Chip di accelerazione ML sviluppato internamente di Google. Sono progettati per offrire un'elaborazione rapida ed efficiente dal punto di vista energetico per attività chiave di machine learning come la moltiplicazione della matrice.

Tipi di capacità di Cloud TPU

Le TPU possono essere create da diversi tipi di capacità (vedi Opzioni di utilizzo in Come funzionano i prezzi delle TPU) :

• Prenotazione: ha come target la quota prenotata. Per utilizzare la quota ripristinata, devi avere un contratto di prenotazione con Google. Utilizza il flag --reserved durante la creazione delle risorse.
• Spot: ha come target la quota prerilasciabile utilizzando le VM spot. Le tue risorse potrebbero essere prerilasciate per fare spazio alle richieste per un job con priorità più elevata. Usa il flag --spot durante la creazione delle risorse.
• On demand: ha come target la quota on demand, che non richiede una prenotazione e non verrà prerilasciata. La richiesta TPU verrà accodata a una coda della quota on demand offerta da Cloud TPU; la disponibilità delle risorse non è garantita. Selezionato per impostazione predefinita, nessun flag necessario.

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

Inizia

Se non hai mai utilizzato TPU, installa Google Cloud CLI e configura il tuo ambiente Cloud TPU. Per utilizzare la multisezione, le risorse TPU devono essere gestite come risorse in coda.

Se sei un utente esistente di TPU v4 e hai una prenotazione, potresti dover eseguire la migrazione della prenotazione a un nuovo sistema di prenotazione. Per saperne di più, contatta il rappresentante del tuo account Google Cloud.

Esempio introduttivo

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

Il codice in shardings.py è progettato per aiutarti a iniziare a sperimentare con diverse opzioni di parallelismo. Ad esempio, parallelismo dei dati, parallelismo dei dati completamente con sharding (FSDP) e parallelismo tensore. Il codice è scalabile da una sezione singola ad ambienti multisezione.

Parallelismo ICI

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

• ici_data_parallelism
• ici_fsdp_parallelism
• ici_tensor_parallelism

I valori specificati per questi parametri determinano il numero di shard per 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 nella sezione.

La tabella seguente mostra input utente di esempio per il parallelismo ICI per i quattro chip disponibili nella versione v4-8:

	ici_data_parallelism	ici_fsdp_parallelism	ici_tensor_parallelism
FSDP a quattro vie	1	4	1
Parallelismo Tensor a 4 vie	1	1	4
FSDP a 2 vie + parallelismo tensore a 2 vie	1	2	2

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

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

1. Configura l'ambiente:

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

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

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

3. Esegui il provisioning delle TPU con il comando seguente:

   $ gcloud alpha 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 comando

   your-qr-id

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

   accelerator-type

   Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per maggiori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni TPU.

   runtime-version

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

   node-id

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

   reserved

   Utilizza la quota prenotata durante la creazione delle sezioni.

   best-effort

   Utilizza la quota best effort durante la creazione delle sezioni [valore predefinito].

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

4. Attendi che il codice QR sia nello stato ACTIVE, a indicare che i nodi worker sono nello stato READY. Una volta avviato il provisioning del codice QR, il completamento potrebbe richiedere da uno a cinque minuti, a seconda delle dimensioni del codice 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) nella VM TPU.

7. All'interno della directory del repository MaxText, esegui lo script di configurazione per installare JAX e altre dipendenze sulla sezione TPU. L'esecuzione dello script di configurazione richiede alcuni minuti.

   $ bash setup.sh
   

8. Esegui questo comando per eseguire shardings.py sulla tua sezione TPU.

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

   Puoi vedere i risultati nei log. Le tue TPU dovrebbero raggiungere circa 260 TFLOP al secondo o un notevole utilizzo di FLOP superiore al 90%. In questo caso, abbiamo selezionato approssimativamente il batch massimo che può rientrare nella memoria a larghezza di banda elevata (HBM) della TPU.

9. Puoi esplorare altre strategie di sharding su ICI, ad esempio puoi provare la seguente combinazione:

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

10. Al termine, elimina la sezione QR e TPU. Devi eseguire questi passaggi di pulizia dall'ambiente in cui hai configurato la sezione (prima esegui exit per uscire dalla sessione SSH). L'eliminazione richiederà da 2 a 5 minuti e può essere eseguita in background con il flag --async facoltativo.

    Importante: i codici QR devono essere eliminati dopo l'utilizzo. In caso contrario, continueranno a utilizzare la tua quota QR, che viene monitorata separatamente dalla quota TPU.

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

Partizionamento orizzontale multisezione utilizzando il parallelismo DCN

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

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

I valori di questi parametri devono essere vincolati in modo che dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism equivalga 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 sezioni
Parallelismo dei dati bidirezionale	2	1	1	2

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

Man mano che aumenti il numero di sezioni (supponendo che le dimensioni della sezione e il batch per sezione siano costanti), aumenti la quantità di parallelismo dei dati.

shardings.py in esecuzione in un ambiente multisezione

Puoi eseguire shardings.py in un ambiente multisezione utilizzando multihost_runner.py o eseguendo shardings.py su ogni VM TPU. In questo caso usiamo multihost_runner.py. I seguenti passaggi sono molto simili a quelli di Guida introduttiva: esperimenti rapidi su più sezioni del repository MaxText, ad eccezione del fatto che qui eseguiamo shardings.py anziché l'LLM più complesso in train.py.

Lo strumento multihost_runner.py è ottimizzato per esperimenti rapidi, riutilizzando ripetutamente le stesse TPU. Poiché lo script multihost_runner.py dipende da connessioni SSH di lunga durata, lo sconsigliamo per 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 esegui lo script multihost_runner.py. Utilizziamo il termine workers per indicare le VM TPU che compongono le tue sezioni. Puoi eseguire multihost_runner.py su una macchina locale o su qualsiasi VM di Compute Engine nello stesso progetto delle sezioni. L'esecuzione di multihost_runner.py su un worker non è supportata.

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

In questo esempio, eseguiamo shardings.py su due sezioni 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 tua macchina runner.

2. Vai alla directory del repository.

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

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

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

     $ export SLICE_COUNT=2
     

5. Crea un ambiente multisezione utilizzando queued-resources create.

   Nota: assicurati di disporre della rispettiva quota prima di selezionare --reserved, --best_effort o la quota on demand predefinita. Per informazioni sui tipi di quota, consulta i criteri per le quote.

   Il seguente comando mostra come creare una TPU multisezione v4. Per utilizzare la versione 5e, specifica un valore per la versione 5e accelerator-type (ad esempio v5litepod-16) e la classe v5e runtime-version (v2-alpha-tpuv5-lite).

     $ gcloud alpha 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 flag comando

   your-qr-id

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

   accelerator-type

   Il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per maggiori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni TPU.

   runtime-version

   La versione software Cloud TPU.

   node-count

   Il numero di sezioni da creare.

   node-prefix

   Il prefisso utilizzato per generare i nomi per ogni sezione. Viene aggiunto un numero al prefisso per ogni sezione. Ad esempio, se imposti node-prefix su mySlice, le sezioni hanno il nome: mySlice-0, mySlice-1 e così via.

   reserved

   Utilizza la quota prenotata durante la creazione delle sezioni.

   best-effort

   Utilizza la quota best effort durante la creazione delle sezioni [valore predefinito].

6. Quando viene avviato il provisioning del codice QR, il completamento può richiedere fino a cinque minuti, a seconda delle dimensioni del codice QR. Attendi che la risorsa in coda (QR) sia nello stato ACTIVE. Puoi controllare lo stato di una richiesta QR utilizzando il seguente comando:

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

   Questo dovrebbe generare 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 del tuo account Google Cloud se lo stato del codice QR è in stato WAITING_FOR_RESOURCES o PROVISIONING per più di 15 minuti.

7. Installare 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"
   

   Vedrai circa 230 TFLOP al secondo di prestazioni nei file di log.

9. Esegui la pulizia delle TPU e del QR al termine dell'operazione. L'eliminazione richiederà da 2 a 5 minuti e può essere eseguita in background con il flag facoltativo --async.

Scalabilità di un carico di lavoro a sezioni multiple

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

• Utilizza jax.experimental.mesh_utils.create_hybrid_device_mesh anziché jax.experimental.mesh_utils.create_device_mesh durante la creazione del mesh.

Queste sono le uniche modifiche necessarie al codice durante il passaggio alla sezione multisezione. Per ottenere prestazioni elevate, la DCN deve essere mappata su assi paralleli di dati e con dati completamente con sharding in parallelo o in parallelo delle pipeline. Le considerazioni sulle prestazioni e le strategie di suddivisione in segmenti sono discusse in dettaglio nella sezione Suddivisione in segmenti con più sezioni per le massime prestazioni.

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

Scelta delle dimensioni delle sezioni per gli ambienti multisezione

Per aumentare la velocità lineare, aggiungi nuove sezioni delle stesse dimensioni della sezione esistente. Ad esempio, se utilizzi una sezione v4-512, Multislice raggiungerà circa il doppio delle prestazioni aggiungendo una seconda sezione v4-512 e raddoppiando le dimensioni del batch globali. Per ulteriori informazioni, consulta la sezione Suddivisione in segmenti con più sezioni per le massime prestazioni.

Esecuzione del job su più sezioni

Esistono tre diversi approcci per eseguire il carico di lavoro personalizzato in un ambiente multisezione:

1. Utilizzo dello script esecutivo della sperimentazione, multihost_runner.py
2. multihost_job.py con lo script esecutivo di produzione
3. Utilizzo di un approccio manuale

Script esecutivo della sperimentazione

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

Poiché multihost_runner.py mantiene connessioni SSH persistenti, è adatto solo per esperimenti di dimensioni modeste e a breve esecuzione. Puoi adattare i passaggi del tutorial multihost_runner.py al tuo carico di lavoro e alla configurazione dell'hardware.

Script runner di produzione

Per i job di produzione che necessitano di resilienza contro guasti hardware e altre prerilasci, è consigliabile integrare direttamente l'API Create Queued Resource. Come esempio funzionante, forniamo multihost_job.py, che attiva la chiamata API Created Queued Resource con lo script di avvio appropriato per eseguire l'addestramento e riprendere con prerilascio. Lo script multihost_job.py è documentato nel file README di MaxText.

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

Approccio manuale

Ti consigliamo di utilizzare o adattare multihost_runner.py o multihost_job.py per eseguire il carico di lavoro personalizzato nella configurazione multislice. Tuttavia, se preferisci eseguire il provisioning e gestire l'ambiente direttamente tramite i comandi QR, consulta Gestione di un ambiente multisezione.

Gestire un ambiente multisezione

Per eseguire il provisioning manuale dei codici QR e gestirli senza utilizzare gli strumenti forniti nel repository MaxText, leggi le sezioni seguenti.

Crea 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 del codice QR assegnato dall'utente.
PROGETTO	Nome progetto Google Cloud
MARGINE	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. Massimo 256 sezioni.
STARTUP_SCRIPT	Se viene aggiunto alla richiesta di creazione, uno script di avvio può essere eseguito ogni volta che viene eseguito o riavviato una sezione TPU e se la sezione TPU viene riparata o reimpostata.

Crea una richiesta QR utilizzando gcloud

$ gcloud alpha 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 flag comando

your-qr-id

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

project

Una stringa definita dall'utente che identifica la richiesta di 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 maggiori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni TPU.

runtime-version

La versione software 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 prenotata durante la creazione delle sezioni.

spot

Utilizza la quota delle VM spot durante la creazione delle sezioni.

Assicurati di disporre della rispettiva quota prima di selezionare --reserved, --spot o la quota on demand predefinita. Per informazioni sui tipi di quota, consulta la pagina 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 tuo codice QR
• accelerator-type: la versione e le dimensioni di una singola sezione
• tpu-vm-runtime-version - Le versioni runtime VM TPU
• your-network-name: facoltativo, una rete a cui verrà collegato il codice QR
• your-subnetwork-name: facoltativo, una subnet a cui verrà collegato il codice QR
• example-tag-1: facoltativo, una stringa tag arbitraria
• your-startup-script: uno script di avvio che verrà eseguito all'allocazione del QR
• slice-count: il numero di sezioni TPU nel tuo ambiente multisezione
• your-qr-id: l'ID fornito dall'utente per il codice QR

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

Nota: sebbene node_spec sia un elenco, è supportata una sola voce node_spec.

Per utilizzare la capacità Spot, sostituisci:

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

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

Invia la richiesta di creazione del codice 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 tuo 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 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 codice QR
• your-qr-guid: il GUID che segue name nell'output della richiesta di creazione 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 viene creato correttamente ("done = true"), lo stato all'interno del campo response sarà WAITING_FOR_RESOURCES o FAILED. Se il QR è nello stato WAITING_FOR_RESOURCES, significa che è stato accodato e inizierà il provisioning quando le risorse saranno sufficienti. Se il codice QR è nello stato FAILED, il motivo dell'errore sarà riportato nell'output. Per ulteriori informazioni sugli altri stati possibili, consulta la guida dell'utente sulle risorse in coda.

Al termine dell'operazione, utilizza la descrizione dei QR per monitorare le fasi del QR.

In rari casi, potresti trovare il codice QR nello stato FAILED, mentre alcune sezioni sono ACTIVE. In questo caso, elimina le risorse create e riprova tra qualche minuto oppure contatta il team di Cloud TPU per risolvere il problema.

SSH e installa le dipendenze

Esegui codice JAX sulle pod di TPU TPU descrive come connettersi alle VM TPU utilizzando SSH in un'unica sezione. Per connetterti a tutte le VM TPU nel tuo ambiente multisezione tramite SSH e installare 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 in QR utilizzando SSH. Il comando viene suddiviso in gruppi di quattro persone e viene inviato contemporaneamente. Il batch successivo di comandi viene inviato quando il batch attuale completa l'esecuzione. In caso di errore con uno dei comandi, l'elaborazione si interrompe e non vengono inviati ulteriori batch. Per ulteriori informazioni, consulta il riferimento sull'API Queued Resource. Se il numero di sezioni in uso supera il limite di threading del computer locale (detto anche limite di batch), riscontrerai un deadlock. Ad esempio, supponiamo che il limite di batch sulla tua macchina locale sia 64. Se provi a eseguire uno script di addestramento su più di 64 sezioni, ad esempio 100, il comando SSH suddividerà le sezioni in batch. Verrà eseguito lo script di addestramento sul primo batch di 64 sezioni e attenderà il completamento degli script prima di eseguire lo script sul batch rimanente di 36 sezioni. Tuttavia, il primo batch di 64 sezioni non può completare finché le 36 sezioni rimanenti non iniziano a eseguire lo script, causando un deadlock.

Per evitare questo scenario, puoi eseguire lo script di addestramento in background su ogni VM aggiungendo una e commerciale (&) al comando di script specificato con il flag --command. In questo modo, dopo aver avviato lo script di addestramento sul primo batch di sezioni, il controllo tornerà immediatamente al comando SSH. Il comando SSH può quindi avviare lo script di addestramento sul batch rimanente di 36 sezioni. Quando esegui i comandi in background, devi collegare i flussi stdout e stderr in modo appropriato. Per aumentare il parallelismo nello stesso 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 ciascuna delle sezioni. Per maggiori informazioni, consulta Eseguire il codice JAX sulle pod di TPU TPU. Dichiara che len(jax.devices()) è uguale al numero di chip nel tuo ambiente multislice. Per farlo, esegui:

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

Se esegui questo codice su quattro sezioni delle versioni v4-16, ci sono otto chip per sezione e quattro, jax.devices() restituisce un totale di 32 chip (dispositivi).

Elenca QR

Puoi visualizzare lo stato dei tuoi codici 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 QR

Per visualizzare la configurazione dettagliata e lo stato di un QR, usa l'API Descrivi 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 rappresenta lo stato di un codice QR. Per ulteriori informazioni sui possibili stati dei QR, consulta Risorse in coda.

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

Puoi eseguire manualmente i carichi di lavoro connettendoti a tutti gli host in ogni sezione tramite SSH ed 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 in corso...

L'API ResetQueuedResource può essere utilizzata per reimpostare 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 archiviati localmente rimarranno intatti e lo script di avvio verrà richiamato dopo un ripristino. L'API ResetQueuedResource può essere utile quando vuoi riavviare tutte le TPU. Ad esempio, quando l'addestramento si blocca e la reimpostazione di tutte le VM è più facile che eseguire il debug.

Il ripristino di tutte le VM viene eseguito in parallelo e il completamento di un'operazione ResetQueuedResource richiede uno o due minuti. Per richiamare l'API, utilizza il seguente comando:

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

Eliminazione dei codici QR in corso...

Per rilasciare risorse al termine della sessione di addestramento, elimina la risorsa in coda con il flag --force. L'eliminazione richiederà da 2 a 5 minuti e può essere eseguita in background con il flag --async facoltativo.

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

Ripristino automatico da errori

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

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

Il seguente script di avvio (utilizzato in Crea QR) ti consente di ripristinare automaticamente gli errori e di 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 repository MaxText prima di provare.

Profilazione e debug

La profilazione è la stessa negli ambienti a sezione singola e multisezione. Per ulteriori informazioni, consulta la sezione Profilazione dei programmi JAX.

Ottimizza la formazione

Sharding con multisezione per ottenere le massime prestazioni

Ottenere le massime prestazioni negli ambienti multisezione richiede la valutazione di come eseguire lo sharding tra più sezioni. In genere sono disponibili tre opzioni: parallelismo dei dati, parallelismo dei dati completamente segmentato e parallelismo delle pipeline. Sconsigliamo l'attivazione di sharding nelle 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 dello sharding all'interno di una sezione che ha funzionato per te in passato.

Ti consigliamo di iniziare con il puro parallelismo dei dati. Il parallelismo dei dati con sharding è utile per liberare spazio in memoria. Lo svantaggio è che la comunicazione tra le sezioni utilizza la rete DCN e rallenta il carico di lavoro. Utilizza il parallelismo delle pipeline solo quando necessario in base alle dimensioni del batch (come analizzato di seguito).

Quando utilizzare il parallelismo dei dati

Il parallelismo dei dati puro funzionerà bene nei casi in cui tu abbia un carico di lavoro che funziona bene, ma tu voglia migliorarne le prestazioni scalando su più sezioni.

Per ottenere una scalabilità elevata su più sezioni, il tempo necessario per eseguire la riduzione completa sulla DCN deve essere inferiore al tempo necessario per eseguire un passaggio a ritroso. La DCN viene utilizzata per la comunicazione tra le sezioni ed è un fattore limitante della velocità effettiva del carico di lavoro.

Ogni chip TPU v4 si comporta a 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 è 4 * 275 * 10¹² FLOPS / 50 Gbps = 22.000 FLOPS / bit.

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

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

Il numero di FLOPS per parametro è 6 * B, mentre il numero di FLOPS per parametro durante il passaggio a ritroso è 4 * B.

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

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

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

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

Per altre architetture di modelli, devi stimare il runtime del passaggio a ritroso per sezione (tramite sincronizzazione con un profiler o conteggio dei FLOPS). Quindi, puoi confrontarlo con il tempo di esecuzione previsto per ridurre del tutto sulla rete DCN e ottenere una buona stima di se il parallelismo dei dati ha un senso per te.

Quando utilizzare il parallelismo dei dati con sharding (FSDP)

Il parallelismo dei dati completamente partizionato (FSDP) combina il parallelismo dei dati (suddivisione dei dati tra i nodi) con lo sharding delle ponderazioni tra i nodi. Per ogni operazione nei passaggi avanti e indietro, i pesi vengono raccolti in modo che ogni sezione abbia i pesi necessari. Invece di sincronizzare i gradienti con tutto-riduzione, i gradienti vengono ridotti man mano che vengono prodotti. In questo modo, ogni sezione riceve solo i gradienti per le ponderazioni a cui è associata.

Analogamente al parallelismo dei dati, FSDP richiederà la scalabilità lineare delle dimensioni globali del batch con il numero di sezioni. Il valore FSDP ridurrà la pressione della memoria man mano che aumenti il numero di sezioni. Questo perché il numero di ponderazioni e stato di ottimizzazione per sezione diminuisce, ma ciò avviene al prezzo di un aumento del traffico di rete e della maggiore possibilità di blocco a causa di un collettivo ritardato.

In pratica, il valore FSDP nelle sezioni è l'ideale se stai aumentando il batch per sezione, archiviando più attivazioni per ridurre al minimo la rimaterializzazione durante il trasferimento a ritroso o l'aumento del numero di parametri nella rete neurale.

Le operazioni di raccolta e riduzione in FSDP funzionano in modo simile a quelle in DP, pertanto puoi determinare se il tuo carico di lavoro FSDP è limitato dalle prestazioni di DCN, come descritto nella sezione precedente.

Quando utilizzare il parallelismo delle pipeline

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

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

Il parallelismo delle pipeline deve essere usato solo se le altre strategie di parallelismo richiedono una dimensione batch globale troppo grande. Prima di provare il parallelismo delle pipeline, è utile sperimentare per verificare in modo empirico se la convergenza per campione rallenta alle dimensioni del batch necessarie per ottenere un FSDP ad alte prestazioni. FSDP tende a raggiungere un utilizzo più elevato di FLOP del modello, ma se la convergenza per campione rallenta con l'aumento delle dimensioni del batch, il parallelismo delle pipeline potrebbe essere comunque la scelta migliore. La maggior parte dei carichi di lavoro può tollerare dimensioni batch sufficientemente grandi da non trarre vantaggio dal parallelismo della pipeline, ma il carico di lavoro potrebbe essere diverso.

Se è necessario il parallelismo delle pipeline, ti consigliamo di combinarlo con il parallelismo dei dati o FSDP. Ciò ti consentirà di ridurre al minimo la profondità della pipeline aumentando al contempo la dimensione del batch per pipeline fino a quando la latenza DCN non diventerà un fattore inferiore alla velocità effettiva. Concretamente, se hai N sezioni, considera le pipeline di profondità 2 e N/2 repliche di parallelismo dei dati, poi le pipeline di profondità 4 e N/4 repliche di parallelismo dei dati e così via, finché il batch per pipeline diventa abbastanza grande da nascondere i collettivi DCN dietro l'aritmetica nel passaggio a ritroso. Ciò ridurrà al minimo il rallentamento introdotto dal parallelismo della pipeline, consentendoti di scalare oltre il limite globale di dimensione del batch.

Best practice per la multisezione

Caricamento dei dati

Durante l'addestramento, carichiamo ripetutamente batch da un set di dati per inserirli nel modello. Avere un caricatore di dati efficiente e asincrono che esegua lo sharding del batch tra gli host è importante per evitare di esaurire le TPU di lavoro. Il caricatore dati attuale in MaxText ha a ogni carico dell'host un sottoinsieme uguale degli esempi. Questa soluzione è adeguata per il testo, ma richiede un filtro all'interno del modello. Inoltre, MaxText non offre ancora la creazione di snapshot deterministica che consentirebbe all'iteratore dei dati di caricare gli stessi dati prima e dopo il prerilascio.

Checkpoint

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

Configurazioni supportate

Forme

Tutte le sezioni devono essere della stessa forma (ad esempio, dello stesso AcceleratorType). Le forme delle sezioni eterogenee non sono supportate.

Orchestrazione

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

Framework

Multislice supporta solo i carichi di lavoro JAX e PyTorch.

Parallelismo

Consigliamo agli utenti di testare la sezione multisezione con parallelismo dei dati. Per scoprire di più sull'implementazione del parallelismo delle pipeline con Multislice, contatta il tuo rappresentante dell'account Google Cloud.

Assistenza e feedback

Tutti i feedback sono ben accetti. Per condividere feedback o richiedere assistenza, contattaci utilizzando il modulo di assistenza o feedback per Cloud TPU.