Questa pagina è stata tradotta dall'API Cloud Translation.

Addestramento Cloud TPU v5p

Cloud TPU v5p è la Cloud TPU di quinta generazione di Google Cloud e il successore della v4 TPU. v5p è ottimizzata per l'addestramento su larga scala ed è una piattaforma leader per lo sviluppo di modelli LLM di base, modelli di diffusione e AI generativa. A livello generale, v5p offre fino al doppio delle prestazioni di v4, pacchettizzando al contempo il doppio di TPU in un pod (sezione più grande di 6000 rispetto a 3000 nella v4), per un rendimento fino a 4 volte superiore a livello di pod. Funziona inoltre con una frequenza di orologio più elevata (1,75 GHz rispetto a 1,05 GHz), aggiunge SparseCore per gli incorporamenti su larga scala e triplica la capacità HBM (High Bandwidth Memory).

Concetti di Cloud TPU v5p

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

I concetti di Cloud TPU (ad esempio sezioni, host e TensorCore) e l'architettura di sistema di Cloud TPU per tutte le versioni 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 o l'inferenza. Questi tipi di acceleratori sono descritti nelle configurazioni v5p.

Gestisci risorse TPU

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

Configurazione del framework

Questa sezione descrive il processo di configurazione generale per l'addestramento di modelli utilizzando JAX o PyTorch con TPU v5p.

Configura JAX

Se le forme delle sezioni sono più grandi di 4 chip, avrai più VM in una sezione. In questo caso, devi utilizzare il flag --worker=all per eseguire l'installazione su tutte le VM TPU utilizzando un singolo 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'

Puoi eseguire il seguente comando per controllare il numero di dispositivi (gli output mostrati qui sono stati prodotti con una sezione v5p-32). Questo codice verifica che tutto sia installato correttamente controllando che JAX rilevi i TensorCore di 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.]

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

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 in questo documento per iniziare con l'addestramento v5p utilizzando JAX.

Configurazione di PyTorch

Il runtime PJRT è l'unico runtime supportato per v5p e PyTorch 2.1 e versioni successive utilizza PJRT come runtime predefinito per tutte le versioni di TPU. Questa sezione descrive come iniziare a utilizzare PJRT sui 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 eseguire una convalida dell'installazione in modo da 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 multisezione.

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 l'addestramento v5p utilizzando PyTorch.

Monitora e profila

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

Tutorial di formazione

Questa sezione si concentra sui tutorial di formazione per un'unica sezione. Puoi adattare questi tutorial all'addestramento Multislice aggiungendo il flag --node=all ai comandi SSH. Per dettagli e best practice, consulta l'introduzione a Multislice.

Diffusione 2.1
MaxText
ResNet su un singolo host v5p
ResNet su una versione 5p multi-host

Tutorial JAX

Addestramento della diffusione 2.1

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

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

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 comando

Variabile	Description
PROJECT_ID	Nome progetto Google Cloud
ACCELERATOR_TYPE	Consulta la pagina delle versioni TPU per la versione di TPU.
MARGINE	Consulta il documento Regioni e zone TPU per conoscere le zone supportate.
RUNTIME_VERSION	Per v5p, utilizza v2-alpha-tpuv5 per RUNTIME_VERSION.
SERVICE_ACCOUNT	Questo è l'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 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 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}
```
Nota: il flag QUOTA_TYPE può essere reserved o best-effort. Se nessuna di queste opzioni è specificata, il valore predefinito è on-demand. Vedi quotas per informazioni sui diversi tipi di quote supportati da Cloud TPU.

Potrai connetterti tramite SSH alla VM TPU quando 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 è nello stato ACTIVE, l'output sarà simile al seguente:
```
state: ACTIVE
```

Installa JAX e le sue 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'

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 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'

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='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 richiesta di risorse TPU e in coda al termine della sessione o per rimuovere le richieste di risorse in coda che sono nello stato "Non riuscito". Per eliminare una risorsa in coda, elimina le sezioni, quindi la richiesta di 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 unico passaggio:

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

Risultati benchmark

Lo script di addestramento per la diffusione stabile è stato eseguito su 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 il modello MaxText utilizzando un set di dati sintetico su Cloud TPU.

MaxText è un LLM open source, ben testato e ad alte prestazioni, arbitrario e scritto in Python/JAX puro con targeting per Cloud TPU. MaxText offre a ricercatori e sviluppatori uno strumento accessibile e adattabile per avanzare le frontiere della ricerca e dello sviluppo dell'elaborazione del linguaggio naturale (NLP).

Prima di eseguire questo tutorial, devi configurare il tuo ambiente Cloud TPU.

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 comando

Variabile	Description
PROJECT_ID	Nome progetto Google Cloud
TPU_NAME	Un nome definito dall'utente per la tua TPU.
ACCELERATOR_TYPE	Consulta la pagina delle versioni TPU per la versione di TPU.
MARGINE	Consulta il documento 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 il multislice:

export NETWORK_NAME=your_network_name
export FIREWALL_RULE_NAME=your_firewall_rule_name

Se esegui carichi di lavoro multisezione e vuoi prestazioni di rete ottimali, valuta la possibilità di creare una rete dedicata con una massima unità di trasmissione (MTU) di 8896 byte e di configurare regole firewall appropriate. Sebbene facoltativo, questo passaggio può migliorare notevolmente le prestazioni, soprattutto in caso di scale up del numero di sezioni sulla rete dei data center (DCN). Nota: per creare una rete è necessaria l'autorizzazione compute.networks.create nel progetto. I seguenti esempi mostrano come creare una rete e una regola firewall dedicate.

Creare 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}

clona il repository MaxText

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

Addestra il modello

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

Opzione 1

Se vuoi che uno script gestisca l'intero flusso di lavoro, dal provisioning delle Cloud TPU all'installazione delle dipendenze, fino all'esecuzione del modello e all'eliminazione 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 l'avvio dello script, nel log dovresti visualizzare un messaggio simile al seguente. Si fa riferimento alla posizione del log nel messaggio di output. Fai clic sul primo link per accedere ai log di tutti i worker una volta completato il provisioning della 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

Per eseguire più volte lo script di addestramento su una Cloud TPU di cui è stato eseguito il provisioning, utilizza lo script multihost_runner.py per utilizzare la risorsa.

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
```
Nota: i seguenti flag sono necessari solo se utilizzi Multislice:
```
--node-count ${NODE_COUNT} \
--node-prefix ${NODE_PREFIX} # optional, the default is QUEUED_RESOURCE_ID
```
Nota: se specifichi --node-count e --node-prefix, non specificare anche --node-id quando crei la risorsa TPU nel passaggio seguente; verrà generato un errore se vengono specificati tutti e tre i flag.
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}
```
Nota: il flag QUOTA_TYPE può essere reserved o best-effort. Se nessuna di queste opzioni è specificata, il valore predefinito è on-demand. Vedi quotas per informazioni sui diversi tipi di quote supportati da Cloud TPU.

Potrai connetterti alle VM TPU utilizzando 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
```

Connettiti alla TPU tramite SSH

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

Installa le dipendenze

Nota: dopo aver utilizzato ssh per connetterti alla VM TPU, devi impostare le variabili ${TPU_NAME} e ${MAXTEXT_OUTPUT_PATH} per utilizzarle nella VM TPU:
```
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'
```
Esegui il modello con vari script di configurazione, ad esempio 32b.sh e 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"
```
Nota: imposta --TPU_PREFIX=${NODE_PREFIX} per multisezione anziché --TPU_PREFIX=${TPU_NAME}.

Esegui la pulizia

Elimina le risorse TPU e in coda.

Risultati benchmark

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

N. di parametri	Tipo di acceleratore	TFLOP/chip/sec	Utilizzo flop 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 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%
1.160 MLD	v5p-7680	2,95E + 02	64,27%
1.160 MLD	v5p-12288	3,04E+02	66,23%

Il modello con parametro 256B è stato testato su v5p-512 e v5p-1024 utilizzando i pesi bf16 e int8. La seguente tabella 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 del modello (MFU)	66,98%	88,85%	67,09%	90,23%

Tutorial su 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. Se vuoi utilizzare un set di dati diverso, consulta Preparazione del set di dati.

Configura

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

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}
```
Nota: il flag QUOTA_TYPE può essere reserved o best-effort. Se nessuna di queste opzioni è specificata, il valore predefinito è on-demand. Vedi quotas per informazioni sui diversi tipi di quote supportati da Cloud TPU.

Potrai connetterti alla VM TPU utilizzando SSH 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}
```

Connettiti alla TPU tramite SSH

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

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

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"

Esegui la pulizia

Elimina le risorse TPU e in coda.

Addestramento di ResNet su una versione v5p multi-host

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

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.16.1-pod-pjrt
export QUEUED_RESOURCE_ID=your-queued-resource-id
export QUOTA_TYPE=quota-type

ACCELERATOR_TYPE può essere v5p-16 o superiore.

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}
```
Nota: il flag QUOTA_TYPE può essere reserved o best-effort. Se nessuna di queste opzioni è specificata, il valore predefinito è on-demand. Vedi quotas per informazioni sui diversi tipi di quote supportati da Cloud TPU.

Potrai connetterti alla VM TPU utilizzando SSH quando la risorsa in coda si trova 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}
```

Connettiti alla tua TPU (nessun worker) tramite SSH

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

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

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 \
  --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 le risorse TPU e 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 General and Scalable Parallelization for ML Computation Graphs (GSPMD).

Imposta

Crea variabili per ID progetto, tipo di acceleratore, zona, versione di runtime e 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

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_DURATION} \
--service-account ${SERVICE_ACCOUNT} \
--${QUOTA_TYPE}
```
Nota: il flag QUOTA_TYPE può essere reserved o best-effort. Se nessuna di queste opzioni è specificata, il valore predefinito è on-demand. Vedi quotas per informazioni sui diversi tipi di quote supportati da Cloud TPU.

Potrai connetterti alla VM TPU utilizzando SSH quando il tuo QueuedResource si trova 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
```
Nota: se l'esecuzione avviene in un ambiente multisezione, utilizza gcloud compute tpus queued-resources ssh --node=all ... anziché gcloud compute tpus tpu-vm ssh ... per i comandi. Ciò garantisce che tutti i chip TPU vengano utilizzati correttamente.

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
'

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 -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'

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"

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 viene eseguito in un ambiente multisezione, devi impostare il flag --spmd_dcn_parallelism sul numero di sezioni.

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

Esegui la pulizia

Elimina le risorse TPU e in coda.

Risultati benchmark

La velocità effettiva per tutte e tre le dimensioni dei modelli Llama 2 è inclusa nella seguente tabella.

	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 del 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 assistenza o feedback per Cloud TPU.