Panoramica di più sezioni di Cloud TPU

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

Se il cui deployment viene eseguito in configurazioni multisezione, i chip TPU in ogni sezione comunicano tramite Inter-Chip-Interconnect (ICI). I chip TPU in diverse sezioni comunicano trasferendo i dati alle CPU (host), che a loro volta trasmettono i dati sulla 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 ottenere le massime prestazioni.

Concetti

Tipo di acceleratore

La forma di ogni sezione TPU che comprende una sezione multipla. 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 errore hardware, Cloud TPU crea una nuova sezione. Nel raro caso in cui le risorse non siano sufficienti per creare una nuova sezione, la creazione non verrà completata finché l'hardware non sarà disponibile. Dopo la creazione della nuova sezione, tutte le altre sezioni nell'ambiente Multislice verranno riavviate in modo che l'addestramento possa continuare.Con uno script di avvio configurato correttamente, lo script di addestramento può 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ù alta e velocità effettiva inferiore (rispetto a ICI) che connette le sezioni TPU in una configurazione multisezione.

Programmazione per le gang

Quando viene eseguito contemporaneamente il provisioning di tutte le sezioni TPU, garantendo il corretto provisioning di tutte le sezioni o di nessuna.

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 sui dati.

Interchip Interconnect (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 ciascuna sezione TPU di 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 delle risorse TPU, utilizzata per accodare e gestire una richiesta per un ambiente TPU a una singola sezione o a più sezioni.

Script di avvio

Uno script di avvio di Compute Engine standard che viene eseguito ogni volta che una VM viene avviata o riavviata. Per più sezioni, è specificato nella richiesta di creazione 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 composta da chip TPU. Tutti i chip in una sezione 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 definiamo worker una VM TPU.

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 da Google. Sono progettati per offrire computing veloce ed efficiente dal punto di vista energetico per attività chiave di machine learning come la moltiplicazione delle matrici.

Tipi di capacità di 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 richiesta, devi aver sottoscritto un accordo 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 risorse potrebbero essere prerilasciate per fare spazio alle richieste per un job con priorità più elevata. Utilizza 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 di TPU verrà accodata a una coda di quota on demand offerta da Cloud TPU; la disponibilità delle risorse non è garantita. Selezionato per impostazione predefinita, nessun flag necessario.

Nota: Cloud TPU limita la dimensione della sezione per tutti i tipi di acceleratori. Per informazioni sui limiti delle dimensioni delle sezioni per le richieste on demand, consulta Dimensioni predefinite delle sezioni per i codici QR on demand.

Inizia

Se non hai mai utilizzato TPU, inizia installando Google Cloud CLI e configura 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 della 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 collaudato scritto in Python e Jax. MaxText è 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 con sharding completo (FSDP) e parallelismo dei tensori. Il codice scala da ambienti a sezione singola a ambienti multisettore.

Parallelismo ICI

ICI è l'interconnessione ad alta velocità che collega le TPU in una singola 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 ogni metodo di parallelismo.

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

La tabella seguente mostra esempi di input utente per il parallelismo di 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 ici_data_parallelism nella maggior parte dei casi deve essere lasciato impostato su 1 perché la rete ICI è abbastanza veloce da preferire quasi sempre FSDP al parallelismo dei dati.

Questo esempio presuppone che tu abbia familiarità con l'esecuzione di codice su una singola sezione TPU, ad esempio in Eseguire 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 vuota una password (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 seguente comando:

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

   your-qr-id

   Una stringa definita dall'utente che identifica la richiesta 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 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 QR.

   reserved

   Utilizza la quota prenotata durante la creazione delle sezioni.

   best-effort

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

   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 finché il QR non sia nello stato ACTIVE, il che significa 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 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 sulla sezione TPU. L'esecuzione dello script di configurazione richiede qualche minuto.

   $ bash setup.sh
   

8. Esegui questo comando per eseguire shardings.py sulla sezione 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 utilizzo impressionante di FLOP superiore al 90%. In questo caso, abbiamo selezionato approssimativamente il batch massimo che rientra nella memoria a larghezza di banda elevata (HBM) della TPU.

9. Scopri altre strategie dello sharding tramite 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 questi passaggi di pulizia dall'ambiente in cui hai configurato la sezione (prima esegui exit per uscire dalla sessione SSH). Il completamento dell'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 a utilizzare la tua quota di QR, che viene monitorata separatamente da quella di TPU.

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

Partizionamento orizzontale multisezione utilizzando il parallelismo della DCN

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

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

I valori di questi parametri devono essere vincolati in modo che dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism sia uguale al numero di 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 DCN non è adatta a questo sharding. Per i carichi di lavoro LLM tipici sui chip v4, anche dcn_fsdp_parallelism deve essere impostato su 1 e, di conseguenza, dcn_data_parallelism deve essere impostato sul numero di sezioni, ma questo dipende dall'applicazione.

Aumentando il numero di sezioni (supponendo di mantenere costanti le dimensioni della sezione e il batch per sezione), 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 o eseguendo shardings.py su ogni VM TPU. In questo caso utilizziamo multihost_runner.py. I seguenti passaggi sono molto simili a quelli di Introduzione: esperimenti rapidi su più sezioni del repository MaxText, ma in questo caso eseguiamo shardings.py anziché l'LLM più complesso in train.py.

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

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

In questo esempio, 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. Consigliamo di lasciare una password vuota (premi Invio due volte 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 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, --best_effort oppure 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 Multislice v4. Per utilizzare v5e, specifica v5e accelerator-type (ad esempio v5litepod-16) e 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 di comando

   your-qr-id

   Una stringa definita dall'utente che identifica la richiesta 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 di Cloud TPU.

   node-count

   Il numero di sezioni da creare.

   node-prefix

   Il prefisso utilizzato per generare i nomi di ogni sezione. Per ogni sezione viene aggiunto un numero al prefisso. 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 [predefinita].

6. Quando inizia il provisioning del codice QR, il completamento potrebbe richiedere fino a cinque minuti, a seconda delle dimensioni del codice QR. Attendi fino a quando la risorsa in coda (QR) non è 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
   

   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 del tuo account Google Cloud se lo stato QR è nello 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. Al termine, pulisci le TPU e il QR. Il completamento dell'eliminazione richiederà dai due ai 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 Multislice, apporta le seguenti modifiche al codice:

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

Queste dovrebbero essere le uniche modifiche al codice necessarie per il passaggio a più sezioni. Per ottenere prestazioni elevate, la rete DCN deve essere mappata su assi paralleli dei dati, paralleli dei dati completamente suddivisi o paralleli della pipeline. Le considerazioni sulle prestazioni e le strategie di sharding sono discusse più in dettaglio in 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 nell'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 con più sezioni

Per aumentare la velocità lineare, aggiungi nuove sezioni delle stesse dimensioni di quella esistente. Ad esempio, se utilizzi una sezione v4-512, l'opzione Multisezione raggiungerà circa il doppio delle prestazioni aggiungendo una seconda sezione v4-512 e raddoppiando la dimensione globale del batch. Per ulteriori informazioni, consulta la sezione Sharding con più sezioni per ottenere le massime prestazioni.

Esecuzione del job su più sezioni

Esistono tre diversi approcci per l'esecuzione del carico di lavoro personalizzato in un ambiente Multislice:

1. Utilizzando lo script di esecuzione della sperimentazione, multihost_runner.py
2. Utilizzando lo script runner di produzione, multihost_job.py
3. Adottare un approccio manuale

Script runner della sperimentazione

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

Poiché multihost_runner.py mantiene connessioni SSH permanenti, è adatto solo per esperimenti di dimensioni modeste e relativamente brevi. Puoi adattare i passaggi del tutorial multihost_runner.py al tuo carico di lavoro e alla configurazione hardware.

Script runner di produzione

Per i job di produzione che richiedono resilienza contro guasti hardware e altre prerilazioni, è preferibile integrare direttamente l'API Create Queued Resource. Come esempio pratico, 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 il 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 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 carico di lavoro personalizzato nella configurazione di Multislice. Tuttavia, se preferisci eseguire il provisioning e gestire il tuo ambiente direttamente utilizzando i comandi QR, consulta Gestire un ambiente multisezione.

Gestire un ambiente Multislice

Per eseguire manualmente il provisioning e la gestione dei codici QR senza utilizzare gli strumenti forniti 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

Salvaguardie	Descrizione
your-qr-id	L'ID del QR assegnato dall'utente.
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 viene aggiunto alla richiesta di creazione, può essere eseguito uno script di avvio ogni volta che viene eseguito il provisioning o il riavvio di 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 di comando

your-qr-id

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

project

Una stringa definita dall'utente che identifica la richiesta 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 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 prenotata durante la creazione delle sezioni.

spot

Utilizza la quota VM spot durante la creazione delle sezioni.

Assicurati di disporre della quota corrispondente prima di selezionare --reserved, --spot o la quota on demand predefinita. Per informazioni sui tipi di quota, consulta 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 runtime VM TPU
• your-network-name - Facoltativo, una rete a cui verrà collegato il QR
• your-subnetwork-name - Facoltativo, una subnet a cui verrà collegato il QR
• example-tag-1 - Facoltativo, una stringa tag arbitraria
• your-startup-script: uno script di avvio che verrà eseguito quando viene allocato il 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 sull'API REST Queued Resource per conoscere tutte le opzioni disponibili.

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

Per utilizzare la capacità Spot, sostituisci:

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

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

Invia la richiesta di creazione 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 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 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 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, è stato in coda e inizierà il provisioning quando ci saranno risorse sufficienti. Se il QR è nello stato FAILED, il motivo dell'errore sarà visualizzato nell'output. Per ulteriori informazioni su altri possibili stati, consulta la guida dell'utente sulle risorse in coda.

Una volta completata l'operazione, utilizza la descrizione dei QR per monitorare le fasi del QR.

In uno scenario raro, 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 contatta il team di Cloud TPU per risolvere il problema.

SSH e installazione delle dipendenze

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

  $ gcloud compute tpus queued-resources ssh ${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 ai nodi in QR tramite SSH. Il comando viene suddiviso in gruppi di quattro e viene inviato contemporaneamente. Il batch successivo di comandi viene inviato al termine dell'esecuzione del batch corrente. In caso di errore con uno dei comandi, l'elaborazione si interrompe e non vengono inviati ulteriori batch. Per ulteriori informazioni, consulta il riferimento per l'API delle risorse in coda. Se il numero di sezioni che utilizzi supera il limite di thread del tuo computer locale (chiamato anche limite di batch), verrà visualizzato un deadlock. Ad esempio, supponiamo che il limite per il 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. Lo script di addestramento verrà eseguito 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ò completarsi 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 iniziare a eseguire lo script di addestramento sul batch rimanente di 36 sezioni. Devi eseguire correttamente la pipeline dei flussi stdout e stderr quando esegui i comandi in background. Per aumentare il parallelismo all'interno dello stesso QR, puoi selezionare sezioni specifiche usando 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 Multisezione. Per farlo, su ogni sezione, esegui:

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

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

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, utilizza l'API QR per la descrizione. 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 QR. Per ulteriori informazioni sui possibili stati dei codici QR, consulta la sezione 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 in ogni sezione su SSH ed eseguendo questo 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

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

I ripristini di tutte le VM vengono eseguiti in parallelo e il completamento di un'operazione ResetQueuedResource richiede uno o due minuti. Per richiamare l'API, usa questo 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. Il completamento dell'eliminazione richiederà da due a cinque minuti e può essere eseguito 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 senza intervento della sezione interessata e il ripristino di tutte le sezioni in seguito. La sezione interessata viene sostituita con una nuova sezione e le sezioni rimanenti altrimenti integre vengono reimpostate. Se non è disponibile alcuna capacità per l'allocazione di una sezione sostitutiva, l'addestramento si interrompe.

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

Il seguente script di avvio (utilizzato in Crea QR) consente di recuperare 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 a più sezioni. Per maggiori informazioni, consulta la sezione Profilazione dei programmi JAX.

Ottimizza la formazione

Sharding con più sezioni per ottenere le massime prestazioni

Per ottenere le massime prestazioni in ambienti multisettore è necessario considerare come eseguire lo sharding su più sezioni. In genere sono disponibili tre opzioni (parallelismo dei dati, parallelismo dei dati completamente segmentato e parallelismo delle pipeline). Sconsigliamo di partizionare le attivazioni tra le varie dimensioni del modello (a volte chiamato parallelismo tensore) perché richiede un'eccessiva 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.

Consigliamo di iniziare con il parallelismo puro dei dati. L'uso del parallelismo dei dati completamente con partizionato è utile per liberare memoria. Lo svantaggio è che la comunicazione tra le sezioni utilizza la rete DCN e rallenterà il carico di lavoro. Utilizza il parallelismo della 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 funziona bene nei casi in cui hai un carico di lavoro che funziona bene, ma vorresti migliorarne le prestazioni scalando su più sezioni.

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

Ogni chip TPU v4 ha 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 è 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 utilizzi più di due sezioni, il compilatore eseguirà un'operazione di shuffling completa di riduzione completa e utilizzerai fino a 64 bit di larghezza di banda DCN per ciascun parametro per passaggio. La quantità di FLOPS necessaria per ogni parametro varia a seconda del modello. Nello specifico, per i modelli linguistici basati su Transformer, il numero di FLOPS necessari per un passaggio in avanti e all'indietro è pari a 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, 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 l'intensità 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 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 i modelli di linguaggio basati su Transformer. Poiché la rete DCN può abbandonare brevemente i pacchetti, è preferibile 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 del modello, dovrai stimare il runtime del passaggio a ritroso per sezione (temporandolo con un profiler o conteggiando i FLOPS). Quindi puoi confrontare questo valore con il tempo di esecuzione previsto per ridurre tutti su DCN e ottenere una stima approssimativa di se il parallelismo dei dati ha senso per te.

Quando utilizzare il parallelismo dei dati con sharding completo (FSDP)

Il parallelismo dei dati con sharding completo (FSDP) combina il parallelismo dei dati (partizionando i dati tra i nodi) con lo sharding dei pesi tra i nodi. Per ogni operazione nei passaggi avanti e indietro, i pesi vengono tutti raccolti in modo che ogni sezione abbia i pesi necessari. Invece di sincronizzare i gradienti con all-Reduce, i gradienti vengono ridotti man mano che vengono prodotti. In questo modo, ogni sezione riceve solo i gradienti per le ponderazioni di cui è responsabile.

Analogamente al parallelismo dei dati, FSDP richiederà di scalare la dimensione del batch globale in modo lineare in base al numero di sezioni. FSDP ridurrà la pressione di memoria man mano che si aumenta il numero di sezioni. Questo perché il numero di ponderazioni e stato di ottimizzazione per sezione diminuisce, ma lo fa a scapito dell'aumento del traffico di rete e della maggiore possibilità di blocco dovuto a un collettivo in ritardo.

In pratica, FSDP su più sezioni è l'ideale se vuoi aumentare il batch per sezione, archiviare più attivazioni per ridurre al minimo la rimaterializzazione durante il passaggio a ritroso o se 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 carico di lavoro FSDP è limitato dalle prestazioni DCN, come descritto nella sezione precedente.

Quando utilizzare il parallelismo della pipeline

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

1. Si aprirà la "bolla della pipeline" in cui i chip sono inattivi perché sono in attesa di dati.
2. Richiede micro-batching, che riduce la dimensione effettiva del batch, l'intensità aritmetica e, in ultima analisi, l'utilizzo del modello 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 sperimentare con criteri empirici se la convergenza per campione rallenta alla dimensione del batch necessaria per ottenere un FSDP ad alte prestazioni. FSDP tende a ottenere un maggiore utilizzo del FLOP nel modello, ma se la convergenza per campione rallenta man mano che la dimensione del batch aumenta, 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 carico di lavoro potrebbe essere diverso.

Se è necessario il parallelismo della 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 di ciascuna pipeline finché la latenza DCN non sarà un fattore di velocità effettiva. Concretamente, se hai N sezioni, considera le pipeline di profondità 2 e le repliche N/2 del parallelismo dei dati, quindi le pipeline di profondità 4 e N/4 replica del parallelismo dei dati e così via, fino a quando il batch per pipeline non diventa abbastanza grande da consentire ai collettivi DCN di essere nascosti dietro l'aritmetica nel passaggio a ritroso. Questo ridurrà al minimo il rallentamento introdotto dal parallelismo della pipeline, consentendoti allo stesso tempo di scalare oltre il limite globale della dimensione del batch.

Best practice per più sezioni

Caricamento dei dati

Durante l'addestramento, carichiamo ripetutamente i batch da un set di dati da inserire nel modello. Per evitare di esaurire le TPU di lavoro, è importante disporre di un caricatore di dati asincrono efficiente che esegue lo sharding del batch tra gli host. L'attuale caricatore di dati in MaxText prevede che ogni host carichi un sottoinsieme uguale di esempi. Questa soluzione è adeguata per il testo, ma richiede un controllo all'interno del modello. Inoltre, MaxText non offre ancora snapshot deterministici che consentirebbero all'iteratore di dati di caricare gli stessi dati prima e dopo il prerilascio.

Checkpoint

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

Configurazioni supportate

Forme

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

Orchestrazione

L'orchestrazione è supportata con GKE. Per maggiori informazioni, consulta la pagina 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 scoprire di più sull'implementazione del parallelismo della 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 feedback o assistenza per Cloud TPU.