Addestramento Cloud TPU v5e

Cloud TPU v5e è l'acceleratore AI di ultima generazione di Google Cloud. Con un ingombro minimo di 256 chip per pod, v5e è ottimizzato per essere il prodotto con il valore più alto per l'addestramento, l'addestramento, l'ottimizzazione e la gestione della rete neurale convoluzionale (CNN), trasformatore, da testo a immagine e.

Concetti, architettura di sistema e configurazioni di Cloud TPU v5e

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

I concetti generali di Cloud TPU (ad esempio sezioni, host, chip e TensorCore) e l'architettura di sistema di Cloud TPU sono descritti nella pagina Architettura del sistema Cloud TPU.

Ogni versione di Cloud TPU richiede tipi di acceleratori specifici per l'addestramento e l'inferenza. Questi tipi di acceleratori sono descritti nella sezione Configurazioni v5e.

Inferenza

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

Sezioni

Una sezione rappresenta una raccolta di chip situati all'interno dello stesso pod, collegati da interconnessioni tra chip ad alta velocità (ICI). La versione v5e ha forme sezioni 2D. Per le forme di sezione supportate, consulta la tabella nella sezione Configurazioni v5e.

Anche forma chip e topologia chip fanno riferimento alle forme delle sezioni.

Pubblicazione

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.

Host singolo e multihost

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

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

Questa versione supporta l'addestramento multi-host e l'inferenza multi-host (utilizzando SAX).

VM TPU

Una macchina virtuale con Linux 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 è anche nota come worker.

Worker

Vedi TPU VM.

Inizia

Per informazioni sull'hardware TPU v5e, vedi Architettura di sistema.

Capacità di protezione

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 già fatto, crea un nuovo account.
  2. Nella console Google Cloud, seleziona o crea un progetto Google Cloud dalla pagina del selettore di progetti.

  3. La configurazione della fatturazione è obbligatoria per tutti gli utilizzi di Google Cloud, quindi assicurati che la fatturazione sia abilitata per il tuo progetto.

  4. Installa i componenti alpha di gcloud.

  5. Se sei un utente TPU che riutilizza 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 Google Cloud consigliata. 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 la scrittura dei log con l'API Logging
    • Writer metriche Monitoring: necessario per scrivere 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 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

La best practice prevede il provisioning di Cloud TPU v5 come risorse in coda utilizzando il comando queued-resource create. Tuttavia, puoi utilizzare anche l'API Create Node (gcloud alpha compute tpus tpu-vm create) per eseguire il provisioning di Cloud TPU v5.

Crea variabili di ambiente

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

Sostituisci le variabili (in rosso) nel seguente elenco con i valori che utilizzerai per il job di addestramento o di inferenza.

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-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
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d
Variabile Descrizione
PROJECT_ID Nome progetto Google Cloud
ACCELERATOR_TYPE Consulta la sezione Tipi di acceleratore per i tipi di acceleratore supportati.
MARGINE Tutta la capacità è in us-west4-a.
RUNTIME_VERSION Utilizza v2-alpha-tpuv5-lite.
SERVICE_ACCOUNT Questo è l'indirizzo del tuo account di servizio che puoi trovare nella console Google Cloud -> IAM -> Account di servizio. Ad esempio: [email protected]
TPU_NAME L'ID di testo assegnato dall'utente della TPU, creato al momento dell'allocazione della richiesta di risorse in coda.
QUEUED_RESOURCE_ID L'ID testo assegnato dall'utente della richiesta di risorsa in coda. Consulta il documento sulle risorse in coda per informazioni sulle risorse in coda.
QUOTA_TYPE Può essere reserved o best-effort. Se nessuna di queste opzioni è specificata, il valore predefinito di QUOTA_TYPE è on-demand. Consulta la pagina relativa alle quotas per informazioni sui diversi tipi di quote supportati da Cloud TPU.
VALID_UNTIL_DURATION Il periodo di validità della richiesta. Consulta le risorse in coda per informazioni sulle diverse durate valide.

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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

Se la risorsa in coda viene creata correttamente, lo stato all'interno del campo response sarà WAITING_FOR_RESOURCES o FAILED. Se la risorsa in coda è nello stato WAITING_FOR_RESOURCES, significa che ha superato la convalida preliminare ed è in attesa di capacità. Quando la capacità sarà disponibile, la richiesta passerà a PROVISIONING. Lo stato WAITING_FOR_RESOURCES non significa che riceverai la quota allocata e il passaggio dallo stato WAITING_FOR_RESOURCES allo stato ACTIVE potrebbe richiedere del tempo. Se la risorsa in coda è nello stato FAILED, il motivo dell'errore sarà riportato nell'output. La richiesta scadrà se non viene soddisfatta in --valid-until-duration e lo stato diventa "FAILED".

Potrai accedere alla VM TPU utilizzando SSH quando la risorsa in coda sarà nello stato ACTIVE.

Utilizza i comandi [list](/tpu/docs/managing-tpus-tpu-vm) o [describe](/tpu/docs/managing-tpus-tpu-vm) per eseguire una query sullo stato della risorsa in coda.

gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

Lo state rappresenta lo stato di una risorsa in coda. Gli stati sono definiti come:

Stato Descrizione
WAITING_FOR_RESOURCES Il comando create della risorsa in coda è stato ricevuto e avvierà 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 per l'uso. Se viene fornito uno script di avvio, inizierà l'esecuzione su tutte le TPU, quando lo stato della risorsa in coda passerà a ACTIVE.
FAILED Impossibile eseguire il provisioning delle sezioni.
SUSPENDING È in corso l'eliminazione di una o più sezioni.
SUSPENDED Tutte le sezioni sottostanti vengono eliminate, ma la risorsa in coda rimane intatta finché non viene eliminata esplicitamente. Al momento, una risorsa in coda sospesa non può essere ripresa e deve essere eliminata.
DELETING La risorsa in coda è in fase di eliminazione.

Connettiti alla VM TPU tramite SSH

La seguente sezione descrive come installare i 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 determinare quante VM avrà la tua sezione.

Per installare i programmi binari o eseguire codice, connettiti alla VM TPU utilizzando il comando ssh tpu-vm.

gcloud compute tpus tpu-vm ssh ${TPU_NAME}

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

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

Se le forme delle sezioni sono superiori a 8 chip, avrai più VM in una sola sezione. In questo caso, utilizza il flag --worker=all per eseguire l'installazione su tutte le VM TPU senza doverti connettere a ogni VM separatamente. Ad esempio:

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

Gestione

Tutti i comandi che puoi utilizzare per gestire le VM TPU sono descritti in Gestire le 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 è disponibile nelle versioni di runtime TPU tpu-vm-tf-2.15.0-pjrt e tpu-vm-tf-2.15.0-pod-pjrt.

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

Configura JAX

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

gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f http://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 veda i TensorCore Cloud TPU e sia in grado di eseguire operazioni di base:

gcloud 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 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 di PyTorch

Tieni presente che v5e supporta solo il runtime PJRT e PyTorch 2.1 e versioni successive utilizzerà PJRT come runtime predefinito per tutte le versioni di 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 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 libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f http://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f http://storage.googleapis.com/libtpu-releases/index.html'

Se non riesci a installare le ruote per Torch/torch_xla/torchvision e ricevi 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 versione con questo comando:

pip3 install setuptools==62.1.0

Esegui uno script con PJRT:

unset LD_PRELOAD

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

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker all \
   --command='
      export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.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/.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 di 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 compute tpus tpu-vm delete ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --quiet

gcloud 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 si trovano in stato FAILED.

Monitora e profila

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 Profilare il modello con gli strumenti di Cloud TPU per saperne di più sulla profilazione e sul monitoraggio delle VM Cloud TPU.

Esempi di 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=v5litepod-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
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Potrai stabilire una connessione SSH alla VM TPU quando la risorsa in coda si trova nello stato ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Quando QueuedResource è nello stato ACTIVE, l'output sarà simile al seguente:

     state: ACTIVE

  3. Installa la versione più recente di JAX e jaxlib:

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

  4. Clona il modello ImageNet e installa i requisiti corrispondenti:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone http://github.com/google/flax.git && cd flax/examples/imagenet && pip install -r requirements.txt && pip install flax==0.7.4'

  5. Per generare dati falsi, il modello ha bisogno di informazioni sulle dimensioni del set di dati. Questi dati possono essere recuperati dai metadati del set di dati ImageNet:

    gcloud 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 http://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

Dopo aver completato tutti i passaggi precedenti, puoi addestrare il modello.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd flax/examples/imagenet && JAX_PLATFORMS=tpu python3 imagenet_fake_data_benchmark.py'

Elimina la TPU e la risorsa in coda

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

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

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

Modelli 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ù diffusi.

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 addestrare un codificatore Transformer su ImageNet, ottenendo ottimi risultati rispetto alle reti convoluzionali. Per saperne di più, consulta le seguenti risorse:

Configura

  1. Crea le variabili di ambiente:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-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
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Potrai stabilire una connessione SSH alla VM TPU quando la risorsa in coda si trova nello stato ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Quando la risorsa in coda è nello stato ACTIVE, l'output sarà simile al seguente:

     state: ACTIVE

  3. Installa JAX e la relativa libreria:

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

  4. Scarica il repository Hugging Face e installa i requisiti:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone http://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==2.16.1 && pip install -r examples/flax/vision/requirements.txt'

  5. Scarica il set di dati Imagenette:

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

Addestra il modello

Addestra il modello con un buffer premappato a 4 GB.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd transformers && 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 TPU e risorsa in coda

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

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

gcloud 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 v5litepod-4, v5litepod-16 e v5litepod-64. La tabella seguente mostra le velocità effettiva con tipi di acceleratore diversi.

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

Allena la 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 stabile di diffusione è 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=v5litepod-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=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

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

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Quando la risorsa in coda è nello stato ACTIVE, l'output sarà simile al seguente:

     state: ACTIVE

  3. Installa JAX e la relativa libreria.

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

  4. Scarica il repository HuggingFace e installa i requisiti.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone http://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install tensorflow==2.16.1 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 compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --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 le risorse TPU e in coda

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

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

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

Risultati di benchmarking per la diffusione

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

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

Addestramento di 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. Era addestrato per prevedere la parola successiva all'interno di frasi. Per maggiori informazioni, consulta le seguenti risorse:

Configura

  1. Crea le variabili di ambiente:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-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=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

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

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Quando la risorsa in coda è nello stato ACTIVE, l'output sarà simile al seguente:

     state: ACTIVE

  3. Installa JAX e la relativa libreria.

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

  4. Scarica il repository HuggingFace e installa i requisiti.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='git clone http://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 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-tpuvm-artifacts/v5litepod-preview/jax/gpt .'

Addestra il modello

Addestra il modello con un buffer premappato a 4 GB.

gcloud 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 TPU e risorsa in coda

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

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

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

Risultati del benchmarking per GPT2

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

v5litepod-4 v5litepod-16 v5litepod-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. Ecco 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=v5litepod-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
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --{QUOTA_TYPE}

    Potrai accedere tramite SSH alla VM TPU quando la risorsa QueuedResource si trova nello stato ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Quando la risorsa in coda è nello stato ACTIVE, l'output sarà simile al seguente:

     state: ACTIVE

  3. Installa dipendenze specifiche per Torch/XLA

    gcloud 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 libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f http://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f http://storage.googleapis.com/libtpu-releases/index.html'
Addestra il modello ResNet
gcloud 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/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      git clone http://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 TPU e risorsa in coda

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

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

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

La tabella seguente mostra le velocità effettiva del benchmark.

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

Addestrare GPT2 su v5e

Questo tutorial illustra come eseguire GPT2 su v5e utilizzando il repository 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=v5litepod-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=queued_resource_id
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Potrai accedere tramite SSH alla VM TPU quando la risorsa QueuedResource è nello stato ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Quando la risorsa in coda è nello stato ACTIVE, l'output sarà simile al seguente:

    state: ACTIVE

  3. Installa le dipendenze Torch/xla.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      sudo apt-get -y update
      sudo apt install -y libopenblas-base
      pip3 install torchvision
      pip3 uninstall -y torch
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f http://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f http://storage.googleapis.com/libtpu-releases/index.html'

  4. Scarica il repository HuggingFace e installa i requisiti.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      git clone http://github.com/pytorch/xla.git
      pip install --upgrade accelerate
      git clone http://github.com/huggingface/transformers.git
      cd transformers
      git checkout ebdb185befaa821304d461ed6aa20a17e4dc3aa2
      pip install .
      git log -1
      pip install datasets evaluate scikit-learn
      '

  5. Scaricare le configurazioni del modello preaddestrato.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/
      gsutil cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'

Addestra il modello

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

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
      export PJRT_DEVICE=TPU_C_API
      cd transformers/
      export LD_LIBRARY_PATH=/usr/local/lib/
      export PT_XLA_DEBUG=0
      export USE_TORCH=ON
      python3 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 \
         --logging_dir=./tensorboard-metrics \
         --cache_dir=./cache_dir \
         --output_dir=/tmp/test-clm \
         --overwrite_output_dir \
         --cache_dir=/tmp \
         --config_name=examples/pytorch/language-modeling/my_config_2.json \
         --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 TPU e risorsa in coda

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

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

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

Risultato benchmark

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

v5litepod-4 v5litepod-16 v5litepod-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

Addestra ViT su v5e

Questo tutorial illustra come eseguire VIT su v5e utilizzando il repository 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=v5litepod-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
export QUOTA_TYPE=quota_type
export VALID_UNTIL_DURATION=1d

  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} \
   --valid-until-duration=${VALID_UNTIL_DURATION} \
   --service-account=${SERVICE_ACCOUNT} \
   --${QUOTA_TYPE}

    Potrai accedere tramite SSH alla VM TPU quando la risorsa QueuedResource è nello stato ACTIVE:

     gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

    Quando la risorsa in coda è nello stato ACTIVE, l'output sarà simile al seguente:

     state: ACTIVE

  3. Installa le dipendenze Torch/xla

    gcloud 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 libopenblas-dev -y
      pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f http://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f http://storage.googleapis.com/libtpu-releases/index.html'

  4. Scarica il repository HuggingFace e installa i requisiti.

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

Addestra il modello

gcloud 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/.local/lib/python3.10/site-packages/libtpu/libtpu.so
      cd vittransformers
      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 TPU e risorsa in coda

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

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

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

Risultato benchmark

La tabella riportata di seguito mostra le velocità effettiva di benchmark per i diversi tipi di acceleratore.

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

TensorFlow 2.x

Addestra Resnet su un singolo host v5e

Questo tutorial descrive come addestrare ImageNet su v5litepod-4 o v5litepod-8 utilizzando un set di dati falso. Se vuoi utilizzare un set di dati diverso, consulta Preparazione del set di dati.

Configura

  1. Crea le variabili di ambiente:

    export PROJECT_ID=your-project-ID
export ACCELERATOR_TYPE=v5litepod-4
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pjrt
export TPU_NAME=your-tpu-name
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

    ACCELERATOR_TYPE può essere v5litepod-4 o v5litepod-8. /

  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} \
   --${QUOTA_TYPE}

    Potrai stabilire una connessione SSH alla VM TPU quando la risorsa in coda si trova nello stato ACTIVE. Per verificare lo stato della risorsa in coda, utilizza il comando seguente:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  3. Connettiti alla TPU tramite SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  4. Imposta alcune variabili di ambiente

    export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so

  5. Passa alla directory del repository dei modelli e ai requisiti di installazione.

    cd ${MODELS_REPO} && git checkout r2.15.0
pip install -r official/requirements.txt

Addestra il modello

Esegui lo script di addestramento.

python3 official/vision/train.py \
   --tpu=local \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

Elimina TPU e risorsa in coda

  1. Elimina la TPU

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

  2. Elimina la richiesta di risorse in coda

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

Addestramento di Resnet su una v5e multi-host

Questo tutorial descrive come addestrare ImageNet su v5litepod-16 o più grandi utilizzando un set di dati falso. Se vuoi utilizzare un set di dati diverso, consulta Preparazione del set di dati.

  1. Crea le variabili di ambiente:

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-east1-c
export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pod-pjrt
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

    ACCELERATOR_TYPE può essere v5litepod-16 o superiore.

  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} \
   --${QUOTA_TYPE}

    Potrai stabilire una connessione SSH alla VM TPU quando la risorsa in coda si trova nello stato ACTIVE. Per verificare lo stato della risorsa in coda, utilizza il comando seguente:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  3. Connettiti alla tua TPU (nessun worker) tramite SSH

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project=${PROJECT_ID} \
   --zone=${ZONE}

  4. Imposta alcune variabili di ambiente

    export MODELS_REPO=/usr/share/tpu/models
export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
export MODEL_DIR=gcp-directory-to-store-model
export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
export TPU_LOAD_LIBRARY=0
export TPU_NAME=your_tpu_name

  5. Passa alla directory del repository dei modelli e ai requisiti di installazione.

     cd $MODELS_REPO && git checkout r2.15.0
 pip install -r official/requirements.txt

Addestra il modello

Esegui lo script di addestramento.

python3 official/vision/train.py \
   --tpu=${TPU_NAME} \
   --experiment=resnet_imagenet \
   --mode=train_and_eval \
   --model_dir=${MODEL_DIR} \
   --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*, task.validation_data.input_path=${DATA_DIR}/validation*"

Elimina TPU e risorsa in coda

  1. Elimina la TPU

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

  2. Elimina la richiesta di risorse in coda

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