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.

Concetti di Cloud TPU (ad esempio sezioni, host, e TensorCores) e l'architettura di sistema Cloud TPU per tutti i servizi Cloud TPU vengono descritte in questo modo Architettura del sistema Cloud TPU .

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

Gestisci 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 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. eseguire l'installazione su tutte le VM TPU con un solo comando:

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

Per controllare il numero di dispositivi, puoi eseguire il comando seguente: gli output mostrati qui sono stati prodotti con una sezione v5p-32). Questo codice verifica che tutto sia installato correttamente controllando che JAX vede i TensorCore di Cloud TPU e può eseguire le 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 per la sezione specificata. jax.local_device_count() indica 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.]

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 JAX di questo documento per ottenere è iniziato con l'addestramento v5p utilizzando JAX.

Configurazione di PyTorch

Il runtime PJRT è l'unico runtime supportato per v5p, mentre PyTorch 2.1+ utilizza 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
'

Usa uno script Python con PJRT per eseguire una convalida dell'installazione per mostrare i dispositivi TPU disponibili (gli output mostrati qui sono stati prodotti con una sezione 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 PyTorch di questo documento per ottenere è iniziato con l'addestramento 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 si concentra sui tutorial di formazione per una singola sezione. Adattare 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 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. Crea le variabili di ambiente: 

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

    Descrizioni flag di comando

    Variabile Descrizione
    PROJECT_ID Nome progetto Google Cloud
    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 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 creata quando viene allocata la richiesta di risorsa in coda.
    QUEUED_RESOURCE_ID L'ID testo assegnato dall'utente della richiesta di risorsa in coda. Consulta Il documento Queued Resources (Risorse in coda) per informazioni sulle risorse in coda.
    QUOTA_TYPE Può essere reserved o best-effort. Se nessuna di queste opzioni è specificata, QUOTA_TYPE il valore predefinito è on-demand. Vedi quotas 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.

  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 una volta che la risorsa in coda sarà 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

  3. Installa JAX e le rispettive dipendenze.

    # compatible with v5p: only jax version 0.4.19 and later \
# jax 0.4.19 requires py 3.10 \

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

  4. Scarica il repository HuggingFace e requisiti di installazione.

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

  5. Addestra il modello

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

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} --project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='export PATH=$PATH:$HOME/.local/bin && cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=stabilityai/stable-diffusion-2-1 --dataset_name=lambdalabs/pokemon-blip-captions --resolution=256 --center_crop --random_flip --train_batch_size=1 --mixed_precision=bf16 --max_train_steps=150 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model --from_pt'

Esegui la pulizia

Elimina la TPU e la richiesta di risorse in coda al termine della sessione o per rimuovi le richieste di risorse in coda nello stato "FAILED" stato. 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'architettura ad alte prestazioni, arbitrariamente scalabile, open source un LLM ben collaudato scritto in puro Python/JAX che ha come target le Cloud TPU. MaxText offre a ricercatori e sviluppatori un'accessibilità e adattabile per far avanzare le 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 tua 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 più sezioni:

    export NETWORK_NAME=your_network_name
export FIREWALL_RULE_NAME=your_firewall_rule_name

    Se esegui carichi di lavoro con più sezioni e vuoi per ottimizzare le prestazioni di rete, valuta la creazione di una rete dedicata con una unità massima di trasmissione (MTU) di 8896 byte e e configurare regole firewall appropriate. Sebbene sia facoltativo, questo passaggio migliorare significativamente le prestazioni, soprattutto quando si fa lo scale up numero di sezioni sulla rete del data center (DCN). Creazione della 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 che uno script gestisca l'intero flusso di lavoro, dal provisioning Cloud TPU e installazione delle dipendenze per l'esecuzione del modello e rimuovere le risorse, puoi usare 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 vedere un messaggio simile come segue nel log. Viene fatto riferimento alla posizione del log nel messaggio di output. Fai clic sul primo link per accedere di tutti i worker una volta completato il 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. Creare una risorsa TPU.

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

    Potrai connetterti alle VM TPU tramite SSH una volta QueuedResource è 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, 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 dei benchmark

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

N. di parametri

Tipo di acceleratore

TFLOP/chip/sec

Utilizzo flop 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 MLD

v5p-1024

2,94E+02

63,99%

1024 MLD

v5p-2048

2,49E+02

64,05%

1024 MLD

V5P-4096

2,97E+02

64,80%

1160 MLD

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 TensorFlow

Addestra ResNet su un singolo host v5p

Questo tutorial descrive come addestrare ImageNet su una TPU v5p-8 utilizzando un set di dati falso. Per 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=v5p-8
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, usa v5p-8 come ACCELERATOR_TYPE.

  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 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 questo 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.

Addestra ResNet su un 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

    Il valore di ACCELERATOR_TYPE può essere v5p-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 connetterti alla VM TPU tramite SSH una volta la risorsa in coda si trova lo stato ACTIVE.

    Usa 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 ID progetto, tipo di acceleratore, zona, 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 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_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}

    Potrai connetterti alla VM TPU tramite SSH una volta QueuedResource è 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. Installa Pytorch/XLA e le dipendenze obbligatorie.

    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 da 7 miliardi.

    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.

La SPMD_USER_GUIDE fornisce una guida dell'utente più approfondita che illustra tutti i diversi ambienti variabili e pulsanti di attivazione/disattivazione dello script HF. Tieni presente che LIBTPU_INIT_ARGS verrà incorporato in PyTorch/XLA e per impostazione predefinita nelle release future.

Esegui la pulizia

Elimina la TPU e le risorse in coda.

Risultati dei benchmark

La velocità effettiva per tutte e tre le dimensioni dei modelli Llama 2 è inclusa nelle seguenti tabelle .

V5p-8

V5p-128

V5p-128

Dimensione modello

7 MLD

13 MLD

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.