Addestramento Cloud TPU v5e [anteprima pubblica]

Cloud TPU v5e è l'acceleratore AI di ultima generazione di Google Cloud. Con un ingombro inferiore di 256 chip per pod, una v5e è ottimizzata per essere il prodotto di maggior valore per l'addestramento, il perfezionamento e la gestione della rete neurale convoluzionale (CNN), da testo a immagine e convoluzionale.

Concetti di Cloud TPU v5e

Se non hai mai utilizzato le Cloud TPU, consulta la home page della documentazione TPU.

Questa sezione definisce i concetti specificamente correlati alla versione 5e. I componenti utilizzati in altre configurazioni di TPU (ad esempio sezioni, host, chip e TensorCore) sono descritti nella pagina Architettura del sistema TPU.

Inferenza

L'inferenza è il processo di utilizzo di un modello addestrato per fare previsioni su nuovi dati. Viene utilizzato dal processo di pubblicazione.

attiva

La pubblicazione è il processo di deployment di un modello di machine learning addestrato in un ambiente di produzione in cui può essere utilizzato per fare previsioni o prendere decisioni. La latenza e la disponibilità a livello di servizio sono importanti per la gestione.

Sezioni

Una sezione rappresenta una raccolta di chip tutti situati all'interno dello stesso pod, collegati tramite interconnessioni di chip ad alta velocità (ICI). La versione v5e ha forme delle sezioni 2D. Consulta la tabella nella sezione Tipi di acceleratore per le forme di sezione supportate.

Anche forma chip e topologia chip si riferiscono alle forme delle sezioni.

Host singolo e multi host

Un host è un computer fisico (CPU) che esegue le VM. Un host può eseguire più VM contemporaneamente.

Le sezioni con meno di 8 chip utilizzano al massimo un host. Le sezioni con più di 8 chip hanno accesso a più di un singolo host e possono eseguire un addestramento distribuito utilizzando più host. Consulta la pagina relativa all'architettura del sistema TPU per i dettagli su sezioni e chip.

La versione 5e supporta l'addestramento multi-host e l'inferenza su un solo host. Questa anteprima non supporta l'inferenza o la pubblicazione multi-host. Il numero massimo di chip consentiti in un singolo job di pubblicazione è 8 chip.

VM TPU

Una macchina virtuale che esegue Linux e che ha accesso alle TPU sottostanti. Per le TPU v5e, ogni VM TPU ha accesso diretto a 1, 4 o 8 chip a seconda del tipo di acceleratore specificato dall'utente. Una VM TPU è nota anche come worker.

Worker

Vedi VM TPU.

architettura del sistema v5e

Ogni chip v5e contiene un TensorCore. Ogni TensorCore ha 4 Matrix Multiply Unit (MXU), un'unità vettoriale e un'unità scalare. La tabella riportata di seguito mostra le specifiche del chip di chiave e i relativi valori per una versione v5e. Le specifiche dei pod sono incluse nella tabella che segue le specifiche del chip.

Specifiche del chip di chiave Valori v5e
Picco di computing per chip (bf16) 197 TFLOP
Picco di computing per chip (Int8) 393 TFLOP
Capacità e larghezza di banda HBM2 16 GB, 819 GB/s
BW Interchip Interconnect 1600 Gbps
Specifiche chiave dei pod Valori v5e
Dimensioni pod di TPU 256 chip
Topologia di interconnessione Torus 2D
Picco di computing per pod 100 PetaOps(Int8)
Riduce completamente la larghezza di banda per pod 51,2 TB/s
Larghezza di banda di bisezione per pod 1,6 TB/s
Larghezza di banda della rete del data center per pod 6,4 Tbps

Chip TPU v5e

Il seguente diagramma illustra un chip TPU v5e.

Chip pod v5e

Tipi di acceleratori

Cloud TPU v5e è un prodotto combinato di addestramento e inferenza (pubblicazione). Il flag AcceleratorType viene utilizzato per distinguere l'ambiente TPU di cui è stato eseguito il provisioning per l'addestramento da quello di cui è stato eseguito il provisioning per la pubblicazione.

I job di addestramento sono ottimizzati in termini di velocità effettiva e disponibilità, mentre i job di servizio sono ottimizzati per la latenza. Pertanto, un job di addestramento sulle TPU di cui è stato eseguito il provisioning potrebbe avere un'esperienza scadente (disponibilità inferiore) e, allo stesso modo, un job di addestramento eseguito sulle TPU di cui è stato eseguito il provisioning per l'addestramento potrebbe avere un'esperienza scadente (latenza maggiore).

AcceleratorType utilizza il numero di TensorCore nel nodo per descrivere le dimensioni della sezione. AcceleratorType è una stringa con il formato v$VERSION_NUMBER-$CORES_COUNT.

Per ulteriori informazioni sui tipi di acceleratori, consulta Configurazioni TPU supportate.

Una v5e ha un core per chip. Le seguenti forme di sezione 2D sono supportate per v5es:

Topologia Numero di chip TPU Numero di host
1x1 1 8/1
2x2 4 1/2
2x4 8 1
4x4 16 2
4x8 32 4
8x8 64 8
8x16 128 16
16x16 256 32

Per i conteggi delle VM per topologia, consulta Tipi di VM.

Tipi di Cloud TPU v5e

L'addestramento è supportato per un massimo di 256 chip.

Per eseguire il provisioning delle TPU per un job di addestramento v5e, utilizza uno dei seguenti tipi di acceleratori nella richiesta di creazione dell'interfaccia a riga di comando o di TPU API:

  • V5E-1
  • V5E-4
  • V5E-8
  • V5E-16
  • V5E-32
  • V5E-64
  • V5E-128
  • V5E-256

Il provisioning per l'inferenza è descritto nell'introduzione all'inferenza v5e.

Tipi di VM TPU

Ogni VM TPU in una sezione di TPU contiene 1, 4 o 8 chip. Le sezioni più piccole a 4 chip avranno lo stesso nodo NUMA (Non Uniform Memory Access). Per ulteriori informazioni sui nodi NUMA, consulta il documento sull'architettura di sistema. Per le VM TPU a 8 chip, la comunicazione CPU-TPU sarà più efficiente all'interno delle partizioni NUMA. Ad esempio, nella figura seguente, la comunicazione CPU0-Chip0 sarà più veloce della comunicazione CPU0-Chip4.

Comunicazione dei nodi NUMA

Inizia

Capacità di protezione

Cloud TPU v5e è ora in Anteprima pubblica. Contatta il team di vendita Cloud per iniziare a utilizzare Cloud TPU v5e per i tuoi carichi di lavoro di AI.

Prepara un progetto Google Cloud

  1. Accedi al tuo Account Google. Se non l'hai ancora fatto, crea un nuovo account.
  2. Nella console Google Cloud, seleziona o crea un progetto Cloud dalla pagina del selettore dei progetti.
  3. La configurazione della fatturazione è obbligatoria per tutto l'utilizzo di Google Cloud, quindi assicurati che la fatturazione sia abilitata per il tuo progetto.

    La fatturazione per l'utilizzo dell'Anteprima pubblica segue i prezzi regionali standard mostrati nella pagina dei prezzi di Cloud TPU.

  4. Installa i componenti alpha di gcloud.

  5. Se sei un utente TPU che riutilizza i componenti gcloud alpha esistenti, aggiornali per assicurarti che i comandi e i flag pertinenti siano supportati:

    gcloud components update
    
  6. Abilita l'API TPU utilizzando il seguente comando gcloud in Cloud Shell. Puoi abilitarlo anche dalla console Google Cloud.

    gcloud services enable tpu.googleapis.com
    
  7. Abilita l'account di servizio TPU.

    Gli account di servizio consentono al servizio Cloud TPU di accedere ad altri servizi Google Cloud. Un account di servizio gestito dall'utente è una pratica consigliata in Google Cloud. Segui queste guide per creare e concedere i ruoli. Sono necessari i seguenti ruoli:

    • Amministratore TPU
    • Amministratore Storage: necessaria per accedere a Cloud Storage
    • Writer log: necessario per scrivere i log con l'API Logging
    • Monitoring Metric Writer: necessario per scrivere le metriche in Cloud Monitoring
  8. Configurare il progetto e la zona.

    L'ID progetto è il nome del progetto visualizzato nella console Cloud.

    export PROJECT_ID=your_project_ID
    export ZONE=us-west4-a
    
    gcloud alpha compute tpus tpu-vm service-identity create --zone=${ZONE}
    
    gcloud auth login
    gcloud config set project ${PROJECT_ID}
    gcloud config set compute/zone ${ZONE}
    

Esegui il provisioning dell'ambiente Cloud TPU

Le Cloud TPU v5 sono gestite come risorse in coda. Puoi eseguire il provisioning della capacità utilizzando il comando queued-resource create.

Creare variabili di ambiente

Imposta le variabili di ambiente necessarie per la creazione della TPU:

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5e-16
   export ZONE=us-west4-a
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   

Variabile Description
ID PROGETTO Nome progetto Google Cloud
TIPO_ACCELERATOR Consulta la sezione Tipi di acceleratore per le forme delle sezioni supportate.
ZONA Tutte le capacità in Anteprima pubblica sono disponibili in us-west4-a.
RUNTIME_VERSION Utilizza v2-alpha-tpuv5-lite per l'Anteprima pubblica.
ACCOUNT_SERVIZIO Questo è l'indirizzo dell'account di servizio che puoi trovare nella console Google Cloud -> IAM -> Account di servizio. Ad esempio: tpu-service-account@myprojectID.iam.gserviceaccount.com
TPU_NAME L'ID testo assegnato dall'utente della TPU creato al momento dell'allocazione della richiesta di risorsa in coda.
ID_RISORSA_CODA L'ID testo assegnato dall'utente della richiesta di risorsa in coda. Consulta il documento Risorse in coda per informazioni sulle risorse in coda.

Crea una risorsa TPU

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID}
   --node-id ${TPU_NAME} \
   --project ${PROJECT_ID} \
   --zone ${ZONE} \
   --accelerator-type ${ACCELERATOR_TYPE} \
   --runtime-version ${RUNTIME_VERSION} \
   --service-account ${SERVICE_ACCOUNT} \
   --reserved

Se la risorsa in coda viene creata correttamente, lo stato all'interno del campo response sarà ACCEPTED o FAILED. Se la risorsa in coda è nello stato ACCEPTED, significa che la risorsa in coda ha superato la convalida preliminare ed è in attesa di capacità. Quando sarà disponibile la capacità, la richiesta passerà a PROVISIONING. Lo stato ACCEPTED non significa che otterrai la quota allocata e potrebbe essere necessario un po' di tempo per cambiare lo stato da ACCEPTED a ACTIVE. Se la risorsa in coda è nello stato FAILED, il motivo dell'errore sarà riportato nell'output.

Potrai connetterti tramite SSH alla VM TPU una volta che QueuedResource è nello stato ACTIVE.

Utilizza i comandi list o describe per eseguire query sullo stato della risorsa in coda.

    gcloud alpha compute tpus queued-resources describe {QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE}
   
Lo stato rappresenta lo stato di una risorsa in coda. Gli stati sono definiti come:

Stato Description
ACCETTATI Il comando create della risorsa in coda è stato ricevuto e inizierà il provisioning, non appena sarà disponibile la capacità.
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, l'esecuzione su tutte le TPU inizierà a essere eseguita su tutte le TPU, quando lo stato della risorsa in coda passa a ACTIVE.
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 la risorsa in coda rimane intatta finché non viene eliminata esplicitamente. Al momento, una risorsa sospesa in coda non può essere ripresa e deve essere eliminata.
DELETING La risorsa in coda è in fase di eliminazione.

SSH nelle VM TPU

La sezione seguente descrive come installare programmi binari su ogni VM TPU nella sezione TPU ed eseguire il codice. In questo contesto, una VM TPU è nota anche come worker.

Consulta la sezione Tipi di VM per calcolare quante VM avrà la sezione.

Per installare i programmi binari o eseguire il codice, devi accedere alla VM TPU tramite SSH utilizzando tpu-vm ssh command.

$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}

Per connetterti tramite SSH a una VM TPU specifica, utilizza il flag --worker che segue un indice basato su 0:

$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --worker=1

Se le forme delle sezioni superano gli 8 chip, avrai più VM in una sezione. In questo caso, utilizza il flag --worker=all per eseguire l'installazione su tutte le VM TPU in un solo passaggio senza dover accedere tramite SSH a ognuna separatamente. Ad esempio:

$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} /
--zone  ${ZONE} \
--worker=all \
--command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Gestisci

Tutti i comandi che puoi utilizzare per gestire le VM TPU sono descritti in Gestione delle TPU.

Configurazione del framework

Questa sezione descrive il processo di configurazione generale per l'addestramento di modelli personalizzati utilizzando JAX o PyTorch con TPU v5e. Il supporto di TensorFlow verrà abilitato a breve.

Per istruzioni sulla configurazione dell'inferenza, consulta l'introduzione all'inferenza v5e.

Configurazione per JAX

Se le forme delle sezioni superano gli 8 chip, avrai più VM in una sezione. In questo caso, devi utilizzare il flag --worker=all per eseguire l'installazione su tutte le VM TPU in un unico passaggio senza utilizzare SSH per accedere a ciascuna VM separatamente:

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} \
--zone  ${ZONE} \
--worker=all \
--command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'

Puoi eseguire il seguente comando per controllare il numero di dispositivi (gli output mostrati qui sono stati prodotti con una sezione v5litepod-16). Questo codice verifica che tutto sia installato correttamente controllando che JAX rilevi i TensorCore Cloud TPU e possa eseguire operazioni di base:

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

L'output sarà simile al seguente:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4

jax.device_count() mostra il numero totale di chip nella sezione specificata. jax.local_device_count() indica il numero di chip accessibili da una singola VM in questa sezione.

# Check the number of chips in the given slice by summing the count of chips
# from all VMs through the
# jax.local_device_count() API call.
gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='python3 -c "import jax; xs=jax.numpy.ones(jax.local_device_count()); print(jax.pmap(lambda x: jax.lax.psum(x, \"i\"), axis_name=\"i\")(xs))"'

L'output sarà simile al seguente:

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]
[16. 16. 16. 16.]

Prova i tutorial su JAX in questo documento per iniziare l'addestramento v5e utilizzando JAX.

Configurazione per PyTorch

Tieni presente che v5e supporta solo il runtime PJRT e PyTorch 2.1 e versioni successive utilizzeranno PJRT come runtime predefinito per tutte le versioni TPU.

Questa sezione descrive come iniziare a utilizzare PJRT su v5e con PyTorch/XLA con i comandi per tutti i worker.

Installa le dipendenze

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--worker=all \
--command='
   sudo apt-get update -y
   sudo apt-get install libomp5 -y
   pip3 install mkl mkl-include
   pip3 install tf-nightly tb-nightly tbp-nightly
   pip3 install numpy
   sudo apt-get install numactl libopenblas-dev -y
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
   gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
   pip3 install torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl
   pip3 install torch_xla[tpuvm]
'

Se non hai installato le ruote per torch/torch_xla/torchvision e visualizzi un errore come pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222, esegui il downgrade della tua versione con questo comando:

pip3 install setuptools==62.1.0

Esegui un semplice script con PJRT:

Per i modelli con allocazioni frequenti e considerevoli, l'utilizzo di tcmalloc può migliorare notevolmente i tempi di addestramento rispetto all'implementazione malloc predefinita, pertanto il valore malloc predefinito utilizzato su una VM TPU è tcmalloc. Tuttavia, a seconda del carico di lavoro (ad esempio, con DLRM che ha allocazioni molto grandi per le tabelle di incorporamento) tcmalloc potrebbe causare un rallentamento. In questo caso, potresti provare a annullare l'impostazione della seguente variabile utilizzando invece il valore predefinito malloc:

unset LD_PRELOAD

Di seguito è riportato un esempio in cui viene utilizzato uno script Python per eseguire un semplice calcolo su una VM v5e:

gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker all --command='
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/$USER/.local/lib/
export PJRT_DEVICE=TPU_C_API
export PT_XLA_DEBUG=0
export USE_TORCH=ON
unset LD_PRELOAD
export TPU_LIBRARY_PATH=/home/$USER/.local/lib/python3.10/site-packages/libtpu/libtpu.so

python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"
'
Questo genera un output simile al seguente:
SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')

Prova i tutorial su PyTorch in questo documento per iniziare l'addestramento v5e utilizzando PyTorch.

Elimina la TPU e la risorsa in coda alla fine della sessione. Per eliminare una risorsa in coda, elimina la sezione e poi la risorsa in coda in due passaggi:

gcloud alpha compute tpus tpu-vm delete $TPU_NAME \
--project=$PROJECT_ID \
--zone=$ZONE \
--quiet

gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} \
--zone ${ZONE} \
--quiet

Questi due passaggi possono essere utilizzati anche per rimuovere le richieste di risorse in coda che sono in stato FAILED.

Monitoraggio e profilo

Cloud TPU v5e supporta il monitoraggio e la profilazione utilizzando gli stessi metodi delle generazioni precedenti di Cloud TPU. Per saperne di più sul monitoraggio, consulta Profilo del modello con gli strumenti Cloud TPU per saperne di più sulla profilazione e sul monitoraggio delle VM Cloud TPU.

Esempi JAX/FLAX

Addestra ImageNet su v5e

Questo tutorial descrive come addestrare ImageNet su v5e utilizzando dati di input falsi. Se vuoi utilizzare dati reali, consulta il file README su GitHub.

Configura

  1. Crea le variabili di ambiente:

    export PROJECT_ID=your_project_ID
    export ACCELERATOR_TYPE=v5e-16
    export ZONE=us-west4-a
    export RUNTIME_VERSION=v2-test-tpuv5-lite
    export SERVICE_ACCOUNT=your_service_account
    export TPU_NAME=your_tpu_name
    export QUEUED_RESOURCE_ID=your_queued_resource_id
    
  2. Crea una risorsa TPU:

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
    --node-id ${TPU_NAME} \
    --project ${PROJECT_ID} \
    --zone ${ZONE} \
    --accelerator-type ${ACCELERATOR_TYPE} \
    --runtime-version ${RUNTIME_VERSION} \
    --service-account ${SERVICE_ACCOUNT} \
    --reserved
    
    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  --project ${PROJECT_ID} --zone ${ZONE}
    
    When the QueuedResource is in the ACTIVE state, the output will be similar to the following:
    …
    state: ACTIVE
    …
    
  3. Installa l'ultima versione di JAX e jaxlib:

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
    
  4. Clona il modello ImageNet e installa i requisiti corrispondenti:

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='git clone -b throughput_testing https://github.com/gkroiz/flax.git && cd flax && pip install -e . && cd examples/imagenet && pip install -r requirements.txt && pip install --upgrade git+https://github.com/google/flax.git'
    
  5. Per generare dati falsi, il modello ha bisogno di informazioni sulle dimensioni del set di dati. Questi dati possono essere raccolti dai metadati del set di dati ImageNet:

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='mkdir -p $HOME/flax/.tfds/metadata/imagenet2012/5.1.0 && curl https://raw.githubusercontent.com/tensorflow/datasets/v4.4.0/tensorflow_datasets/testing/metadata/imagenet2012/5.1.0/dataset_info.json --output $HOME/flax/.tfds/metadata/imagenet2012/5.1.0/dataset_info.json'
    

Addestra il modello

Una volta completati tutti i passaggi precedenti, puoi addestrare il modello. Nel repository GitHub sono disponibili diversi file di configurazione. In alternativa, puoi personalizzare le configurazioni eseguendo l'override nella riga di comando.

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='cd flax/examples/imagenet && JAX_USE_PJRT_C_API_ON_TPU=1 JAX_PLATFORMS=tpu TF_CPP_MIN_LOG_LEVEL=0 TPU_STDERR_LOG_LEVEL=0 TPU_MIN_LOG_LEVEL=0 python3 imagenet_fake_data_benchmark.py --config=configs/fake_data_16_chips.py'
    

Elimina la TPU e la risorsa in coda

Elimina la TPU e la risorsa in coda alla fine della sessione.

   gcloud alpha compute tpus tpu-vm delete $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --quiet
   
   gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE} --quiet
   

Modelli in FLAX con faccina che abbraccia

I modelli Hugging Face implementati in FLAX funzionano immediatamente su Cloud TPU v5e. Questa sezione fornisce istruzioni per eseguire i modelli più utilizzati.

Addestra ViT su Imagenette

Questo tutorial mostra come addestrare il modello Vision Transformer (ViT) da HuggingFace utilizzando il set di dati imagenette di Fast AI su Cloud TPU v5e.

Il modello ViT è stato il primo ad aver addestrato con successo un codificatore Transformer su ImageNet con risultati eccellenti rispetto alle reti convoluzionali. Per maggiori informazioni, consulta le seguenti risorse:

Configura

  1. Crea le variabili di ambiente:

    export PROJECT_ID=your_project_ID
    export ACCELERATOR_TYPE=v5e-16
    export ZONE=us-west4-a
    export RUNTIME_VERSION=v2-test-tpuv5-lite
    export SERVICE_ACCOUNT=your_service_account
    export TPU_NAME=your_tpu_name
    export QUEUED_RESOURCE_ID=your_queued_resource_id
    
  2. Creare una risorsa TPU.

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
    --node-id ${TPU_NAME} \
    --project ${PROJECT_ID} \
    --zone ${ZONE} \
    --accelerator-type ${ACCELERATOR_TYPE} \
    --runtime-version ${RUNTIME_VERSION} \
    --service-account ${SERVICE_ACCOUNT}
    

    Potrai accedere tramite SSH alla VM TPU una volta che la risorsa in coda si trova nello stato ACTIVE:

    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
    --project ${PROJECT_ID} --zone ${ZONE}
    
    When the queued resource is in the ACTIVE state, the output will be similar to the following:
    …
    state: ACTIVE
    …
    
  3. Installa JAX e la relativa libreria:

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
    
  4. Scarica il repository Hugging Face e installa i requisiti:

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow && pip install -r examples/flax/vision/requirements.txt'
    
  5. Scarica il set di dati Imagenette:

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=${ZONE} --worker=all --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'
    

Addestra il modello

Addestra il modello con un buffer premappato a 4 GB.

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='cd transformers && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'
  

Elimina la TPU e la risorsa in coda

Elimina la TPU e la risorsa in coda alla fine della sessione.

   gcloud alpha compute tpus tpu-vm delete $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --quiet
   
   gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE} --quiet
   

Risultati del benchmarking ViT

Lo script di addestramento è stato eseguito su v5e-4, v5e-16 e v5e-64. La tabella seguente mostra le velocità effettiva con diversi tipi di acceleratore.

Tipo di acceleratore V5E-4 V5E-16 V5E-64
Epoca 3 3 3
Dimensione batch globale 32 128 512
Velocità effettiva (esempi/sec) 263,40 429,34 470,71

Allena Diffusione su Pokémon

Questo tutorial mostra come addestrare il modello di diffusione stabile da HuggingFace utilizzando il set di dati Pokémon su Cloud TPU v5e.

Il modello Diffusione stabile è un modello da testo a immagine latente che genera immagini fotorealistiche da qualsiasi input di testo. Per saperne di più, consulta le seguenti risorse:

Configura

  1. Crea le variabili di ambiente:

    export PROJECT_ID=your_project_ID
    export ACCELERATOR_TYPE=v5e-16
    export ZONE=us-west4-a
    export RUNTIME_VERSION=v2-test-tpuv5-lite
    export SERVICE_ACCOUNT=your_service_account
    export TPU_NAME=your_tpu_name
    export QUEUED_RESOURCE_ID=your_queued_resource_id
    

  2. Crea una risorsa TPU:

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
    --node-id ${TPU_NAME} \
    --project ${PROJECT_ID} \
    --zone ${ZONE} \
    --accelerator-type ${ACCELERATOR_TYPE} \
    --runtime-version ${RUNTIME_VERSION} \
    --service-account ${SERVICE_ACCOUNT} \
    --reserved
    
    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  --project ${PROJECT_ID} --zone ${ZONE}
    
    When the queued resource is in the ACTIVE state, the output will be similar to the following:
    …
    state: ACTIVE
    …
    
  3. Installa JAX e la relativa libreria.

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
    
  4. Scarica il repo di HuggingFace e installa i requisiti.

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='git clone https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install tensorflow clu && pip install -U -r examples/text_to_image/requirements_flax.txt'
    

Addestra il modello

Addestra il modello con un buffer premappato a 4 GB.

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --=$PROJECT=${PROJECT_ID} --zone=$ZONE --worker=all --command='cd diffusers/examples/text_to_image && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --fake_data=1 --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model'
   

Elimina la TPU e la risorsa in coda

Elimina la TPU e la risorsa in coda alla fine della sessione.

   gcloud alpha compute tpus tpu-vm delete $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --quiet
   
   gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE} --quiet
   

Benchmarking dei risultati per la diffusione

Lo script di addestramento è stato eseguito su v5e-4, v5e-16 e v5e-64. La tabella seguente mostra le velocità effettiva.

Tipo di acceleratore V5E-4 V5E-16 V5E-64
Passaggio di addestramento 1500 1500 1500
Dimensione batch globale 32 64 128
Velocità effettiva (esempi/sec) 36,53 43,71 49,36

Addestrare GPT2 sul set di dati OSCAR

Questo tutorial mostra come addestrare il modello GPT2 da HuggingFace utilizzando il set di dati OSCAR su Cloud TPU v5e.

GPT2 è un modello Transformer preaddestrato su testi non elaborati senza etichettatura umana. È stato addestrato per prevedere la parola successiva nelle frasi. Per maggiori informazioni, consulta le seguenti risorse:

Configura

  1. Crea le variabili di ambiente:

    export PROJECT_ID=your_project_ID
    export ACCELERATOR_TYPE=v5e-16
    export ZONE=us-west4-a
    export RUNTIME_VERSION=v2-test-tpuv5-lite
    export SERVICE_ACCOUNT=your_service_account
    export TPU_NAME=your_tpu_name
    export QUEUED_RESOURCE_ID=your_queued_resource_id
    
  2. Creare una risorsa TPU.

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
    --node-id ${TPU_NAME} \
    --project ${PROJECT_ID} \
    --zone ${ZONE} \
    --accelerator-type ${ACCELERATOR_TYPE} \
    --runtime-version ${RUNTIME_VERSION} \
    --service-account ${SERVICE_ACCOUNT} \
    --reserved
    
    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  --project ${PROJECT_ID} --zone ${ZONE}
    
    When the queued resource is in the ACTIVE state, the output will be similar to the following:
    …
    state: ACTIVE
    …
    
  3. Installa JAX e la relativa libreria.

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='pip install "jax[tpu]==0.4.6" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
    
  4. Scarica il repo di HuggingFace e i requisiti di installazione.

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow && pip install -r examples/flax/language-modeling/requirements.txt'
    
  5. Scarica le configurazioni per addestrare il modello.

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='cd transformers/examples/flax/language-modeling && gsutil cp -r gs://cloud-tpu-v2-images-artifacts/config/v5litepod-preview/jax/gpt .'
    

Addestra il modello

Addestra il modello con un buffer premappato a 4 GB.

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --worker=all --command='cd transformers/examples/flax/language-modeling && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 run_clm_flax.py --output_dir=./gpt --model_type=gpt2 --config_name=./gpt --tokenizer_name=./gpt --dataset_name=oscar --dataset_config_name=unshuffled_deduplicated_no --do_train --do_eval --block_size=512 --per_device_train_batch_size=4 --per_device_eval_batch_size=4 --learning_rate=5e-3 --warmup_steps=1000 --adam_beta1=0.9 --adam_beta2=0.98 --weight_decay=0.01 --overwrite_output_dir --num_train_epochs=3 --logging_steps=500 --eval_steps=2500'
   

Elimina la TPU e la risorsa in coda

Elimina la TPU e la risorsa in coda alla fine della sessione.

   gcloud alpha compute tpus tpu-vm delete $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --quiet
   
   gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE} --quiet
   

Benchmarking dei risultati per GPT2

Lo script di addestramento è stato eseguito su v5e-4, v5e-16 e v5e-64. La tabella seguente mostra le velocità effettiva.

V5E-4 V5E-16 V5E-64
Epoca 3 3 3
Dimensione batch globale 64 64 64
Velocità effettiva (esempi/sec) 74,60 72,97 72,62

PyTorch/XLA

Addestra ResNet utilizzando il runtime PJRT

PyTorch/XLA sta eseguendo la migrazione da XRT a PjRt da PyTorch 2.0 e versioni successive. Di seguito sono riportate le istruzioni aggiornate per configurare la versione v5e per i carichi di lavoro di addestramento PyTorch/XLA.

Configurazione
  1. Crea le variabili di ambiente:

    export PROJECT_ID=your_project_ID
    export ACCELERATOR_TYPE=v5e-16
    export ZONE=us-west4-a
    export RUNTIME_VERSION=v2-alpha-tpuv5-lite
    export SERVICE_ACCOUNT=your_service_account
    export TPU_NAME=tpu-name
    export QUEUED_RESOURCE_ID=your_queued_resource_id
    
  2. Crea una risorsa TPU:

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
    --node-id ${TPU_NAME} \
    --project ${PROJECT_ID} \
    --zone ${ZONE} \
    --accelerator-type ${ACCELERATOR_TYPE} \
    --runtime-version ${RUNTIME_VERSION} \
    --service-account ${SERVICE_ACCOUNT} \
    --reserved
    
    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  --project ${PROJECT_ID} --zone ${ZONE}
    
    When the queued resource is in the ACTIVE state, the output will be similar to the following:
    …
    state: ACTIVE
    …
    
  3. Clona PyTorch/XLA

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker=all --command='cd /usr/share/
    sudo git clone --recursive https://github.com/pytorch/pytorch
    cd pytorch/
    sudo git clone --recursive https://github.com/pytorch/xla.git
    cd xla/
    sudo git checkout 3b10663
    sudo pip3 install numpy'
    
  4. Installazione di Torch, Torch/xla, Torchvision

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE} --project ${PROJECT_ID} --worker=all --command='
    pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch-2.0%2B20230222-cp38-cp38-linux_x86_64.whl
    pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-2.0%2B20230222-cp38-cp38-linux_x86_64.whl
    pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torchvision-2.0%2B20230222-cp38-cp38-linux_x86_64.whl
    sudo rm -rf /usr/local/lib/python3.8/dist-packages/libtpu*
    sudo pip3 install torch_xla[tpuvm]
    '
    
  5. Installa dipendenze specifiche per Torch/XLA

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} --zone  ${ZONE} --worker all --command='
    sudo apt-get update -y
    sudo apt-get install libomp5 -y
    pip3 install mkl mkl-include
    pip3 install tf-nightly tb-nightly tbp-nightly
    pip3 install numpy
    sudo apt-get install numactl libopenblas-dev -y
    pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
    pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
    gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
    pip3 install torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl
    pip3 install torch_xla[tpuvm]
    '
    
  6. Scarica libtpu

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker=all --command="
    sudo gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/20230216/libtpu.so /home/$USER/.local/lib/python3.8/site-packages/libtpu/libtpu.so
    "
    
  7. Aggiorna tpu.py con la configurazione v5e

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker=all --command="
    sudo gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/20230216/tpu.py /home/$USER/.local/lib/python3.8/site-packages/torch_xla/experimental/tpu.py
    "
    
Addestra il modello ResNet
      gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --project ${PROJECT_ID} --zone ${ZONE} --worker=all --command='
      date
      export PJRT_DEVICE=TPU_C_API
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      export XLA_USE_BF16=1
      export LIBTPU_INIT_ARGS=--xla_jf_auto_cross_replica_sharding
      export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
      export TPU_LIBRARY_PATH=/home/$USER/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      git clone https://github.com/pytorch/xla.git
      cd xla/
      git reset --hard caf5168785c081cd7eb60b49fe4fffeb894c39d9
      python3 test/test_train_mp_imagenet.py --model=resnet50  --fake_data --num_epochs=1 —num_workers=16  --log_steps=300 --batch_size=64 --profile'

Elimina la TPU e la risorsa in coda

Elimina la TPU e la risorsa in coda alla fine della sessione.

   gcloud alpha compute tpus tpu-vm delete $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --quiet
   gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE} --quiet
Risultato benchmark

La tabella seguente mostra le velocità effettiva dei benchmark.

Tipo di acceleratore Velocità effettiva (esempi/secondo)
V5E-4 4240 ex/s
V5E-16 10.810 ex/s
V5E-64 46.154 ex/s

Addestrare GPT2 su v5e

Questo tutorial illustra come eseguire GPT2 su v5e utilizzando il repo di HuggingFace su PyTorch/XLA utilizzando il set di dati wikitext.

Configurazione

  1. Crea le variabili di ambiente:

    export PROJECT_ID=your_project_ID
    export ACCELERATOR_TYPE=v5e-16
    export ZONE=us-west4-a
    export RUNTIME_VERSION=v2-alpha-tpuv5-lite
    export SERVICE_ACCOUNT=your_service_account
    export TPU_NAME=your_tpu_name
    export QUEUED_RESOURCE_ID=your_queued_resource_id
    
  2. Creare una risorsa TPU.

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
    --node-id ${TPU_NAME} \
    --project ${PROJECT_ID} \
    --zone ${ZONE} \
    --accelerator-type ${ACCELERATOR_TYPE} \
    --runtime-version ${RUNTIME_VERSION} \
    --service-account ${SERVICE_ACCOUNT}
    
    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  --project ${PROJECT_ID} --zone ${ZONE}
    
    When the queued resource is in the ACTIVE state, the output will be similar to the following:
    …
    state: ACTIVE
    …
    
  3. Installa le dipendenze torch/xla.

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} --zone  ${ZONE} --worker all --command='
    sudo apt-get update -y
    sudo apt-get install libomp5 -y
    pip3 install mkl mkl-include
    pip3 install tf-nightly tb-nightly tbp-nightly
    pip3 install numpy
    sudo apt-get install numactl libopenblas-dev -y
    pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
    pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly%2B20230712-cp310-cp310-linux_x86_64.whl
    gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
    pip3 install torchvision-0.16.0a0+fc838ad-cp310-cp310-linux_x86_64.whl
    pip3 install torch_xla[tpuvm]
    '
    
  4. Scarica il repo di HuggingFace e i requisiti di installazione.

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command="
    rm -rf transformers/
    git clone https://github.com/huggingface/transformers.git; \
    cd transformers; \
    sudo pip3 install -e .; \
    git checkout -q ebdb185befaa821304d461ed6aa20a17e4dc3aa2; \
    pip3 install datasets; \
    pip3 install evaluate; \
    pip3 install scikit-learn;"
    
  5. Scarica le configurazioni del modello preaddestrato.

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command="
    cd transformers &&
    gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/GPT2/my_config_2.json examples/pytorch/language-modeling/
    gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/GPT2/fsdp_config.json examples/pytorch/language-modeling/
    "
    

Addestra il modello

Addestra il modello 2 miliardi utilizzando una dimensione batch di 16.

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME --project=$PROJECT_ID --zone=$ZONE --worker=all --command='
   export PJRT_DEVICE=TPU_C_API
   export PT_XLA_DEBUG=0
   export USE_TORCH=ON
   export XLA_USE_BF16=1
   export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
   export TPU_LIBRARY_PATH=/home/$USER/.local/lib/python3.10/site-packages/libtpu/libtpu.so
   cd transformers
   numactl --cpunodebind=0 python3 -u examples/pytorch/xla_spawn.py --num_cores 4
    examples/pytorch/language-modeling/run_clm.py --num_train_epochs 3
    --dataset_name wikitext
    --dataset_config_name wikitext-2-raw-v1
    --per_device_train_batch_size 16
    --per_device_eval_batch_size 16
    --do_train
    --do_eval
    --output_dir /tmp/test-clm
    --overwrite_output_dir
    --config_name examples/pytorch/language-modeling/my_config_2.json --cache_dir /tmp
    --tokenizer_name gpt2
    --block_size 1024
    --optim adafactor
    --adafactor true
    --save_strategy no
    --logging_strategy no
    --fsdp "full_shard"
    --fsdp_config examples/pytorch/language-modeling/fsdp_config.json'
   

Elimina la TPU e la risorsa in coda

Elimina la TPU e la risorsa in coda alla fine della sessione.

   gcloud alpha compute tpus tpu-vm delete $TPU_NAME --project=${PROJECT_ID} --zone=$ZONE --quiet
   
   gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} --project ${PROJECT_ID} --zone ${ZONE} --quiet
   

Risultato benchmark

Lo script di addestramento è stato eseguito su v5e-4, v5e-16 e v5e-64. La tabella seguente mostra le velocità effettiva dei benchmark per i diversi tipi di acceleratore.

V5E-4 V5E-16 V5E-64
Epoca 3 3 3
config 600 Mln 2 MLD 16 MLD
Dimensione batch globale 64 128 256
Velocità effettiva (esempi/sec) 66 77 31

Treno ViT su v5e

Questo tutorial illustra come eseguire VIT su v5e utilizzando il repo HuggingFace su PyTorch/XLA nel set di dati cifar10.

Configurazione

  1. Crea le variabili di ambiente:

    export PROJECT_ID=your_project_ID
    export ACCELERATOR_TYPE=v5e-16
    export ZONE=us-west4-a
    export RUNTIME_VERSION=v2-alpha-tpuv5-lite
    export SERVICE_ACCOUNT=your_service_account
    export TPU_NAME=tpu-name
    export QUEUED_RESOURCE_ID=queued-resource-id
    
  2. Creare una risorsa TPU.

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
    --node-id ${TPU_NAME} \
    --project ${PROJECT_ID} \
    --zone ${ZONE} \
    --accelerator-type ${ACCELERATOR_TYPE} \
    --runtime-version ${RUNTIME_VERSION} \
    --service-account ${SERVICE_ACCOUNT} \
    --reserved
    
    gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  --project ${PROJECT_ID} --zone ${ZONE}
    
    When the queued resource is in the ACTIVE state, the output will be similar to the following:
    …
    state: ACTIVE
    …
    
  3. Installa le dipendenze torch/xla

    gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME}  --project ${PROJECT_ID} --zone  ${ZONE} --worker=all --command='
    sudo apt-get update -y
    sudo apt-get install libomp5 -y
    pip3 install mkl mkl-include
    pip3 install tf-nightly tb-nightly tbp-nightly
    pip3 install numpy
    sudo apt-get install numactl libopenblas-dev -y
    pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch-nightly+20230726-cp310-cp310-linux_x86_64.whl
    pip3 install https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-nightly+20230726-cp310-cp310-linux_x86_64.whl
    gsutil cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/pytorch/wheels/torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl .
    pip3 install torchvision-0.16.0a0+c9ac3a5-cp310-cp310-linux_x86_64.whl
    pip3 install torch_xla[tpuvm]
    '
    
  4. Scarica il repo di HuggingFace e i requisiti di installazione.

    gcloud alpha compute tpus tpu-vm ssh $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE \
    --worker=all \
    --command="
      git clone https://github.com/suexu1025/transformers.git vittransformers; \
      cd vittransformers; \
      pip3 install .; \
      pip3 install datasets; \
      wget https://github.com/pytorch/xla/blob/master/scripts/capture_profile.py"
    

Addestra il modello

   gcloud alpha compute tpus tpu-vm ssh $TPU_NAME 
--project=$PROJECT_ID
--zone=$ZONE
--worker=all
--command=" export PJRT_DEVICE=TPU_C_API export PT_XLA_DEBUG=0 export USE_TORCH=ON export TF_CPP_MIN_LOG_LEVEL=0 export XLA_USE_BF16=1 export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH export TPU_LIBRARY_PATH=/home/$USER/.local/lib/python3.10/site-packages/libtpu/libtpu.so cd vittransformers; numactl --cpunodebind=0 python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/image-pretraining/run_mae.py --dataset_name=cifar10 --remove_unused_columns=False --label_names=pixel_values --mask_ratio=0.75 --norm_pix_loss=True --do_train=true --do_eval=true --base_learning_rate=1.5e-4 --lr_scheduler_type=cosine --weight_decay=0.05 --num_train_epochs=3 --warmup_ratio=0.05 --per_device_train_batch_size=8 --per_device_eval_batch_size=8 --logging_strategy=steps --logging_steps=30 --evaluation_strategy=epoch --save_strategy=epoch --load_best_model_at_end=True --save_total_limit=3 --seed=1337 --output_dir=MAE --overwrite_output_dir=true --logging_dir=./tensorboard-metrics --tpu_metrics_debug=true"

Elimina la TPU e la risorsa in coda

Elimina la TPU e la risorsa in coda alla fine della sessione.

   gcloud alpha compute tpus tpu-vm delete $TPU_NAME 
--project=${PROJECT_ID}
--zone=$ZONE
--quiet

   gcloud alpha compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} 
--project ${PROJECT_ID}
--zone ${ZONE}
--quiet

Risultato benchmark

La tabella seguente mostra le velocità effettiva dei benchmark per i diversi tipi di acceleratore.

V5E-4 V5E-16 V5E-64
Epoca 3 3 3
Dimensione batch globale 32 128 512
Velocità effettiva (esempi/sec) 201 657 2.844

Assistenza e feedback

Tutti i feedback sono ben accetti. Per condividere feedback o richiedere assistenza, contattaci qui o invia un'email all'indirizzo cloudtpu-preview-support@google.com

Termini

Tutte le informazioni che Google ti ha fornito in merito a questa Anteprima pubblica sono informazioni riservate di Google e sono soggette alle disposizioni di riservatezza previste dai Termini di servizio di Google Cloud Platform (o altri accordi che regolano l'utilizzo di Google Cloud Platform da parte dell'utente).