Addestramento su Cloud TPU v5p
Cloud TPU v5p è la TPU di quinta generazione di Google Cloud e il successore della TPU v4. v5p è ottimizzata per l'addestramento su larga scala e per essere una piattaforma leader per lo sviluppo di modelli di base LLM, modelli di diffusione e AI generativa. A un livello generale, la versione v5p offre fino al doppio delle prestazioni della versione v4, con il doppio delle TPU in un pod (6000 sezioni più grandi rispetto a 3000 nella versione v4), con un aumento delle prestazioni fino a quattro volte a livello di pod. Inoltre, funziona a una frequenza di clock più elevata (1,75 GHz rispetto a 1,05 GHz), aggiunge SparseCore per gli embedding su larga scala e triplica la capacità della memoria HBM (High Bandwidth Memory).
Concetti di Cloud TPU v5p
Se non hai mai utilizzato Cloud TPU, consulta la home page della documentazione di 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 di sistema di 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.
Gestire le risorse TPU
Tutti i comandi in questo documento presuppongono che tu stia creando una VM TPU v5p. Per ulteriori informazioni sui comandi per creare VM TPU, consulta Gestire le TPU o la Guida utente per le risorse in coda per gestire le risorse in coda. Per semplificare l'esecuzione dei comandi, gli esempi di codice in questo documento utilizzano le seguenti variabili di ambiente:
export PROJECT_ID=your-project export ACCELERATOR_TYPE=v5p-8 export ZONE=us-east5-a export RUNTIME_VERSION=v2-alpha-tpuv5 export TPU_NAME=your-tpu-name
Descrizioni delle variabili di ambiente
PROJECT_ID- Il progetto Google Cloud in cui stai creando la 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.
ZONE- La zona in cui prevedi di creare la tua Cloud TPU.
RUNTIME_VERSION- La versione del software TPU.
TPU_NAME- Il nome definito dall'utente della TPU in uso.
Configurazione del 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 le forme dei segmenti sono superiori a 4 chip, avrai più VM
in un segmento. 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
nella frazione 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 la versione v5p e PyTorch 2.1 e versioni successive utilizzano PJRT come runtime predefinito per tutte le versioni 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 pip install numpy pip install torch torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu '
Utilizza uno script Python con PJRT per convalidare l'installazione. Lo script mostra 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 profilazione
Cloud TPU v5p supporta il monitoraggio e il profiling utilizzando gli stessi metodi delle generazioni precedenti di Cloud TPU. Puoi leggere l'articolo Profilazione del modello con gli strumenti Cloud TPU per saperne di più sulla profilazione e Monitoraggio delle VM Cloud TPU per saperne di più sul monitoraggio.
Tutorial sulla 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 la
Introduzione a Multislice.
Tutorial su JAX
Train Diffusion 2.1
Questo tutorial mostra come addestrare il modello di diffusione stabile di HuggingFace utilizzando il set di dati Pokémon su Cloud TPU v5p.
Il modello di diffusione stabile è un modello di testo a immagine latente che genera immagini fotorealistiche da qualsiasi input di testo. Per ulteriori informazioni, consulta le seguenti risorse:
Configura
Crea le variabili di ambiente:
export GCS_BUCKET_NAME=your-bucket export PROJECT_ID=your-project-ID export ACCELERATOR_TYPE=v5p-32 export ZONE=europe-west4-b export LOCATION=europe-west4 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 dei flag dei comandi
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. LOCALITÀ La regione Google Cloud in cui creare il bucket di archiviazione Cloud Storage. RUNTIME_VERSION Per la versione v5p, utilizza 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 reservedospot. Se non viene specificato nessuno di questi valori, il valore predefinito di QUOTA_TYPE èon-demand. Consulta la sezione quote per informazioni sui diversi tipi di quote supportati da Cloud TPU.VALID_UNTIL_DURATION La durata di validità della richiesta. Consulta la sezione Risorse in coda per informazioni sulle diverse durate valide. Configura un bucket di archiviazione per l'output del modello.
gcloud storage buckets create gs://$GCS_BUCKET_NAME \ --project=$PROJECT_ID \ --location=$LOCATION
-
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 connetterti tramite SSH alla VM TPU quando la risorsa in coda sarà in stato
ACTIVE. Controlla lo stato della risorsa in coda eseguendo il seguente 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
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 alla fine della sessione o rimuovere le richieste di risorse in coda nello stato "FAILED". Per eliminare una risorsa in coda, elimina i vari slice e poi la richiesta della 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 i segmenti 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 del benchmark
Lo script di addestramento di Stable Diffusion è stato eseguito su v5p-8, v5p-32 e v5p-128. La tabella seguente mostra la velocità in Mbps.
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, ad alte prestazioni, arbitrariamente scalabile e ben testato, scritto in puro Python/JAX e indirizzato alle Cloud TPU. MaxText offre a ricercatori e sviluppatori uno strumento accessibile e adattabile per far avanzare le frontiere della ricerca e dello sviluppo dell'elaborazione del linguaggio naturale (NLP).
Prima di eseguire questo tutorial, devi configurare l'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 dei flag dei comandi
Variabile Descrizione PROJECT_ID Nome del progetto Google Cloud TPU_NAME Un nome definito dall'utente per la TPU. 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 v2-alpha-tpuv5 per la versione del runtime. RUN_NAME Nome dell'esecuzione dell'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 le 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). Tieni presente che la creazione di una rete richiede l'autorizzazione
compute.networks.createnel progetto. Gli esempi riportati di seguito mostrano come creare una regola di rete e firewall dedicata.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}
Clona il repository MaxText
git clone https://github.com/google/maxtext.git
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. La posizione del log viene richiamata 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
Per eseguire lo script di addestramento più volte 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
--node-count ${NODE_COUNT} \ --node-prefix ${NODE_PREFIX} # optional, the default is QUEUED_RESOURCE_ID
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
QueuedResourcesarà nello statoACTIVE:Utilizza il comando
describeper 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
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, 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) |
|---|---|---|---|
32B |
v5p-128 |
3,28E+02 |
71,47% |
64B |
v5p-128 |
3,23E+02 |
70,31% |
128B |
v5p-256 |
3,15E+02 |
68,68% |
128B |
v5p-512 |
3,15E+02 |
68,53% |
256B |
v5p-1024 |
3.16E+02 |
68,82% |
512B |
v5p-1024 |
2,94E+02 |
63,99% |
1024B |
v5p-2048 |
2,49E+02 |
64,05% |
1024B |
v5p-4096 |
2,97E+02 |
64,80% |
1160B |
v5p-7680 |
2,95E+02 |
64,27% |
1160B |
v5p-12288 |
3,04E+02 |
66,23% |
Il modello di parametro da 256 B è stato testato su v5p-512 e v5p-1024 utilizzando sia i pesi bf16 che 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 dei FLOPS del 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 TPU v5p-8 utilizzando un set di dati falso. Se vuoi utilizzare un altro set di dati, consulta la sezione Preparazione del set di dati.
Configura
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.18.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-8comeACCELERATOR_TYPE.-
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 non appena la risorsa in coda sarà nello 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}
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
Vai alla 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
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 versioni successive utilizzando un set di dati falso. Se vuoi utilizzare un altro set di dati, consulta la sezione 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.18.0-pod-pjrt export QUEUED_RESOURCE_ID=your-queued-resource-id export QUOTA_TYPE=quota-type
ACCELERATOR_TYPEpuò esserev5p-16o superiore.-
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 non appena la risorsa in coda sarà nello stato
ACTIVE.Utilizza il comando
describeper 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 TPU (worker zero) 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
Vai alla 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
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 la parallellizzazione generale e scalabile per i grafici di calcolo ML (GSPMD).
Configurazione
Crea le variabili di ambiente.
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
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
QueuedResourcesarà nello statoACTIVE:Utilizza il comando
describeper 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
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 torch_xla[tpu]~=2.5.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 esegui l'operazione in un ambiente multislice, devi impostare il flag --spmd_dcn_parallelism sul numero di slice.
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 del 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 della griglia di sharding |
(4, 1) |
(64, 1) |
(16, 4) |
Utilizzo dei FLOPS del modello (MFU) |
56,67% |
55,80% |
51,85% |
Assistenza e feedback
Accogliamo con entusiasmo tutti i feedback. Per condividere un feedback o richiedere assistenza, compila il modulo di assistenza o feedback per Cloud TPU.