Addestramento di Cloud TPU v5p

Cloud TPU v5p è la Cloud TPU di quinta generazione di Google Cloud e successore della TPU v4. La versione v5p è ottimizzata per l'addestramento su larga scala ed essere una piattaforma leader per lo sviluppo di modelli LLM, modelli di diffusione e AI generativa. A livello generale, v5p offre prestazioni fino al doppio rispetto alla versione v4, TPU in più in un pod (sezione più grande di 6.000 rispetto alle 3.000 della versione v4), con prestazioni fino a 4 volte superiori a livello di pod. Viene inoltre eseguito a una velocità frequenza di clock (1.75Ghz contro 1.05Ghz), aggiunge SparseCore per incorporamenti su larga scala e triplica la capacità di memoria a larghezza di banda elevata (HBM).

Concetti di Cloud TPU v5p

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

I concetti di Cloud TPU (ad esempio, slice, host e TensorCore) e l'architettura di sistema di Cloud TPU per tutte le versioni di Cloud TPU sono descritti nella pagina Architettura di sistema di Cloud TPU.

Ogni versione di Cloud TPU richiede tipi di acceleratori specifici per la formazione o l'inferenza. Questi tipi di acceleratori sono descritti nelle configurazioni v5p.

Gestire le risorse TPU

Tutti i comandi che puoi utilizzare per gestire le VM TPU sono descritti in Gestione delle TPU o Guida dell'utente alle risorse in coda per la gestione in coda.

Configurazione framework

Questa sezione descrive la procedura di configurazione generale per l'addestramento del modello utilizzando JAX o PyTorch con TPU v5p.

Configurazione per JAX

Se hai forme di sezioni maggiori di 4 chip, avrai più VM in un'unica sezione. In questo caso, devi utilizzare il flag --worker=all per eseguire l'installazione su tutte le VM TPU utilizzando un unico comando:

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'

Puoi eseguire il seguente comando per controllare il numero di dispositivi (gli output mostrati qui sono stati prodotti con uno slice v5p-32). 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 in la sezione specificata. jax.local_device_count() indica il numero di chip accessibili da una singola VM in questo slice.

# 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.]

Utilizza --node=all per eseguire il comando su tutti i worker Multislice.

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

Prova i tutorial su JAX in questo documento per iniziare con l'addestramento di v5p utilizzando JAX.

Configurazione per PyTorch

Il runtime PJRT è l'unico runtime supportato per v5p, mentre PyTorch 2.1+ usa PJRT come runtime predefinito per tutte le versioni di TPU. Questa sezione descrive come iniziare a utilizzare PJRT su pod v5p con PyTorch/XLA 2.2.0 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
sudo apt-get install libopenblas-dev -y
pip3 install numpy
pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html
'

Utilizza uno script Python con PJRT per convalidare l'installazione e mostrare i dispositivi TPU disponibili (gli output mostrati qui sono stati prodotti con una fetta v5p-32).

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} --zone ${ZONE} --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'
 
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...
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']

Utilizza --node=all per eseguire il comando su tutti i worker Multislice.

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

Prova i tutorial di PyTorch in questo documento per iniziare con l'addestramento di v5p utilizzando PyTorch.

Monitoraggio e profilo

Cloud TPU v5p supporta il monitoraggio e la profilazione utilizzando gli stessi metodi delle precedenti generazioni di Cloud TPU. Puoi leggi Profilare il modello con gli strumenti di Cloud TPU per saperne di più sulla profilazione e sul monitoraggio delle VM Cloud TPU per scoprire di più sul monitoraggio.

Tutorial di formazione

Questa sezione è incentrata sui tutorial di formazione per un singolo segmento. L'adattamento di questi tutorial all'addestramento multislice può essere ottenuto aggiungendo il flag --node=all ai comandi SSH. Per dettagli e best practice, consulta Introduzione a più sezioni.

Tutorial su JAX

Train Diffusion 2.1

Questo tutorial mostra come addestrare il modello Stable Diffusion HuggingFace utilizzando Pokémon su Cloud TPU v5p.

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

Configura

  1. Configura un bucket di archiviazione per l'output del tuo modello.

    gcloud storage buckets create gs://your_bucket \
 --project=your_project \
 --location=us-east5-a

  2. Crea le variabili di ambiente:

    export GCS_BUCKET_NAME=your-bucket
export PROJECT_ID=your-project-ID
export ACCELERATOR_TYPE=v5p-32
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export SERVICE_ACCOUNT=your-service-account
export TPU_NAME=your-tpu-name
export QUEUED_RESOURCE_ID=your-qr-name
export QUOTA_TYPE=spot
export VALID_UNTIL_DURATION=1d

    Descrizioni flag di comando

    Variabile Descrizione
    PROJECT_ID Nome del progetto Google Cloud
    ACCELERATOR_TYPE Consulta la pagina Versioni TPU per la tua versione TPU.
    ZONA Consulta il documento Regioni e zone TPU per le zone supportate.
    RUNTIME_VERSION Per la versione v5p, utilizza la versione v2-alpha-tpuv5 per RUNTIME_VERSION.
    SERVICE_ACCOUNT Si tratta dell'indirizzo del tuo 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 che viene creato quando viene allocata la richiesta di risorse in coda.
    QUEUED_RESOURCE_ID L'ID testo assegnato dall'utente della richiesta di risorsa in coda. Per informazioni sulle risorse in coda, consulta il documento Risorse in coda.
    QUOTA_TYPE Può essere reserved o spot. Se nessuna di queste opzioni è specificata, QUOTA_TYPE il valore predefinito è on-demand. Consulta la sezione quote per informazioni sui diversi tipi di quote supportati da Cloud TPU.
    VALID_UNTIL_DURATION Il periodo di validità della richiesta. Consulta: Risorse in coda per informazioni sulle diverse durate valide.

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

    Potrai accedere tramite SSH alla VM TPU una volta che la risorsa in coda nello stato ACTIVE. Controlla lo stato della risorsa in coda eseguendo questo comando:

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

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

    state: ACTIVE

  4. Addestra il modello

    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
git reset --hard 57629bcf4fa32fe5a57096b60b09f41f2fa5c35d # This identifies the GitHub commit to use.
pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html # Install the latest version of JAX
pip3 install -r requirements.txt
pip3 install .
export LIBTPU_INIT_ARGS=""
python -m src.maxdiffusion.models.train src/maxdiffusion/configs/base_2_base.yml run_name=my_run base_output_directory=gs://$GCS_BUCKET_NAME enable_profiler=False"

Esegui la pulizia

Eliminare la richiesta di TPU e risorse in coda al termine della sessione o rimuovere le richieste di risorse in coda nello stato "NON ANDATA A BUON FINE". Per eliminare una risorsa in coda, elimina le sezioni 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

In alternativa, utilizza --force per eliminare le sezioni e la richiesta di risorse in coda in un solo passaggio:

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

Risultati dei benchmark

Lo script di addestramento di Stable Diffusion veniva eseguito sulle versioni v5p-8, v5p-32 e v5p-128. La tabella seguente mostra la velocità effettiva.

V5p-8

V5p-32

v5p-128

Passaggio di addestramento

150

150

150

Dimensione batch globale

32

64

64

Velocità effettiva (esempi/sec)

12:10

18:08

19.10

MaxText

Questo tutorial mostra come addestrare MaxText utilizzando un modello sintetico su Cloud TPU.

MaxText è un LLM open source, ad alte prestazioni, arbitrariamente scalabile e ben testato, scritto in puro Python/JAX e indirizzato alle Cloud TPU. MaxText offre a ricercatori e sviluppatori un'accessibilità e adattabile per avanzare alle frontiere del linguaggio naturale ricerca e sviluppo nell'ambito dell'elaborazione (NLP).

Prima di eseguire questo tutorial, devi per configurare l'ambiente Cloud TPU.

  1. Imposta le variabili di ambiente

    export PROJECT_ID=your_project_ID
export TPU_NAME=your_tpu_name # user defined TPU name
export ACCELERATOR_TYPE=v5p-256
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export RUN_NAME=your_experiment_run_name # user defined name for this run
export GCS_BUCKET_NAME=your_bucket_name # Output cloud folder. Should start with gs://
export MAXTEXT_OUTPUT_PATH=${GCS_BUCKET_NAME}/your_experiment_output_path
export NUM_SLICES=1 # Update the value to a number >1 for Multislice.

    Descrizioni flag di comando

    Variabile Descrizione
    PROJECT_ID Nome progetto Google Cloud
    TPU_NAME Un nome definito dall'utente per la TPU.
    ACCELERATOR_TYPE Consulta le versioni TPU per la tua versione di TPU.
    ZONA Consulta il documento su regioni e zone TPU per conoscere le zone supportate.
    RUNTIME_VERSION Per v5p usare v2-alpha-tpuv5 per la versione runtime.
    RUN_NAME Nome esecuzione esperimento fornito dall'utente.

    Configurazione facoltativa consigliata per Multislice:

    export NETWORK_NAME=your_network_name
export FIREWALL_RULE_NAME=your_firewall_rule_name

    Se esegui carichi di lavoro Multislice e vuoi ottimizzare le prestazioni di rete, ti consigliamo di creare una rete dedicata con un'unità di trasmissione massima (MTU) di 8896 byte e di configurare regole firewall appropriate. Anche se facoltativo, questo passaggio può migliorare notevolmente le prestazioni, in particolare quando si aumenta il numero di slice sulla rete del data center (DCN). Creazione nota in corso... una rete richiede l'autorizzazione compute.networks.create nel progetto. I seguenti esempi mostrano come creare una rete e un firewall dedicati personalizzata.

    Crea una rete dedicata:

    gcloud compute networks create ${NETWORK_NAME} \
--mtu=8896 \
--project=${PROJECT_ID} \
--subnet-mode=auto \
--bgp-routing-mode=regional

    Crea una regola firewall:

    gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
--network ${NETWORK_NAME} --allow tcp,icmp,udp --project=${PROJECT_ID}

  2. clona il repository MaxText

    git clone https://github.com/google/maxtext.git

  3. Addestra il modello

    Le sezioni seguenti descrivono due opzioni per l'addestramento di MaxText.

    Opzione 1

    Se vuoi uno script per gestire l'intero flusso di lavoro, dal provisioning delle TPU Cloud e dall'installazione delle dipendenze all'esecuzione del modello e al ritiro delle risorse, puoi utilizzare multihost_job.py.

    cd maxtext && python3 multihost_job.py --PROJECT=${PROJECT_ID} --ZONE=${ZONE} \
--NUM_SLICES=${NUM_SLICES} --TPU_TYPE=${ACCELERATOR_TYPE} \
--VERSION=${RUNTIME_VERSION} --RUN_NAME=${RUN_NAME} #user defined run name \
--BUCKET_NAME=${GCS_BUCKET_NAME} \ #used to store logs and configs
--COMMAND="bash setup.sh && bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME} OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

    Dopo aver avviato lo script, dovresti visualizzare un messaggio simile al seguente nel log. Viene fatto riferimento alla posizione del log nel messaggio di output. Fai clic sul primo link per accedere ai log di tutti i worker al termine del provisioning delle TPU.

    ------------------------------------

multihost_job finished running, TPUs are starting up to run your job remotely.

Logs for your job are displayed here:
https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22_log%22%2529;?project=PROJECT_ID

To see the output of a single host, you may edit the slice and worker
number in the `log_file_path` property here:

https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22RUN_NAME_log%22%2529%20AND%0Alabels.%22agent.googleapis.com%2Flog_file_path%22%3D%20%22%2FRUN_NAME%2Fmain_command_log_slice_0_worker_0%22;?project=PROJECT_ID

When your job is finished, the main command log is in your Cloud Storage
bucket:
https://console.cloud.google.com/storage/browser/YOUR_BUCKET_NAME/RUN_NAME?project=PROJECT_ID

View the status of the created TPUs using:
gcloud compute tpus queued-resources list --filter=RUN_NAME
--zone=ZONE --project=PROJECT_ID
Opzione 2

Eseguire lo script di addestramento più volte su un ambiente Cloud TPU, utilizza lo script multihost_runner.py per utilizzare la risorsa.

  1. Configura le variabili per creare una TPU.

    export SERVICE_ACCOUNT=your_service_account
export TPU_NAME=your_tpu_name
export QUEUED_RESOURCE_ID=your_queued_resource_id
export VALID_DURATION=1d
export QUOTA_TYPE=quota_type
     
    --node-count ${NODE_COUNT} \
--node-prefix ${NODE_PREFIX} # optional, the default is QUEUED_RESOURCE_ID

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

    Potrai connetterti alle VM TPU tramite SSH quando il tuo QueuedResource sarà nello stato ACTIVE:

    Utilizza il comando describe per eseguire una query sullo stato della risorsa in coda.

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

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

     state: ACTIVE

  3. Connettiti alla TPU tramite SSH

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

  4. Installa le dipendenze

    export TPU_NAME=your_tpu_name
export MAXTEXT_OUTPUT_PATH=output-path
     
    cd maxtext && python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
--COMMAND='bash setup.sh'

  5. Esegui il modello con vari script di configurazione, ad esempio 32b.sh, 64b.sh. Se esegui lo script da una VM TPU, devi aggiungere il flag --INTERNAL_IP=true.

    python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
--COMMAND="bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME}
OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

Esegui la pulizia

Elimina la TPU e le risorse in coda.

Risultati del benchmark

Lo script di addestramento di MaxText è stato eseguito da 32 B a 1160 B con precisione bf16. I risultati di queste esecuzioni sono riportati nella tabella seguente.

N. di parametri

Tipo di acceleratore

TFLOP/chip/sec

Utilizzo dei FLOPS del modello

(MFU)

32 MLD

V5p-128

3,28E+02

71,47%

64 MLD

v5p-128

3,23E+02

70,31%

128 MLD

V5p-256

3,15E+02

68,68%

128 MLD

v5p-512

3,15E+02

68,53%

256 MLD

v5p-1024

3,16E+02

68,82%

512 B

v5p-1024

2,94E+02

63,99%

1024 B

v5p-2048

2,49E+02

64,05%

1024 B

V5P-4096

2.97E+02

64,80%

1160 B

V5P-7680

2,95E+02

64,27%

1160 MLD

v5p-12288

3,04E+02

66,23%

Il modello di parametri 256B è stato testato su v5p-512 e v5p-1024 usando entrambi i pesi bf16 e int8. La tabella seguente mostra i risultati di questi test.

v5p-512

v5p-512

v5p-1024

v5p-1024

Dimensione batch globale

(token)

5,24E + 05

5,24E + 05

1,05E + 06

1,05E+06

Precisione

bf16

int8

bf16

int8

TFLOP/chip/sec

307

408

308

414

Utilizzo flop modello

(MFU)

66,98%

88,85%

67,09%

90,23%

Tutorial su TensorFlow

Addestramento di ResNet su un singolo host v5p

Questo tutorial descrive come addestrare ImageNet su una v5p-8 TPU utilizzando un set di dati falso. Se vuoi utilizzare un altro set di dati, consulta la sezione Preparazione del set di dati.

Configura

  1. Crea le variabili di ambiente:

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

    Per questo tutorial, utilizza v5p-8 come ACCELERATOR_TYPE.

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

    Potrai connetterti alla VM TPU tramite SSH una volta la risorsa in coda si trova lo stato ACTIVE. Per controllare lo stato della risorsa in coda, utilizza il seguente comando:

    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. Modifica la directory del repository dei modelli e installa i requisiti.

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

Addestra il modello

  1. 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"

Esegui la pulizia

Elimina la TPU e le risorse in coda.

Addestramento di ResNet su un cluster v5p multi-host

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

  1. Crea le variabili di ambiente:

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

    ACCELERATOR_TYPE può essere v5p-16 o superiore.

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

    Potrai connetterti alla VM TPU tramite SSH una volta la risorsa in coda si trova lo stato ACTIVE.

    Utilizza il comando describe per eseguire una query sullo stato della risorsa in coda:

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

  3. Connettiti alla TPU (worker zero) tramite SSH

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

  4. Imposta alcune variabili di ambiente

    export TPU_NAME=your_tpu_name
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

  5. Modifica la directory del repository dei modelli e installa i requisiti.

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

Addestra il modello

  1. Esegui lo script di addestramento.

    python3 official/vision/train.py \
  --tpu=${TPU_NAME} \
  --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"

Esegui la pulizia

Elimina la TPU e le risorse in coda.

PyTorch/XLA

Llama 2

Questo tutorial illustra come addestrare il modello Llama 2 7B su v5p utilizzando un fork del repository HuggingFace su PyTorch/XLA con le opzioni General e Scalable Parallelizzazione per i grafici GSPMD (ML Computation Graphs).

Configurazione

  1. Crea variabili per l'ID progetto, il tipo di acceleratore, la zona, la versione di runtime e il nome della TPU.

    export PROJECT_ID=your_project_ID
export ACCELERATOR_TYPE=v5p-8
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
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_DURATION=1d

  2. Creare 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} \
--valid-until-duration ${VALID_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}

    Potrai connetterti alla VM TPU tramite SSH quando il tuo QueuedResource sarà nello stato ACTIVE:

    Utilizza il comando describe per eseguire una query sullo stato della risorsa in coda.

    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 Pytorch/XLA e le dipendenze richieste.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
--project ${PROJECT_ID} \
--zone  ${ZONE} \
--worker=all \
--command='
sudo apt-get update
sudo apt-get install libopenblas-dev -y
pip3 install numpy
pip3 install typing-extensions
pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f https://storage.googleapis.com/libtpu-releases/index.html
'

  4. Scarica il repository HuggingFace e i requisiti di installazione.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME}
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='
git clone -b llama2-google-next-training https://github.com/pytorch-tpu/transformers.git
cd transformers
pip3 install git+file://$PWD
pip3 install datasets accelerate evaluate scikit-learn'

  5. Scarica la configurazione del modello 7B.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command="curl https://huggingface.co/TheBloke/Llama-2-7B-fp16/raw/main/config.json --output ~/config.json"

  6. Addestra il modello

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='
export PJRT_DEVICE=TPU
export XLA_USE_BF16=1
export XLA_IR_DEBUG=1
export XLA_HLO_DEBUG=1

export LIBTPU_INIT_ARGS="--xla_enable_async_collective_permute=true
--xla_tpu_enable_async_collective_fusion_multiple_steps=true
--xla_tpu_enable_async_collective_fusion=true
--xla_tpu_overlap_compute_collective_tc=true
--xla_enable_async_all_gather=true
--xla_jf_spmd_threshold_for_windowed_einsum_mib=0"

export PROFILE_EPOCH=0
export PROFILE_STEP=3
export PROFILE_DURATION_MS=20000
export PROFILE_LOGDIR=/tmp/home/

cd transformers
python examples/pytorch/language-modeling/run_clm.py \
 --tokenizer_name hf-internal-testing/llama-tokenizer \
 --dataset_name wikitext \
 --dataset_config_name wikitext-2-raw-v1 \
 --per_device_train_batch_size 96 \
 --per_device_eval_batch_size 8 \
 --num_train_epochs 1 \
 --do_train \
 --output_dir /tmp/output \
 --overwrite_output_dir \
 --config_name ~/config.json \
 --save_strategy no \
 --logging_strategy no \
 --remove_unused_columns no \
 --optim adafactor \
 --torch_dtype bfloat16 \
 --dataloader_drop_last yes \
 --block_size 2048 \
 --spmd_2d_sharding 1 \
 --spmd_grad_chkpt
'

Se l'ambiente è in esecuzione in più sezioni, devi imposta il flag --spmd_dcn_parallelism sul numero di sezioni.

SPMD_USER_GUIDE fornisce una guida utente più approfondita che spiega tutte le diverse variabili e i diversi pulsanti di attivazione/disattivazione dell'ambiente dello script HF. Tieni presente che LIBTPU_INIT_ARGS verrà incorporato in PyTorch/XLA e attivato per impostazione predefinita nelle release future.

Esegui la pulizia

Elimina la TPU e le risorse in coda.

Risultati dei benchmark

La tabella seguente include il throughput per tutti e tre i modelli di Llama 2.

v5p-8

v5p-128

v5p-128

Dimensione modello

7 miliardi

13B

70 mld

Dimensione batch globale

96

1024

128

Forma mesh di sharding

(4, 1)

(64, 1)

(16, 4)

Utilizzo flop modello

(MFU)

56,67%

55,80%

51,85%

Assistenza e feedback

Tutti i feedback sono ben accetti. Per condividere feedback o richiedere assistenza: compila il modulo di feedback o assistenza per Cloud TPU.