Panoramica multisezione di Cloud TPU [anteprima pubblica]

Anteprima

Questo prodotto o funzionalità è soggetto ai "Termini delle Offerte pre-GA" nella sezione Termini generali del servizio dei Termini specifici del servizio. I prodotti e le funzionalità pre-GA sono disponibili "così come sono" e potrebbero avere un supporto limitato. Per saperne di più, consulta le descrizioni della fase di lancio.

Cloud TPU multislice è una tecnologia di scalabilità delle prestazioni full stack che consente a un job di addestramento di utilizzare più sezioni TPU all'interno di un singolo pod o su sezioni in più pod con 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 prestazioni migliori. Tuttavia, più sezioni più piccole sono più facilmente disponibili, consentendo un tempo di avvio più rapido quando la funzionalità Multislice viene utilizzata con sezioni più piccole.

Quando il deployment è stato eseguito in configurazioni multislice, i chip TPU in ogni sezione comunicano tramite l'interconnessione tra chip (ICI). I chip TPU in sezioni diverse comunicano trasferendo i dati alle CPU (host), che a loro volta li trasmettono sulla rete di data center (DCN).

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

Concetti

AcceleratorType

La forma di ogni sezione di TPU che comprende una sezione multislice. Ogni sezione in una richiesta multislice ha lo stesso AcceleratorType. AcceleratorType è composto da un tipo di TPU (v2, v3, v4, v5e e così via) 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. Nei rari casi in cui le risorse non siano sufficienti per creare una nuova sezione, la creazione non verrà completata finché l'hardware non sarà disponibile. 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ò essere riavviato automaticamente senza l'intervento dell'utente, con caricamento e ripristino 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 superiore e velocità effettiva inferiore (rispetto a ICI) che connette le sezioni TPU in una configurazione multisezione.

Programmazione gang

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

Host

Un host è un computer fisico che esegue le VM. Un host può eseguire al massimo quattro VM alla volta. 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)

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

Multistrato

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

Nodo

Nel contesto multislice, il termine nodo si riferisce a una singola sezione TPU. A ogni sezione TPU in un multislice 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 single-slice o Multislice.

Script di avvio

Uno script di avvio di Compute Engine standard che viene eseguito a ogni avvio o riavvio di una VM. Per il multislice, viene specificato nella richiesta di creazione della risorsa in coda. 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 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.

Tensore

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 il computing più veloce ed efficiente dal punto di vista energetico per le attività chiave di machine learning, come la moltiplicazione di matrici.

Tipi di capacità di Cloud TPU

Le TPU possono essere create a partire da tre 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 recuperata, devi avere un contratto di prenotazione con Google. Utilizza il flag --reserved durante la creazione delle risorse.
• Prerilasciabile: ha come target la quota prerilasciabile. Le tue risorse potrebbero essere prerilasciate per fare spazio alle richieste di un job con priorità più elevata. Utilizza il flag --best-effort 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 di quote on demand offerta da Cloud TPU. Selezionata per impostazione predefinita, non sono necessari flag.

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

Inizia

Se non hai mai utilizzato le TPU, per prima cosa installa i componenti gcloud e abilita l'API Cloud TPU. Quindi, configura il comando gcloud per impostare il progetto. Segui tutte le istruzioni riportate in Configura un progetto e interrompi quando raggiungi la sezione Gestione delle TPU.

Per utilizzare la sezione multisezione, 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 team Cloud TPU. In GA, la richiesta di creazione di risorse in coda potrà accedere a tutte le prenotazioni e gli utenti non dovranno trasferire alcuna capacità prima di utilizzare le risorse multislice o in coda.

Esempio introduttivo

Questo tutorial utilizza il codice all'interno del repository GitHub di MaxText. shardings.py (all'interno del repository MaxText) è progettato per semplificare la sperimentazione di diverse opzioni di parallelismo. È il punto di partenza migliore se vuoi esplorare il parallelismo dei dati, il parallelismo dei dati completamente segmentato (FSDP) e/o il parallelismo dei tensori. Consente di estendere facilmente il codice alle impostazioni sia single-slice che multislice.

Non è specifico del multislice e dovrebbe essere disponibile altrove.

Partizionamento orizzontale a sezione singola utilizzando il parallelismo ICI

Lo sharding ICI corrisponde allo sharding all'interno di ogni sezione. shardings.py fornisce tre input per il parallelismo ICI:

• ici_data_parallelism
• ici_fsdp_parallelism
• ici_tensor_parallelism

I valori specificati per questi parametri determinano il numero di shard per ogni metodo di parallelismo.

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

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

	parallelism_ici_data	parallelismo_ici_fsdp	parallelismo_ici_tensor
FSDP a 4 vie	1	4	1
Parallelismo tensore a 4 vie	1	1	4
FSDP a 2 vie + parallelismo tensore a 2 vie	1	2	2

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

Questo esempio presuppone che tu abbia familiarità con l'esecuzione del codice su una singola sezione TPU, ad esempio in Eseguire un calcolo su una VM Cloud TPU utilizzando JAX. Questo esempio è tratto da quella guida per eseguire shardings.py su una singola sezione.

1. Configura l'ambiente:

   $ QUEUED_RESOURCE_ID=user-defined-queued-resource-id
   PROJECT=your-gcp-project-name
   ZONE=us-central2-b
   RUNTIME_VERSION=tpu-ubuntu2204-base
   
   gcloud auth login
   gcloud config set project ${PROJECT}
   gcloud config set compute/zone ${ZONE}
   

2. Crea chiavi SSH per gcloud. Ti consigliamo di lasciare una password vuota (premi Invio due volte dopo aver eseguito il comando riportato di seguito). 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 $QUEUED_RESOURCE_ID \
   --accelerator-type v4-8 \
   --runtime-version tpu-ubuntu2204-base \
   --node-id $QUEUED_RESOURCE_ID \
   [--reserved |--best-effort]
   

   Nota: attualmente, l'interfaccia a riga di comando gcloud non supporta tutte le opzioni di creazione di risorse in coda, come i tag. Per saperne di più, vedi Creare risorse in coda.

4. Attendi che lo stato del codice QR sia ACTIVE, il che significa che i nodi worker sono in stato READY. Una volta avviato il provisioning del codice QR, il completamento può richiedere da uno a cinque minuti, a seconda delle dimensioni del codice QR. Puoi controllare lo stato utilizzando il seguente comando:

   $ gcloud alpha compute tpus queued-resources \
   list --filter=$QUEUED_RESOURCE_ID
   

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

   $ gcloud compute tpus tpu-vm ssh $QUEUED_RESOURCE_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.

   $ bash setup.sh # These installations will take
   a few minutes.
   

8. Esegui il comando seguente per eseguire shardings.py sulla sezione TPU.

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

   Puoi visualizzare i risultati nei log. Le tue TPU dovrebbero raggiungere circa 260 TFLOP/s o un utilizzo impressionante di oltre il 90%di FLOP. In questo caso, abbiamo selezionato approssimativamente il batch massimo adatto alla memoria HBM (High Bandwidth Memory).

9. Non esitare a 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 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 QR, che viene monitorata separatamente dalla quota TPU.

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

Partizionamento orizzontale con più sezioni utilizzando il parallelismo DCN

Lo script shardings.py accetta tre parametri che specificano il parallelismo DCN, corrispondente al numero di shard di ciascun tipo di parallelismo dei dati:

• Parallelismo_dati_dcn
• Parallelismo_dcn_fsdp
• Parallelismo_tensor_dcn

I valori di questi parametri devono essere vincolati in modo che dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism corrisponda al numero di sezioni.

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

	Parallelismo_dati_dcn	Parallelismo_dcn_fsdp	Parallelismo_tensor_dcn	N. di sezioni
Parallelismo dei dati bidirezionale	2	1	1	2

dcn_tensor_parallelism deve essere sempre impostato su 1 perché la rete DCN non è adatta a questo sharding. Per i carichi di lavoro LLM tipici su chip v4, anche dcn_fsdp_parallelism deve essere impostato su 1 e quindi 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 rimangano costanti), aumenti la quantità di parallelismo dei dati. I passaggi qui sono elencati anche nel file README di MaxText.

Esecuzione di shardings.py in un ambiente multisezione

Puoi eseguire shardings.py in un ambiente multislice utilizzando multihost_runner.py o eseguendo shardings.py su ogni VM TPU. Qui usiamo multihost_runner.py. I passaggi seguenti sono molto simili a quelli di Guida introduttiva: esperimenti rapidi su più sezioni del repository MaxText, tranne per il fatto che qui eseguiamo shardings.py anziché il più complesso LLM 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 sulla quale viene eseguito lo script multihost_runner.py. Utilizziamo il termine workers per indicare le VM TPU che compongono le sezioni. Puoi eseguire multihost_runner.py su una macchina locale o su qualsiasi VM di Compute Engine nello stesso progetto 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 facilmente l'esempio per eseguirlo su più TPU.

Configura l'ambiente

1. Clona MaxText sulla tua macchina runner.
2. Vai alla directory repo.
3. Configura l'ambiente:

   $ export QUEUED_RESOURCE_ID=user-defined-queued-resource-id
   $ export PROJECT=your-gcp-project-name
   $ export ZONE=us-central2-b
   $ export RUNTIME_VERSION=tpu-ubuntu2204-base
   $ gcloud config set project ${PROJECT}
   $ gcloud config set compute/zone ${ZONE}
   

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

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

5. Aggiungi una variabile di ambiente per impostare il conteggio delle sezioni TPU su 2.

   $ export SLICE_COUNT=2
   

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

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

   Il comando seguente mostra come creare una TPU multislice v4. TPU v5e è supportata anche nell'anteprima pubblica multisezione. Per utilizzare la v5e, specifica una v5e accelerator-type (ad esempio v5e-16) e la v5e runtime-version (v2-alpha-tpuv5-lite).

   $ gcloud alpha compute tpus queued-resources \
   create $QUEUED_RESOURCE_ID \
   --accelerator-type=v4-16 \
   --runtime-version=tpu-ubuntu2204-base \
   --node-count=$SLICE_COUNT \
   --node-prefix=$QUEUED_RESOURCE_ID \
   [--reserved|--best-effort]
   

   Descrizioni flag comando

   QUEUED_RESOURCE_ID

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

   accelerator-type

   Il tipo di Cloud TPU da creare.

   runtime-version

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

   node-count

   Il numero di sezioni da creare.

   node-prefix

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

   reserved

   Utilizza la quota riservata durante la creazione delle sezioni.

   best-effort

   Utilizza la quota secondo il criterio del "best effort" per la creazione delle sezioni [valore predefinito].

7. Attendi che la risorsa in coda (QR) sia nello stato ACTIVE, il che significa che i nodi worker sono in stato READY. Una volta avviato il provisioning del codice QR, il completamento può richiedere da uno a cinque minuti, a seconda delle dimensioni del codice QR. Puoi controllare lo stato utilizzando il seguente comando:

   $ gcloud alpha compute tpus queued-resources list \
   --filter=$QUEUED_RESOURCE_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
   ...
   

   Contattaci se lo stato del codice QR è nello stato ACCEPTED (WAITING_FOR_RESOURCES) o PROVISIONING per più di 15 minuti.

8. Installare le dipendenze.

   $ pip3 install absl-py
   $ python3 multihost_runner.py \
     --TPU_PREFIX=$QUEUED_RESOURCE_ID \
     --COMMAND="bash setup.sh"
   

9. Esegui shardings.py su ciascun worker utilizzando multihost_runner.py.

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

   Nei file di log vedrai circa 230 TFLOP/s di prestazioni.

10. Al termine, pulisci le TPU e il QR. L'eliminazione richiederà da due a cinque minuti e può essere eseguita in background con il flag facoltativo --async.

MaxText: un LLM multisezione, scalabile e ad alte prestazioni

Per iniziare, abbiamo fornito train.py, un'implementazione LLM leggibile, scalabile e ad alte prestazioni, pronta all'uso. Ti invitiamo a eseguirlo autonomamente e a utilizzare il codice nelle tue implementazioni. MaxText contiene i propri risultati README e prestazioni.

Sono inclusi i risultati relativi al rendimento di MaxText, copiati dal file README di MaxText. MaxText contiene anche tutto ciò che consente di riprodurre i risultati. Qui supponiamo i FLOPS a 6 * BATCH * PARAMETRI.

Modello di 22 miliardi di parametri

Articoli di ferramenta	TFLOP/sec/chip	Utilizzo del modello FLOP
1x v4-128	165	60%
2x v4-128	158	57,4%
4x v4-128	157	57,1%

Modello di 52 miliardi di parametri

Articoli di ferramenta	TFLOP/sec/chip	Utilizzo del modello FLOP (MFU)
1x v4-384	154	56,0%
2x v4-384	162	58,9%

I miglioramenti da 1x v4-384 a 2x v4-384 sono guidati da un'anomalia del compilatore. Anche se di dimensioni diverse, potresti notare una riduzione del numero di MFU con l'aumentare del numero di sezioni.

Scalabilità di un carico di lavoro a più sezioni

Modifiche al codice

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 durante la creazione del mesh.

• Imposta alcune variabili di ambiente:

  os.environ["JAX_USE_PJRT_C_API_ON_TPU"] = "1"
  os.environ["LIBTPU_INIT_ARGS"] = "--xla_tpu_enable_megascale_barrier=true"
  

  Nota: LIBTPU_INIT_ARGS è un elenco separato da spazi, ad esempio LIBTPU_INIT_ARGS="--flag1=val1 --flag2=var2".

Queste dovrebbero essere le uniche modifiche al codice necessarie per il passaggio alla sezione multisezione. Per ottenere prestazioni elevate, la rete DCN deve essere mappata su assi paralleli di dati paralleli, completamente suddivisi in shard o paralleli della pipeline. Le considerazioni sulle prestazioni e le strategie di suddivisione in segmenti sono discusse in maggiore dettaglio nell'articolo Suddivisione in partizionamento con multislice per ottenere le massime prestazioni.

Per convalidare che il tuo codice può accedere a tutti i dispositivi, puoi dichiarare che len(jax.devices()) corrisponde 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 di quelle esistenti. Ad esempio, se al momento utilizzi una sezione v4-512, Multislice raggiungerà circa il doppio delle prestazioni aggiungendo una seconda sezione v4-512 e raddoppiando le dimensioni del batch globale. Per ulteriori informazioni, consulta la sezione Suddivisione con più sezioni per ottenere le massime prestazioni.

Esecuzione del job su più sezioni

Esistono tre diversi approcci all'esecuzione di un carico di lavoro personalizzato in un ambiente multislice:

1. Utilizzo dello script dell'esecuzione della sperimentazione, multihost_runner.py
2. Utilizzando lo script dell'esecutore di produzione, multihost_job.py
3. Utilizzo di un approccio manuale

Script esecutore esperimento

Per una sperimentazione rapida, forniamo multihost_runner.py (ospitato in MaxText), un semplice script per distribuire il codice in un ambiente multislice esistente, eseguire il comando su ciascun host, copiare i log e monitorare lo stato di errore di ciascun comando. Lo script multihost_runner.py è documentato in modo più completo nel README di MaxText.

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

Script del runner di produzione

Per i job di produzione che richiedono resilienza contro guasti hardware e altre prerilasci, è preferibile l'integrazione diretta con l'API Create Queued Resource. Come esempio funzionante, forniamo multihost_job.py (ospitato all'interno di MaxText), che attiva la chiamata all'API Created Queued Resource con lo script di avvio appropriato per eseguire l'addestramento e riprenderlo in caso di prerilascio. Lo script multihost_job.py è più documentato in MaxText README.

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 desiderata. Tuttavia, se preferisci eseguire il provisioning e gestire l'ambiente utilizzando direttamente i comandi delle risorse in coda, fai riferimento a Gestire un ambiente multislice.

Gestire un ambiente multisezione

Per eseguire il provisioning e gestire manualmente le risorse in coda senza utilizzare gli strumenti forniti nel repository MaxText, leggi le sezioni seguenti.

Crea risorse in coda

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

  $ export QUEUED_RESOURCE_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
ID_RISORSA_CODA	L'ID assegnato dall'utente della risorsa in coda.
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
TIPO_ACCELERATOR	versione 4-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 32 sezioni, ma può essere aumentato durante l'anteprima pubblica.
AVVIO_SCRIPT	Se viene aggiunto alla richiesta di creazione, uno script di avvio può essere eseguito ogni volta che viene eseguito il provisioning/riavvio di una sezione TPU e se viene riparata o reimpostata una sezione TPU.

Crea una richiesta di risorse in coda utilizzando gcloud

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

Assicurati di disporre della quota corrispondente prima di selezionare --reserved o --best_effort oppure la quota on demand predefinita. Per informazioni sui tipi di quote, consulta i criteri per le quote.

Crea una richiesta di risorse in coda utilizzando curl

Crea un file denominato queuedresourcereq.json e copia al suo interno il seguente JSON.

{
  "guaranteed": { "reserved": true },
  "tpu": {
    "node_spec": [
      {
        "parent": "projects/<PROJECT>/locations/<ZONE>",
        "node": {
          "accelerator_type": "<ACCELERATOR_TYPE>",
          "runtime_version": "<RUNTIME_VERSION>",
          "network_config": {
            "network": "<NETWORK_NAME>",
            "subnetwork": "<SUBNETWORK_NAME>",
            "enable_external_ips": true
          },
     "tags" : ["<EXAMPLE_TAG_1>","<EXAMPLE_TAG_2>"]
          "metadata": {
            "startup-script": "<STARTUP_SCRIPT>"
          }
        },
        "multi_node_params": {
          "node_count": <SLICE_COUNT>,
          "node_id_prefix": "<QUEUED_RESOURCE_ID>"
        }
      }
    ]
  }
}

Sostituisci gli elementi in "< >" con i valori appropriati. Per ulteriori informazioni, consulta la documentazione sull'API REST Queued Resource per tutte le opzioni disponibili.

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

Per utilizzare la capacità interrompibile, sostituisci:

"guaranteed": { "reserved": true }

con

"best_effort": {}

In alternativa, 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/${PROJECT}/locations/${ZONE}/queuedResources\?queued_resource_id\=${QUEUED_RESOURCE_ID}

Dovresti vedere un output simile al seguente con il nome dell'operazione per la richiesta di creazione del codice QR:

{
  "name": "projects/123456789012/locations/us-central2-b/operations/operation-0000000000000-0000000000000-00000000-00000000",
  "metadata": {
    "@type": "type.googleapis.com/google.cloud.common.OperationMetadata",
    "createTime": "2023-03-09T07:35:11.142588587Z",
    "target": "projects/123456789012/locations/us-central2-b/queuedResources/your-queued-resource-id",
    "verb": "create",
    "cancelRequested": false,
    "apiVersion": "v2alpha1"
  },
  "done": false
}

Puoi utilizzare l'operazione name per ottenere lo stato del comando QR create.

Recupera operazione

Per esaminare i dettagli dell'operazione per creare un codice QR, utilizza:

  $ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2alpha1/projects/$PROJECT/locations/$ZONE/operations/operation-00000000000000-0000000000000-00000000-00000000

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

{
  "name": "projects/project/locations/zone/operations/operation-00000000000000-0000000000000-00000000-00000000",
  "metadata": {...},
  "done": true,
  "response": {
    "@type": "type.googleapis.com/google.cloud.tpu.v2alpha1.QueuedResource",
    ...
    "state": {
      "state": "WAITING_FOR_RESOURCES"
    }
  }
}

Se il codice QR è stato creato correttamente ("done = true"), lo stato all'interno del campo response sarà ACCEPTED (WAITING_FOR_RESOURCES) o FAILED. Se lo stato del codice QR è ACCEPTED (WAITING_FOR_RESOURCES), significa che è stato accodato e inizierà il provisioning quando saranno disponibili risorse sufficienti. Se il QR è nello stato FAILED, il motivo dell'errore sarà nell'output. Per ulteriori informazioni sugli altri stati possibili, consulta la guida dell'utente sulle risorse in coda.

Una volta completata l'operazione, utilizza la descrivi le risorse in coda per monitorare le fasi del codice QR.

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

SSH e installazione delle dipendenze

Esegui codice JAX su sezioni di pod di 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 ${QUEUED_RESOURCE_ID} \
  --zone ${ZONE} \
  --node=all \
  --worker=all \
  --command="some-command"
  --batch-size=4

Questo comando gcloud invia il comando specificato a tutti i worker e nodi in QR tramite SSH. Il comando è suddiviso in gruppi di quattro e viene inviato contemporaneamente. Il batch di comandi successivo viene inviato quando il batch corrente completa l'esecuzione. In caso di errore con uno dei comandi, l'elaborazione si interrompe e non vengono inviati ulteriori batch. Per saperne di più, consulta il riferimento per l'API della risorsa in coda.

Configurazione della rete

Assicurati che le sezioni TPU possano comunicare tra loro eseguendo i passaggi riportati di seguito.

Installa JAX su ognuno delle sezioni - Esegui il codice JAX sulle sezioni di pod di TPU.

Dichiara che len(jax.devices()) è uguale al numero di chip nel tuo ambiente Multislice. Per farlo, esegui su ogni sezione:

export JAX_USE_PJRT_C_API_ON_TPU=1 && python3 -c 'import jax; print(jax.devices())'

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

Elenca risorse in coda

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

$ gcloud alpha compute tpus queued-resources list

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

Descrivi le risorse in coda

Per visualizzare la configurazione e lo stato dettagliati di un codice QR, utilizza l'API Describe QR. Puoi chiamare questa API utilizzando gcloud o curl.

In uso: gcloud:

$ gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}
...state:
 state: ACTIVE
...

In uso: curl:

$ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2alpha1/projects/${PROJECT}/locations/${ZONE}/queuedResources/${QUEUED_RESOURCE_ID}
{
  "name": your-queued-res,
  "tpu": {
    "nodeSpec": [
      {
        ... // node 1
      },
      {
        ... // node 2
      },
      ...
    ]
  },
  ...
  "state": "ACTIVE"
  }
}

"state" rappresenta lo stato di un QR. Gli stati sono definiti come:

ACCEPTED (deprecato, sostituito da WAITING_FOR_RESOURCES)	Il comando QR create è stato ricevuto e inizierà il provisioning, non appena saranno disponibili risorse.
PROVISIONING	È in corso il provisioning delle sezioni TPU.
ACTIVE	È stato eseguito il provisioning di tutte le TPU, che sono pronte all'uso. Se viene fornito uno script di avvio, verrà avviato l'esecuzione su tutte le TPU, quando lo stato QR passa a ACTIVE. Tieni presente che al momento è presente un piccolo bug: lo stato QR continua a essere visualizzato come ACTIVE, mentre alcune sezioni vengono eliminate. In questo caso, il QR non funzionerà come previsto e tutte le risorse devono essere pulite.
FAILED	Impossibile eseguire il provisioning delle sezioni.
SUSPENDING	Una o più sezioni sono in fase di eliminazione.
SUSPENDED	Tutte le sezioni sottostanti vengono eliminate, ma il codice QR rimane intatto finché non vengono eliminate esplicitamente. Al momento non è possibile ripristinare un codice QR sospeso, che deve essere eliminato.
DELETING	È in corso l'eliminazione del codice QR.

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 tramite SSH ed eseguendo il seguente comando su tutti gli host. In generale, devi eseguire un comando come il seguente una volta per ogni sezione.

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

Reimpostazione delle risorse in coda

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

Le reset di tutte le VM vengono eseguite in parallelo e il completamento di un'operazione ResetQueuedResource richiede da 1 a 2 minuti. Per richiamare l'API, utilizza il seguente comando:

$ gcloud alpha compute tpus queued-resources reset $QUEUED_RESOURCE_ID

Eliminazione delle risorse in coda

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

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

Ripristino automatico da errori

In caso di interruzione, Multislice offre una riparazione senza intervento della sezione interessata e la reimpostazione di tutte le sezioni in un secondo momento. La sezione interessata viene sostituita con una nuova e le altre sezioni altrimenti integre vengono reimpostate. Se non è disponibile 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 Creare risorse in coda) consente di recuperare automaticamente dagli errori e di riprendere l'addestramento dai checkpoint archiviati in un bucket Cloud Storage durante l'addestramento MaxText:

{
 "tpu": {
   "node_spec": [
     {
      ...
         "metadata": {
               "startup-script": "#! /bin/bash \n pwd \n runuser -l user1 -c 'cd /home/user1/MaxText && python3 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 multislice. Per ulteriori informazioni, consulta la sezione Profilazione dei programmi JAX.

Ottimizzazione della formazione

Sharding con multisezione per ottenere le massime prestazioni

Per ottenere le massime prestazioni negli ambienti multislice, è necessario valutare 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 di attivare lo sharding tra 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 sharding all'interno di una sezione che ha funzionato in passato.

Ti consigliamo di iniziare con il parallelismo dei dati puro. L'utilizzo del parallelismo dei dati completamente segmentato è utile per liberare la memoria utilizzata. Lo svantaggio è che la comunicazione tra le sezioni utilizza la rete DCN e rallenta 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 vuoi migliorarne le prestazioni scalando su più sezioni.

Per ottenere una scalabilità elevata su più sezioni, il tempo necessario per eseguire la riduzione totale su 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 nella velocità effettiva dei carichi di lavoro.

Ogni chip TPU v4 si comporta con 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 ciascun 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-riduce completa e utilizzerai fino a 64 bit di larghezza di banda DCN per ciascun parametro per passaggio.

La quantità di FLOPS necessaria per ciascun parametro varia in base al modello.

In particolare, per i modelli linguistici basati su Transformer, il numero di FLOPS richiesti per un passaggio in avanti e un passaggio indietro è di 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, devi assicurarti 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 all'indietro per la larghezza di banda della rete (in bit) per parametro per passaggio:

Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

Pertanto, per un modello linguistico basato su Transformer, se utilizzi due sezioni:

Operational intensity = 4 * B  / 32

Oppure, se utilizzi più di due sezioni:

Operational intensity = 4 * B  / 64

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

Per altre architetture di modelli, devi stimare il runtime del pass backwards per sezione (calcolando il codice utilizzando un profiler o conteggiando i FLOPS). Puoi quindi confrontare il valore con il tempo di esecuzione previsto per ridurre del tutto sulla rete DCN e ottenere una buona stima del fatto che il parallelismo dei dati sia adatto alla tua attività.

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

Il parallelismo dei dati completamente con sharding (FSDP) combina il parallelismo dei dati (suding dei dati tra i nodi) con lo sharding dei pesi tra i nodi. Per ogni operazione nei passaggi in avanti e all'indietro, i pesi vengono raccolti in modo che ogni sezione abbia i pesi necessari. Invece di sincronizzare i gradienti utilizzando All-Ridurre, i gradienti vengono ridotti dispersi man mano che vengono prodotti. In questo modo, ogni sezione riceve solo i gradienti per i pesi a cui è responsabile.

Analogamente al parallelismo dei dati, FSDP richiederà la scalabilità lineare delle dimensioni globali del batch in base al numero di sezioni. FSDP ridurrà la pressione della memoria all'aumento del numero di sezioni. Questo perché il numero di ponderazioni e lo stato di ottimizzazione per sezione diminuisce, ma lo fa al prezzo dell'aumento del traffico di rete e della maggiore possibilità di blocco a causa di un collettivo ritardato.

In pratica, il valore FSDP tra le sezioni è l'ideale se stai aumentando il batch per sezione, archiviando più attivazioni per ridurre al minimo la rimaterializzazione durante il passaggio all'indietro o aumentando il numero di parametri nella rete neurale.

Le operazioni di raccolta e riduzione complessiva in FSDP funzionano in modo simile a quelle in DP, quindi puoi determinare se il tuo carico di lavoro FSDP è limitato dalle prestazioni della rete 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 del batch globale superiore alla dimensione massima del batch preferita. Il parallelismo delle pipeline consente alle sezioni che comprendono una pipeline di "condividere" un batch. Tuttavia, il parallelismo delle pipeline ha due svantaggi significativi:

1. Corrisponde alla "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 e l'intensità aritmetica e, in ultima analisi, modella l'utilizzo di FLOP.

Il parallelismo delle pipeline deve essere utilizzato solo se le altre strategie richiedono una dimensione batch globale troppo grande. Prima di provare il parallelismo delle pipeline, vale la pena sperimentare per verificare empiricamente se la convergenza per campione rallenta alla dimensione del batch necessaria per ottenere un FSDP ad alte prestazioni. FSDP tende a raggiungere un utilizzo più elevato del modello FLOP, ma se la convergenza per campione rallenta con l'aumento delle dimensioni del batch, il parallelismo della 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 le dimensioni del batch per pipeline fino a quando la latenza DCN non sarà più un fattore in termini di velocità effettiva. In concreto, se hai N sezioni, considera le pipeline di profondità 2 e N/2 repliche del parallelismo dei dati, poi le pipeline di profondità 4 e N/4 delle repliche dei dati di parallelismo e così via, fino a quando il batch per pipeline diventa abbastanza grande da far nascondere i collettivi DCN dietro l'aritmetica nel passaggio indietro. Ciò ridurrà al minimo il rallentamento introdotto dal parallelismo della pipeline, consentendoti al contempo di scalare oltre il limite globale delle dimensioni del batch.

Best practice per il multisezione

Caricamento dei dati

Durante l'addestramento carichiamo ripetutamente batch da un set di dati da inserire nel modello. È importante disporre di un caricatore di dati asincrono efficiente che esegue lo sharding del batch tra gli host per evitare di perdere il lavoro delle TPU. Il caricatore dati attuale in MaxText ha a ciascun carico dell'host un sottoinsieme uguale degli esempi. Questa soluzione è adeguata per il testo, ma richiede un reshard all'interno del modello. Inoltre, MaxText non offre ancora uno snapshot deterministico, che consentirebbe all'iteratore di dati di caricare gli stessi dati prima e dopo il prerilascio.

Verifica

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

Ottimizzazione della formazione

Sharding con multisezione per ottenere le massime prestazioni

Per ottenere le massime prestazioni negli ambienti multislice, è necessario valutare 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 di attivare lo sharding tra 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 sharding all'interno di una sezione che ha funzionato in passato.

Ti consigliamo di iniziare con il parallelismo dei dati puro. L'utilizzo del parallelismo dei dati completamente segmentato è utile per liberare la memoria utilizzata. Lo svantaggio è che la comunicazione tra le sezioni utilizza la rete DCN e rallenta 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 vuoi migliorarne le prestazioni scalando su più sezioni.

Per ottenere una scalabilità elevata su più sezioni, il tempo necessario per eseguire la riduzione totale su 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 nella velocità effettiva dei carichi di lavoro.

Ogni chip TPU v4 si comporta con 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 ciascun 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-riduce completa e utilizzerai fino a 64 bit di larghezza di banda DCN per ciascun parametro per passaggio.

La quantità di FLOPS necessaria per ciascun parametro varia in base al modello.

In particolare, per i modelli linguistici basati su Transformer, il numero di FLOPS richiesti per un passaggio in avanti e un passaggio indietro è di 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, devi assicurarti che l'intensità operativa superi quella aritmetica dell'hardware TPU. Abbiamo calcolato l'intensità aritmetica all'inizio di questa sezione. Per calcolare l'intensità operativa, dividiamo 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 utilizziamo due sezioni:

Operational intensity = 4 * B  / 32

Oppure, se ne utilizziamo più di due sezioni:

Operational intensity = 4 * B  / 64

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

Per altre architetture di modelli, devi stimare il runtime del pass backwards per sezione (calcolando il codice utilizzando un profiler o conteggiando i FLOPS). Puoi quindi confrontare il valore con il tempo di esecuzione previsto per ridurre del tutto sulla rete DCN e ottenere una buona stima del fatto che il parallelismo dei dati sia adatto alla tua attività.

Configurazioni supportate

Forme

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

Generazioni di TPU

La funzionalità multislice è attualmente supportata solo sui pod v4.

Orchestrazione

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

Framework

L'anteprima Multislice v4 supporta solo i carichi di lavoro JAX.

Parallelismo

Consigliamo agli utenti di testare la funzionalità Multisezione con parallelismo dei dati in questa anteprima. Contattaci se vuoi saperne di più sull'implementazione del parallelismo delle pipeline con Multislice.

Assistenza e feedback

Tutti i feedback sono ben accetti. Per condividere feedback o richiedere assistenza, contattaci qui o invia un'email all'indirizzo [e-mail]

Termini

Tutte le informazioni che Google ti ha fornito in merito a questa anteprima privata sono informazioni riservate di Google e sono soggette alle disposizioni in materia di riservatezza indicate nei Termini di servizio della piattaforma Google Cloud (o altri accordi che regolano l'utilizzo di Google Cloud da parte dell'utente).