Addestramento di Cloud TPU v5e
Cloud TPU v5e è l'acceleratore di AI di ultima generazione di Google Cloud. Con un ingombro inferiore di 256 chip per pod, v5e è ottimizzato per essere il prodotto di maggior valore per l'addestramento, l'ottimizzazione e la pubblicazione di transformer, conversione di testo in immagini e reti neurali convoluzionali (CNN). Per saperne di più sull'utilizzo di Cloud TPU v5e per la pubblicazione, consulta Inferenza con v5e.
Per ulteriori informazioni sull'hardware e sulle configurazioni di Cloud TPU v5e, consulta TPU v5e.
Inizia
Le sezioni seguenti descrivono come iniziare a utilizzare TPU v5e.
Quota per le richieste
Per utilizzare TPU v5e per l'addestramento, devi disporre di una quota. Esistono diversi tipi di quote per le TPU on demand, le TPU riservate e le VM TPU Spot. Sono richieste quote separate se utilizzi TPU v5e per l'inferenza. Per ulteriori informazioni sulle quote, consulta Quote. Per richiedere una quota TPU v5e, contatta il team di Cloud Sales.
Crea un account e un progetto Google Cloud
Per utilizzare Cloud TPU, devi disporre di un account e di un progetto Google Cloud. Per maggiori informazioni, consulta Configurare un ambiente Cloud TPU.
Crea una Cloud TPU
La best practice consiste nel eseguire il provisioning di Cloud TPU v5 come risorse in coda utilizzando il comando queued-resource create
. Per saperne di più, consulta Gestire le risorse in coda.
Puoi anche utilizzare l'API Create Node (gcloud compute tpus tpu-vm create
) per eseguire il provisioning di Cloud TPU v5. Per saperne di più, consulta Gestire le risorse TPU.
Per ulteriori informazioni sulle configurazioni v5e disponibili per l'addestramento, consulta Tipi di Cloud TPU v5e per l'addestramento.
Configurazione del framework
Questa sezione descrive la procedura di configurazione generale per l'addestramento di modelli personalizzati utilizzando JAX o PyTorch con TPU v5e. Il supporto di TensorFlow è disponibile nelle versioni di runtime TPU tpu-vm-tf-2.18.0-pjrt
e tpu-vm-tf-2.18.0-pod-pjrt
.
Per le istruzioni di configurazione dell'inferenza, consulta l'introduzione all'inferenza v5e.
Configurazione per JAX
Se le forme dei segmenti sono superiori a 8 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 in un unico passaggio senza utilizzare SSH per accedere a ciascuna singola VM:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
Descrizioni dei flag dei comandi
Variabile | Descrizione |
TPU_NAME | L'ID testo assegnato dall'utente della TPU che viene creato quando viene allocata la richiesta di risorse in coda. |
PROJECT_ID | Nome del progetto Google Cloud. Utilizza un progetto esistente o creane uno nuovo in Configurare il progetto Google Cloud |
ZONA | Consulta il documento Regioni e zone TPU per le zone supportate. |
worker | La VM TPU che ha accesso alle TPU sottostanti. |
Puoi eseguire il seguente comando per controllare il numero di dispositivi (gli output mostrati qui sono stati prodotti con una fetta v5litepod-16). Questo codice verifica che tutto sia installato correttamente controllando che JAX veda i TensorCore di Cloud TPU e possa eseguire operazioni di base:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'
L'output sarà simile al seguente:
SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
16
4
16
4
16
4
16
4
jax.device_count()
mostra il numero totale di chip nell'intervallo specificato.
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.]
Prova i tutorial JAX in questo documento per iniziare con l'addestramento v5e utilizzando JAX.
Configurazione per PyTorch
Tieni presente che la versione v5e supporta solo il runtime PJRT e che PyTorch 2.1 e versioni successive utilizzeranno PJRT come runtime predefinito per tutte le versioni di TPU.
Questa sezione descrive come iniziare a utilizzare PJRT su v5e con PyTorch/XLA con comandi 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 -y
sudo apt-get install libomp5 -y
pip3 install mkl mkl-include
pip3 install tf-nightly tb-nightly tbp-nightly
pip3 install numpy
sudo apt-get install libopenblas-dev -y
pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html
pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
Se ricevi un errore durante l'installazione dei wheel per torch
, torch_xla
o
torchvision
come
pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end
or semicolon (after name and no valid version specifier) torch==nightly+20230222
,
esegui il downgrade della versione con questo comando:
pip3 install setuptools==62.1.0
Eseguire uno script con PJRT
unset LD_PRELOAD
Di seguito è riportato un esempio che utilizza uno script Python per eseguire un calcolo su una VM v5e:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker all \
--command='
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$HOME/.local/lib/
export PJRT_DEVICE=TPU_C_API
export PT_XLA_DEBUG=0
export USE_TORCH=ON
unset LD_PRELOAD
export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
python3 -c "import torch; import torch_xla; import torch_xla.core.xla_model as xm; print(xm.xla_device()); dev = xm.xla_device(); t1 = torch.randn(3,3,device=dev); t2 = torch.randn(3,3,device=dev); print(t1 + t2)"'
Viene generato un output simile al seguente:
SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
xla:0
tensor([[ 1.8611, -0.3114, -2.4208],
[-1.0731, 0.3422, 3.1445],
[ 0.5743, 0.2379, 1.1105]], device='xla:0')
Prova i tutorial di PyTorch in questo documento per iniziare con l'addestramento di v5e utilizzando PyTorch.
Elimina la TPU e la risorsa in coda al termine della sessione. Per eliminare una risorsa in coda, elimina il segmento e poi la risorsa in coda in due passaggi:
gcloud compute tpus tpu-vm delete ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--quiet
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--quiet
Questi due passaggi possono essere utilizzati anche per rimuovere le richieste di risorse in coda nello stato FAILED
.
Esempi di JAX/FLAX
Le sezioni seguenti descrivono esempi di come addestrare i modelli JAX e FLAX su TPU v5e.
Addestra ImageNet su v5e
Questo tutorial descrive come addestrare ImageNet su v5e utilizzando dati di input falsi. Se vuoi utilizzare dati reali, consulta il file README su GitHub.
Configura
Crea le variabili di ambiente:
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-west4-a export RUNTIME_VERSION=v2-alpha-tpuv5-lite 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_UNTIL_DURATION=1d
-
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 quando la risorsa in coda sarà nello stato
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Quando QueuedResource è nello stato
ACTIVE
, l'output sarà simile al seguente:state: ACTIVE
Installa la versione più recente di JAX e jaxlib:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
Clona il modello ImageNet e installa i requisiti corrispondenti:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='git clone https://github.com/google/flax.git && cd flax/examples/imagenet && pip install -r requirements.txt && pip install flax==0.7.4'
Per generare dati falsi, il modello ha bisogno di informazioni sulle dimensioni del set di dati. Questi dati possono essere raccolti dai metadati del set di dati ImageNet:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='mkdir -p $HOME/flax/.tfds/metadata/imagenet2012/5.1.0 && curl https://raw.githubusercontent.com/tensorflow/datasets/v4.4.0/tensorflow_datasets/testing/metadata/imagenet2012/5.1.0/dataset_info.json --output $HOME/flax/.tfds/metadata/imagenet2012/5.1.0/dataset_info.json'
Addestra il modello
Una volta completati tutti i passaggi precedenti, puoi addestrare il modello.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='cd flax/examples/imagenet && JAX_PLATFORMS=tpu python3 imagenet_fake_data_benchmark.py'
Elimina la TPU e la risorsa in coda
Elimina la TPU e la risorsa in coda al termine della sessione.
gcloud compute tpus tpu-vm delete ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--quiet
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--quiet
Modelli FLAX Hugging Face
I modelli Hugging Face implementati in FLAX funzionano subito su Cloud TPU v5e. Questa sezione fornisce istruzioni per eseguire i modelli più diffusi.
Addestramento di ViT su Imagenette
Questo tutorial mostra come addestrare il modello Vision Transformer (ViT) di HuggingFace utilizzando il set di dati Imagenette di Fast AI su Cloud TPU v5e.
Il modello ViT è stato il primo ad addestrare correttamente un codificatore Transformer su ImageNet con risultati eccellenti rispetto alle reti convoluzionali. Per maggiori informazioni, consulta le seguenti risorse:
Configura
Crea le variabili di ambiente:
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-west4-a export RUNTIME_VERSION=v2-alpha-tpuv5-lite 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_UNTIL_DURATION=1d
-
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 quando la risorsa in coda sarà nello stato
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Quando la risorsa in coda è nello stato
ACTIVE
, l'output sarà simile al seguente:state: ACTIVE
Installa JAX e la relativa libreria:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
Scarica il repository di Hugging Face 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/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install tensorflow==2.18.0 && pip install -r examples/flax/vision/requirements.txt'
Scarica il set di dati Imagenette:
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='cd transformers && wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'
Addestra il modello
Addestra il modello con un buffer premappato di 4 GB.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='cd transformers && JAX_PLATFORMS=tpu python3 examples/flax/vision/run_image_classification.py --train_dir "imagenette2/train" --validation_dir "imagenette2/val" --output_dir "./vit-imagenette" --learning_rate 1e-3 --preprocessing_num_workers 32 --per_device_train_batch_size 8 --per_device_eval_batch_size 8 --model_name_or_path google/vit-base-patch16-224-in21k --num_train_epochs 3'
Elimina la TPU e la risorsa in coda
Elimina la TPU e la risorsa in coda al termine della sessione.
gcloud compute tpus tpu-vm delete ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--quiet
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--quiet
Risultati del benchmarking di ViT
Lo script di addestramento è stato eseguito su v5litepod-4, v5litepod-16 e v5litepod-64. La tabella seguente mostra le portate con diversi tipi di acceleratori.
Tipo di acceleratore | v5litepod-4 | v5litepod-16 | v5litepod-64 |
Epoca | 3 | 3 | 3 |
Dimensione batch globale | 32 | 128 | 512 |
Velocità effettiva (esempi/sec) | 263,40 | 429,34 | 470,71 |
Allenare la diffusione su Pokémon
Questo tutorial mostra come addestrare il modello di diffusione stabile di HuggingFace utilizzando il set di dati Pokémon su Cloud TPU v5e.
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
Configura un bucket di archiviazione per l'output del modello.
gcloud storage buckets create gs://your_bucket
--project=your_project
--location=us-west1
export GCS_BUCKET_NAME=your_bucketCreare variabili di ambiente
export GCS_BUCKET_NAME=your_bucket export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-west1-c export RUNTIME_VERSION=v2-alpha-tpuv5-lite 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 dei flag dei comandi
Variabile Descrizione GCS_BUCKET_NAME Visualizzato nella console Google Cloud -> Cloud Storage -> Bucket PROJECT_ID Nome del progetto Google Cloud. Utilizza un progetto esistente o creane uno nuovo in Configurare il 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 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 reserved
ospot
. 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. -
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 quando la risorsa in coda sarà nello stato
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Quando la risorsa in coda è nello stato
ACTIVE
, l'output sarà simile al seguente:state: ACTIVE
Installa JAX e la relativa libreria.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='pip install "jax[tpu]==0.4.16" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
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/RissyRan/diffusers.git && cd diffusers && pip install . && pip install tensorflow==2.18.0 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'
Addestra il modello
Addestra il modello con un buffer premappato di 4 GB.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE} --project ${PROJECT_ID} --worker=all --command="
git clone https://github.com/google/maxdiffusion
cd maxdiffusion
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
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=your_run base_output_directory=gs://${GCS_BUCKET_NAME}/ enable_profiler=False"
Elimina la TPU e la risorsa in coda
Elimina la TPU e la risorsa in coda al termine della sessione.
gcloud compute tpus tpu-vm delete ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--quiet
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--quiet
Risultati del benchmarking per la diffusione
Lo script di addestramento è stato eseguito su v5litepod-4, v5litepod-16 e v5litepod-64. La tabella seguente mostra le portate.
Tipo di acceleratore | v5litepod-4 | v5litepod-16 | v5litepod-64 |
Passaggio di addestramento | 1500 | 1500 | 1500 |
Dimensione del batch globale | 32 | 64 | 128 |
Velocità effettiva (esempi/sec) | 36,53 | 43,71 | 49,36 |
Addestrare GPT2 sul set di dati OSCAR
Questo tutorial mostra come addestrare il modello GPT2 di HuggingFace utilizzando il set di dati OSCAR su Cloud TPU v5e.
GPT2 è un modello di transformer preaddestrato su testi non elaborati senza etichettatura da parte di persone. È stato addestrato a prevedere la parola successiva nelle frasi. Per maggiori informazioni, consulta le seguenti risorse:
Configura
Crea le variabili di ambiente:
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-west4-a export RUNTIME_VERSION=v2-alpha-tpuv5-lite 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
-
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 quando la risorsa in coda sarà nello stato
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Quando la risorsa in coda è nello stato
ACTIVE
, l'output sarà simile al seguente:state: ACTIVE
Installa JAX e la relativa libreria.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='pip install -U "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html'
Scarica i requisiti del repository HuggingFace e installali.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='git clone https://github.com/huggingface/transformers.git && cd transformers && pip install . && pip install -r examples/flax/_tests_requirements.txt && pip install --upgrade huggingface-hub urllib3 zipp && pip install TensorFlow && pip install -r examples/flax/language-modeling/requirements.txt'
Scarica le configurazioni per addestrare il modello.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='cd transformers/examples/flax/language-modeling && gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt . --recursive'
Addestra il modello
Addestra il modello con un buffer premappato di 4 GB.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='cd transformers/examples/flax/language-modeling && TPU_PREMAPPED_BUFFER_SIZE=4294967296 JAX_PLATFORMS=tpu python3 run_clm_flax.py --output_dir=./gpt --model_type=gpt2 --config_name=./gpt --tokenizer_name=./gpt --dataset_name=oscar --dataset_config_name=unshuffled_deduplicated_no --do_train --do_eval --block_size=512 --per_device_train_batch_size=4 --per_device_eval_batch_size=4 --learning_rate=5e-3 --warmup_steps=1000 --adam_beta1=0.9 --adam_beta2=0.98 --weight_decay=0.01 --overwrite_output_dir --num_train_epochs=3 --logging_steps=500 --eval_steps=2500'
Elimina la TPU e la risorsa in coda
Elimina la TPU e la risorsa in coda al termine della sessione.
gcloud compute tpus tpu-vm delete ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--quiet
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--quiet
Risultati del benchmarking per GPT-2
Lo script di addestramento è stato eseguito su v5litepod-4, v5litepod-16 e v5litepod-64. La tabella seguente mostra le portate.
v5litepod-4 | v5litepod-16 | v5litepod-64 | |
Epoca | 3 | 3 | 3 |
Dimensione batch globale | 64 | 64 | 64 |
Velocità effettiva (esempi/sec) | 74,60 | 72,97 | 72,62 |
PyTorch/XLA
Le sezioni seguenti descrivono esempi di come addestrare i modelli PyTorch/XLA su TPU v5e.
Addestra ResNet utilizzando il runtime PJRT
PyTorch/XLA esegue la migrazione da XRT a PjRt da PyTorch 2.0 e versioni successive. Di seguito sono riportate le istruzioni aggiornate per configurare la versione 5e per i carichi di lavoro di addestramento PyTorch/XLA.
Configura
Crea le variabili di ambiente:
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-west4-a export RUNTIME_VERSION=v2-alpha-tpuv5-lite export SERVICE_ACCOUNT=your_service_account export TPU_NAME=tpu-name export QUEUED_RESOURCE_ID=queued_resource_id export QUOTA_TYPE=quota_type export VALID_UNTIL_DURATION=1d
-
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 quando la risorsa in coda sarà nello stato
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Quando la risorsa in coda è nello stato
ACTIVE
, l'output sarà simile al seguente:state: ACTIVE
Installa le dipendenze specifiche di Torch/XLA
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=' sudo apt-get update -y sudo apt-get install libomp5 -y pip3 install mkl mkl-include pip3 install tf-nightly tb-nightly tbp-nightly pip3 install numpy sudo apt-get install libopenblas-dev -y pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
Addestrare il modello ResNet
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='
date
export PJRT_DEVICE=TPU_C_API
export PT_XLA_DEBUG=0
export USE_TORCH=ON
export XLA_USE_BF16=1
export LIBTPU_INIT_ARGS=--xla_jf_auto_cross_replica_sharding
export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
git clone https://github.com/pytorch/xla.git
cd xla/
git reset --hard caf5168785c081cd7eb60b49fe4fffeb894c39d9
python3 test/test_train_mp_imagenet.py --model=resnet50 --fake_data --num_epochs=1 —num_workers=16 --log_steps=300 --batch_size=64 --profile'
Elimina la TPU e la risorsa in coda
Elimina la TPU e la risorsa in coda al termine della sessione.
gcloud compute tpus tpu-vm delete ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--quiet
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--quiet
Risultato del benchmark
La tabella seguente mostra le portate di benchmark.
Tipo di acceleratore | Velocità effettiva (esempi/secondo) |
v5litepod-4 | 4240 ex/s |
v5litepod-16 | 10.810 ex/s |
v5litepod-64 | 46.154 ex/s |
Addestra GPT2 su v5e
Questo tutorial illustra come eseguire GPT2 su v5e utilizzando il repository HuggingFace su PyTorch/XLA utilizzando il set di dati wikitext.
Configura
Crea le variabili di ambiente:
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-west4-a export RUNTIME_VERSION=v2-alpha-tpuv5-lite 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
-
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 quando la risorsa in coda è nello stato
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Quando la risorsa in coda è nello stato
ACTIVE
, l'output sarà simile al seguente:state: ACTIVE
Installa le dipendenze PyTorch/XLA.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=' sudo apt-get -y update sudo apt install -y libopenblas-base pip3 install torchvision pip3 uninstall -y torch pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
Scarica i requisiti del repository HuggingFace e installali.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=' git clone https://github.com/pytorch/xla.git pip install --upgrade accelerate git clone https://github.com/huggingface/transformers.git cd transformers git checkout ebdb185befaa821304d461ed6aa20a17e4dc3aa2 pip install . git log -1 pip install datasets evaluate scikit-learn '
Scarica le configurazioni del modello preaddestrato.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=' gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/ --recursive gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'
Addestra il modello
Addestra il modello 2B utilizzando una dimensione del batch di 16.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='
export PJRT_DEVICE=TPU_C_API
cd transformers/
export LD_LIBRARY_PATH=/usr/local/lib/
export PT_XLA_DEBUG=0
export USE_TORCH=ON
python3 examples/pytorch/xla_spawn.py \
--num_cores=4 \
examples/pytorch/language-modeling/run_clm.py \
--num_train_epochs=3 \
--dataset_name=wikitext \
--dataset_config_name=wikitext-2-raw-v1 \
--per_device_train_batch_size=16 \
--per_device_eval_batch_size=16 \
--do_train \
--do_eval \
--logging_dir=./tensorboard-metrics \
--cache_dir=./cache_dir \
--output_dir=/tmp/test-clm \
--overwrite_output_dir \
--cache_dir=/tmp \
--config_name=examples/pytorch/language-modeling/my_config_2.json \
--tokenizer_name=gpt2 \
--block_size=1024 \
--optim=adafactor \
--adafactor=true \
--save_strategy=no \
--logging_strategy=no \
--fsdp=full_shard \
--fsdp_config=examples/pytorch/language-modeling/fsdp_config.json'
Elimina la TPU e la risorsa in coda
Elimina la TPU e la risorsa in coda al termine della sessione.
gcloud compute tpus tpu-vm delete ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--quiet
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--quiet
Risultato del benchmark
Lo script di addestramento è stato eseguito su v5litepod-4, v5litepod-16 e v5litepod-64. La tabella seguente mostra i throughput del benchmark per diversi tipi di acceleratori.
v5litepod-4 | v5litepod-16 | v5litepod-64 | |
Epoca | 3 | 3 | 3 |
config | 600 MLN | 2 Mld | 16B |
Dimensione del batch globale | 64 | 128 | 256 |
Velocità effettiva (esempi/sec) | 66 | 77 | 31 |
Addestrare ViT su v5e
Questo tutorial spiega come eseguire VIT su v5e utilizzando il repository HuggingFace su PyTorch/XLA sul set di dati cifar10.
Configura
Crea le variabili di ambiente:
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-west4-a export RUNTIME_VERSION=v2-alpha-tpuv5-lite export SERVICE_ACCOUNT=your_service_account export TPU_NAME=tpu-name export QUEUED_RESOURCE_ID=queued_resource_id export QUOTA_TYPE=quota_type export VALID_UNTIL_DURATION=1d
-
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 QueuedResource sarà nello stato
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Quando la risorsa in coda è nello stato
ACTIVE
, l'output sarà simile al seguente:state: ACTIVE
Installa le dipendenze PyTorch/XLA
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all --command=' sudo apt-get update -y sudo apt-get install libomp5 -y pip3 install mkl mkl-include pip3 install tf-nightly tb-nightly tbp-nightly pip3 install numpy sudo apt-get install libopenblas-dev -y pip3 install torch~=2.1.0 torchvision torch_xla[tpu]~=2.1.0 -f https://storage.googleapis.com/libtpu-releases/index.html pip3 install torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html'
Scarica i requisiti del repository HuggingFace e installali.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=" git clone https://github.com/suexu1025/transformers.git vittransformers; \ cd vittransformers; \ pip3 install .; \ pip3 install datasets; \ wget https://github.com/pytorch/xla/blob/master/scripts/capture_profile.py"
Addestra il modello
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='
export PJRT_DEVICE=TPU_C_API
export PT_XLA_DEBUG=0
export USE_TORCH=ON
export TF_CPP_MIN_LOG_LEVEL=0
export XLA_USE_BF16=1
export LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH
export TPU_LIBRARY_PATH=$HOME/.local/lib/python3.10/site-packages/libtpu/libtpu.so
cd vittransformers
python3 -u examples/pytorch/xla_spawn.py --num_cores 4 examples/pytorch/image-pretraining/run_mae.py --dataset_name=cifar10 \
--remove_unused_columns=False \
--label_names=pixel_values \
--mask_ratio=0.75 \
--norm_pix_loss=True \
--do_train=true \
--do_eval=true \
--base_learning_rate=1.5e-4 \
--lr_scheduler_type=cosine \
--weight_decay=0.05 \
--num_train_epochs=3 \
--warmup_ratio=0.05 \
--per_device_train_batch_size=8 \
--per_device_eval_batch_size=8 \
--logging_strategy=steps \
--logging_steps=30 \
--evaluation_strategy=epoch \
--save_strategy=epoch \
--load_best_model_at_end=True \
--save_total_limit=3 \
--seed=1337 \
--output_dir=MAE \
--overwrite_output_dir=true \
--logging_dir=./tensorboard-metrics \
--tpu_metrics_debug=true'
Elimina la TPU e la risorsa in coda
Elimina la TPU e la risorsa in coda al termine della sessione.
gcloud compute tpus tpu-vm delete ${TPU_NAME}
--project=${PROJECT_ID}
--zone=${ZONE}
--quiet
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
--project=${PROJECT_ID}
--zone=${ZONE}
--quiet
Risultato del benchmark
La tabella seguente mostra i throughput del benchmark per i diversi tipi di acceleratori.
v5litepod-4 | v5litepod-16 | v5litepod-64 | |
Epoca | 3 | 3 | 3 |
Dimensione batch globale | 32 | 128 | 512 |
Velocità effettiva (esempi/sec) | 201 | 657 | 2844 |
TensorFlow 2.x
Le sezioni seguenti descrivono esempi di come addestrare i modelli TensorFlow 2.x su TPU v5e.
Addestramento di ResNet su un singolo host v5e
Questo tutorial descrive come addestrare ImageNet su v5litepod-4
o v5litepod-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=v5litepod-4 export ZONE=us-east1-c export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pjrt export TPU_NAME=your-tpu-name export QUEUED_RESOURCE_ID=your-queued-resource-id export QUOTA_TYPE=quota-type
ACCELERATOR_TYPE
può esserev5litepod-4
ov5litepod-8
.-
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 tramite SSH alla VM TPU quando 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"
Elimina la TPU e la risorsa in coda
Eliminare la TPU
gcloud compute tpus tpu-vm delete ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet
Eliminare la richiesta di risorsa in coda
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet
Addestramento di ResNet su un modello v5e multi-host
Questo tutorial descrive come addestrare ImageNet su v5litepod-16
o versioni successive utilizzando un set di dati falso. Se vuoi utilizzare un altro set di dati, consulta la sezione Preparare il set di dati.
Crea le variabili di ambiente:
export PROJECT_ID=your_project_ID export ACCELERATOR_TYPE=v5litepod-16 export ZONE=us-east1-c export RUNTIME_VERSION=tpu-vm-tf-2.15.0-pod-pjrt export TPU_NAME=your_tpu_name export QUEUED_RESOURCE_ID=your-queued-resource-id export QUOTA_TYPE=quota-type
ACCELERATOR_TYPE
può esserev5litepod-16
o 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 tramite SSH alla VM TPU quando 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 (worker zero) 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 TPU_LOAD_LIBRARY=0 export TPU_NAME=your_tpu_name
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 \
--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*"
Elimina la TPU e la risorsa in coda
Eliminare la TPU
gcloud compute tpus tpu-vm delete ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet
Eliminare la richiesta di risorsa in coda
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet