Addestra un modello utilizzando TPU v5e

Attenzione: questa pagina si applica all'API Cloud TPU, che non è più in fase di sviluppo attivo. L'API Cloud TPU riceverà solo correzioni di bug e aggiornamenti della sicurezza. Per le funzionalità più recenti e il supporto delle versioni più recenti delle TPU, ti consigliamo di utilizzare Compute Engine o Google Kubernetes Engine per gestire le risorse TPU. Per saperne di più, consulta Risorse TPU in Compute Engine e Informazioni sulle TPU in GKE. L'API Cloud TPU non è supportata su TPU7x e versioni successive di TPU.

Con un ingombro ridotto di 256 chip per pod, TPU v5e è ottimizzato per essere un prodotto di alto valore per l'addestramento, l'ottimizzazione e la distribuzione di Transformer, modelli da testo a immagine e reti neurali convoluzionali (CNN). Per saperne di più sull'utilizzo di Cloud TPU v5e per la distribuzione, consulta Inferenza utilizzando v5e.

Per ulteriori informazioni sull'hardware e sulle configurazioni di Cloud TPU v5e, consulta TPU v5e.

Inizia

Le sezioni seguenti descrivono come iniziare a utilizzare TPU v5e.

Richiedi quota

Per utilizzare TPU v5e per l'addestramento, devi avere una quota. Esistono diversi tipi di quote per TPU on demand, TPU prenotate e VM spot TPU. Se utilizzi TPU v5e per l'inferenza, sono necessarie quote separate. Per saperne di più sulle quote, consulta Quote. Per richiedere la quota per TPU v5e, contatta il team di vendita Cloud.

Crea un account e un progetto Google Cloud

Per utilizzare Cloud TPU, devi disporre di un account e di un progetto Google Cloud . Per saperne di più, consulta Configura un ambiente Cloud TPU.

Crea una Cloud TPU

La best practice consiste nell'eseguire il provisioning di Cloud TPU v5e come risorse in coda utilizzando il comando queued-resource create. Per saperne di più, consulta Gestisci le risorse in coda.

Puoi anche utilizzare l'API Create Node (gcloud compute tpus tpu-vm create) per eseguire il provisioning di Cloud TPU v5e. Per saperne di più, consulta Gestisci le risorse TPU.

Per saperne di più sulle configurazioni v5e disponibili per l'addestramento, consulta Tipi di Cloud TPU v5e per l'addestramento.

Configurazione del framework

Questa sezione descrive la procedura di configurazione generale per l'addestramento di modelli personalizzati utilizzando JAX con TPU v5e. Per le opzioni di addestramento di PyTorch, consulta Addestra un modello utilizzando PyTorch su Cloud TPU.

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

Definisci alcune variabili di ambiente:

export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5litepod-16
export ZONE=us-west4-a
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id

Configurazione per JAX

Se hai forme di slice con più di 8 chip, avrai più VM in uno slice. 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 separatamente:

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

Descrizioni dei flag dei comandi

• TPU_NAME: l'ID di testo assegnato dall'utente della TPU, che viene creato quando viene allocata la richiesta di risorsa in coda.
• PROJECT_ID:il nome del progetto Google Cloud . Utilizza un progetto esistente o creane uno nuovo in Configura il progetto Google Cloud
• ZONE: Consultare il documento Regioni e zone TPU per le zone supportate.
• worker: la VM TPU che ha accesso alle TPU sottostanti.

Puoi eseguire il seguente comando per controllare il numero di dispositivi (gli output mostrati qui sono stati generati con una sezione v5litepod-16). Questo codice verifica che tutto sia installato correttamente controllando che JAX veda i TensorCore di Cloud TPU e possa 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 di JAX in questo documento per iniziare con l'addestramento v5e utilizzando JAX.

Esegui la pulizia

Elimina la TPU e la risorsa in coda al termine 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 nello stato FAILED.

Esempi di JAX/FLAX

Le sezioni seguenti descrivono esempi di come addestrare i modelli JAX e FLAX su TPU v5e.

Addestra ImageNet su v5e

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

Configura

1. Crea variabili di ambiente:

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=us-west4-a
   export ACCELERATOR_TYPE=v5litepod-8
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your-service-account
   export QUEUED_RESOURCE_ID=your-queued-resource-id

   Descrizioni delle variabili di ambiente

   • PROJECT_ID: l'ID progetto Google Cloud . Utilizza un progetto esistente o creane uno nuovo.
   • TPU_NAME: il nome della TPU.
   • ZONE: la zona in cui creare la VM TPU. Per saperne di più sulle zone supportate, consulta Regioni e zone TPU.
   • ACCELERATOR_TYPE: il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni di TPU.
   • RUNTIME_VERSION: la versione software di Cloud TPU.
   • SERVICE_ACCOUNT: l'indirizzo email del account di servizio. Puoi trovarlo nella pagina Service account nella console Google Cloud .

     Ad esempio: tpu-service-account@PROJECT_ID.iam.

   • QUEUED_RESOURCE_ID: l'ID di testo assegnato dall'utente della richiesta di risorse in coda.

2. Crea una risorsa TPU:

   gcloud 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
   

   Nota: per utilizzare una prenotazione, aggiungi il flag --reserved. Per utilizzare le VM spot TPU, aggiungi il flag --spot.

   Potrai connetterti tramite SSH alla tua VM TPU quando la risorsa in coda è 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 -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

   Nota: se visualizzi un errore che suggerisce di utilizzare jax>=0.4.8, puoi ignorare il messaggio.

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 https://github.com/coolkp/flax.git && cd flax && git checkout pmap-orbax-conversion && git pull"
   

5. Per generare dati fittizi, il modello ha bisogno di informazioni sulle dimensioni del set di dati. Queste informazioni possono essere raccolte dai metadati del set di dati ImageNet:

   gcloud compute tpus tpu-vm ssh $TPU_NAME \
      --project=$PROJECT_ID \
      --zone=$ZONE \
      --worker=all \
      --command="cd flax/examples/imagenet && pip install -r requirements-cloud-tpu.txt"
   

Addestra il modello

Una volta completati 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 && bash ../../tests/download_dataset_metadata.sh && JAX_PLATFORMS=tpu python imagenet_fake_data_benchmark.py"

Elimina la TPU e la risorsa in coda

Elimina la TPU e la risorsa in coda al termine 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 Hugging Face

I modelli Hugging Face implementati in FLAX funzionano immediatamente su Cloud TPU v5e. Questa sezione fornisce istruzioni per l'esecuzione di modelli popolari.

Addestra ViT su Imagenette

Questo tutorial mostra come addestrare il modello Vision Transformer (ViT) di Hugging Face utilizzando il set di dati Imagenette di Fast AI su Cloud TPU v5e.

Il modello ViT è stato il primo ad addestrare correttamente un codificatore Transformer su ImageNet con risultati eccellenti rispetto alle reti convoluzionali. Per saperne di più, consulta la panoramica di ViT.

Configura

1. Crea variabili di ambiente:

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=us-west4-a
   export ACCELERATOR_TYPE=v5litepod-16
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your-service-account
   export QUEUED_RESOURCE_ID=your-queued-resource-id

   Descrizioni delle variabili di ambiente

   • PROJECT_ID: l'ID progetto Google Cloud . Utilizza un progetto esistente o creane uno nuovo.
   • TPU_NAME: il nome della TPU.
   • ZONE: la zona in cui creare la VM TPU. Per saperne di più sulle zone supportate, consulta Regioni e zone TPU.
   • ACCELERATOR_TYPE: il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni di TPU.
   • RUNTIME_VERSION: la versione software di Cloud TPU.
   • SERVICE_ACCOUNT: l'indirizzo email del account di servizio. Puoi trovarlo nella pagina Service account nella console Google Cloud .

     Ad esempio: tpu-service-account@PROJECT_ID.iam.

   • QUEUED_RESOURCE_ID: l'ID di testo assegnato dall'utente della richiesta di risorse in coda.

2. Crea una risorsa TPU:

   gcloud 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
   

   Nota: per utilizzare una prenotazione, aggiungi il flag --reserved. Per utilizzare le VM spot TPU, aggiungi il flag --spot.

   Potrai connetterti tramite SSH alla tua VM TPU quando la risorsa in coda è 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 -U "jax[tpu]" -f https://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 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==2.19 && sed -i "s/torchvision==0.12.0+cpu/torchvision==0.22.1/" examples/flax/vision/requirements.txt && pip install -r examples/flax/vision/requirements.txt && pip install tf-keras'
   

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 https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'
   

Addestra il modello

Addestra il modello con un buffer pre-mappato di 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 la TPU e la risorsa in coda

Elimina la TPU e la risorsa in coda al termine 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 benchmark ViT

Lo script di addestramento è stato eseguito su v5litepod-4, v5litepod-16 e v5litepod-64. La tabella seguente mostra i throughput con diversi tipi di acceleratore.

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

Addestra Diffusion su Pokémon

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

Il modello Stable Diffusion è un modello latente da testo a immagine che genera immagini fotorealistiche da qualsiasi input di testo. Per maggiori informazioni, consulta le seguenti risorse:

• Panoramica di Stable Diffusion
• Codice sorgente di Stable Diffusion

Configura

1. Imposta una variabile di ambiente per il nome del bucket di archiviazione:

   export GCS_BUCKET_NAME=your_bucket_name
   

2. Configura un bucket di archiviazione per l'output del modello:

   Google Cloud CLI storage buckets create gs://$GCS_BUCKET_NAME \
        --project=$PROJECT_ID \
        --location=$ZONE
   

3. Crea variabili di ambiente:

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=us-west1-c
   export ACCELERATOR_TYPE=v5litepod-16
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   export SERVICE_ACCOUNT=your-service-account
   export QUEUED_RESOURCE_ID=your-queued-resource-id

   Descrizioni delle variabili di ambiente

   • PROJECT_ID: l'ID progetto Google Cloud . Utilizza un progetto esistente o creane uno nuovo.
   • TPU_NAME: il nome della TPU.
   • ZONE: la zona in cui creare la VM TPU. Per saperne di più sulle zone supportate, consulta Regioni e zone TPU.
   • ACCELERATOR_TYPE: il tipo di acceleratore specifica la versione e le dimensioni della Cloud TPU che vuoi creare. Per ulteriori informazioni sui tipi di acceleratori supportati per ogni versione di TPU, consulta Versioni di TPU.
   • RUNTIME_VERSION: la versione software di Cloud TPU.
   • SERVICE_ACCOUNT: l'indirizzo email del account di servizio. Puoi trovarlo nella pagina Service account nella console Google Cloud .

     Ad esempio: tpu-service-account@PROJECT_ID.iam.

   • QUEUED_RESOURCE_ID: l'ID di testo assegnato dall'utente della richiesta di risorse in coda.

4. Crea una risorsa TPU:

   gcloud 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
   

   Nota: per utilizzare una prenotazione, aggiungi il flag --reserved. Per utilizzare le VM spot TPU, aggiungi il flag --spot.

   Potrai connetterti tramite SSH alla tua VM TPU quando la risorsa in coda sarà 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
   

5. 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 https://storage.googleapis.com/jax-releases/libtpu_releases.html'
   

6. 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 https://github.com/RissyRan/diffusers.git && cd diffusers && pip install . && pip install -U -r examples/text_to_image/requirements_flax.txt && pip install tensorflow==2.17.1 clu && pip install tensorboard==2.17.1'
   

Addestra il modello

Addestra il modello con un buffer pre-mappato di 4 GB.

gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --project=$PROJECT_ID --worker=all --command="
    git clone https://github.com/google/maxdiffusion
    cd maxdiffusion
    pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
    pip3 install -r requirements.txt
    pip3 install .
    pip3 install gcsfs
    export LIBTPU_INIT_ARGS=''
    python -m src.maxdiffusion.train src/maxdiffusion/configs/base_2_base.yml run_name=my_run \
    jax_cache_dir=gs://${GCS_BUCKET_NAME} activations_dtype=bfloat16 weights_dtype=bfloat16 \
    per_device_batch_size=1 precision=DEFAULT dataset_save_location=gs://${GCS_BUCKET_NAME} \
    output_dir=gs://${GCS_BUCKET_NAME}/ attention=flash"

Esegui la pulizia

Elimina la TPU, la risorsa in coda e il bucket Cloud Storage al termine della sessione.

1. Elimina la TPU:

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

2. Elimina la risorsa in coda:

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

3. Elimina il bucket Cloud Storage:

   gcloud storage rm -r gs://$GCS_BUCKET_NAME
   

Risultati del benchmark per la diffusione

Lo script di addestramento è stato eseguito su v5litepod-4, v5litepod-16 e v5litepod-64. La tabella seguente mostra i throughput.

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

Addestra un modello utilizzando PyTorch su Cloud TPU

Puoi utilizzare TorchTPU direttamente per addestrare il tuo modello o riutilizzare la libreria di riferimento TorchTitan. Per iniziare a utilizzare una VM TPU, consulta Inizia a utilizzare TorchTPU su una VM TPU. Per saperne di più sull'ecosistema PyTorch su Cloud TPU, consulta PyTorch su Cloud TPU.