Entrenamiento de Cloud TPU v5e
Cloud TPU v5e es el acelerador de IA de última generación de Google Cloud. Con una huella más pequeña de 256 chip por Pod, la v5e está optimizada para ser el producto de mayor valor en entrenamiento, ajuste y entrega de transformadores, texto a imagen y red neuronal convolucional (CNN). Si deseas obtener más información sobre el uso de Cloud TPU v5e para la entrega, consulta Inferencia con v5e.
Para obtener más información sobre el hardware y las configuraciones de la TPU de Cloud TPU v5e, consulta TPU v5e.
Comenzar
En las siguientes secciones, se describe cómo comenzar a usar TPU v5e.
Cuota de solicitudes
Necesitas cuota si quieres usar la TPU v5e para el entrenamiento. Existen diferentes tipos de cuotas para las TPU según demanda, las TPU reservadas y las VM Spot de TPU. Se requieren cuotas independientes si usas la TPU v5e para la inferencia. Para obtener más información sobre las cuotas, consulta Cuotas. Para solicitar la cuota de TPU v5e, comunícate con Ventas de Cloud.
Crea una cuenta y un proyecto de Google Cloud
Necesitas una cuenta y un proyecto de Google Cloud para usar Cloud TPU. Para obtener más información, consulta Configura un entorno de Cloud TPU.
Crear una Cloud TPU
La práctica recomendada es aprovisionar Cloud TPU v5 como recursos en cola mediante el comando queued-resource create
. Para obtener más información, consulta Administra recursos en cola.
También puedes usar la API de Crear nodo (gcloud alpha compute tpus tpu-vm create
) para aprovisionar Cloud TPU v5es. Para obtener más información, consulta Administra recursos de TPU.
Si deseas obtener más información sobre las configuraciones de v5e disponibles para el entrenamiento, consulta Tipos de Cloud TPU v5e para el entrenamiento.
Configuración del framework
En esta sección, se describe el proceso de configuración general para el entrenamiento de modelos personalizados mediante JAX o PyTorch con TPU v5e. La compatibilidad con TensorFlow está disponible en las versiones del entorno de ejecución de TPU tpu-vm-tf-2.16.1-pjrt
y tpu-vm-tf-2.16.1-pod-pjrt
.
Para obtener instrucciones de configuración de la inferencia, consulta Introducción a la inferencia v5e.
Configuración para JAX
Si tienes formas de porciones de más de 8 chips, tendrás varias VMs en una porción. En este caso, debes usar la marca --worker=all
para ejecutar la instalación en todas las VM de TPU en un solo paso, sin usar SSH para acceder a cada una por separado:
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'
Puedes ejecutar el siguiente comando para verificar la cantidad de dispositivos (los resultados que se muestran aquí se produjeron con un segmento v5litepod-16). Este código prueba que todo esté instalado correctamente verificando que JAX detecte los tensores de Cloud TPU y pueda ejecutar operaciones básicas:
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())"'
El resultado será similar al siguiente ejemplo:
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()
muestra la cantidad total de chips en la porción determinada.
jax.local_device_count()
indica el recuento de chips a los que puede acceder una sola VM en esta porción.
# 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))"'
El resultado será similar al siguiente ejemplo:
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.]
Prueba los instructivos de JAX en este documento para comenzar con el entrenamiento de la v5e con JAX.
Configuración para PyTorch
Ten en cuenta que v5e solo es compatible con el entorno de ejecución PJRT, y PyTorch 2.1 y versiones posteriores usará PJRT como el entorno de ejecución predeterminado para todas las versiones de TPU.
En esta sección, se describe cómo comenzar a usar PJRT en v5e con PyTorch/XLA con comandos para todos los trabajadores.
Instala dependencias
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'
Si recibes un error cuando instalas las ruedas para torch
, torch_xla
o torchvision
, como pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end
or semicolon (after name and no valid version specifier) torch==nightly+20230222
, regresa tu versión con este comando:
pip3 install setuptools==62.1.0
Ejecuta una secuencia de comandos con PJRT
unset LD_PRELOAD
El siguiente es un ejemplo en el que se usa una secuencia de comandos de Python para realizar un cálculo en 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)"'
Esto genera un resultado similar al que se muestra a continuación:
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')
Prueba los instructivos de PyTorch de este documento para comenzar con el entrenamiento de la versión 5e mediante PyTorch.
Borra la TPU y el recurso en cola al final de la sesión. Para borrar un recurso en cola, borra la porción y, luego, el recurso en cola en 2 pasos:
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
Estos dos pasos también se pueden usar para quitar las solicitudes de recursos en cola que tengan el estado FAILED
.
Ejemplos de JAX/FLAX
En las siguientes secciones, se describen ejemplos de cómo entrenar modelos JAX y FLAX en la TPU v5e.
Entrena ImageNet en v5e
En este instructivo, se describe cómo entrenar ImageNet en v5e con datos de entrada falsos. Si quieres usar datos reales, consulta el archivo README en GitHub.
Configurar
Crea variables de entorno:
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 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}
Podrás establecer una conexión SSH a tu VM de TPU una vez que tu recurso en cola esté en el estado
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Cuando QueuedResource se encuentra en el estado
ACTIVE
, el resultado será similar al siguiente:state: ACTIVE
Instala las versiones más recientes de JAX y jaxlib:
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'
Clona el modelo de ImageNet y, luego, instala los requisitos correspondientes:
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'
Para generar datos falsos, el modelo necesita información sobre las dimensiones del conjunto de datos. Esto se puede recopilar a partir de los metadatos del conjunto de datos de 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'
Entrenar el modelo
Una vez que hayas completado todos los pasos anteriores, puedes entrenar el modelo.
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'
Borra la TPU y el recurso en cola
Borra la TPU y el recurso en cola al final de la sesión.
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
Modelos FLAX de Hugging Face
Los modelos de Hugging Face implementados en FLAX funcionan de inmediato en Cloud TPU v5e. En esta sección, se proporcionan instrucciones para ejecutar modelos populares.
Entrena ViT en Imagenette
En este instructivo, se muestra cómo entrenar el modelo de Vision Transformer (ViT) de HuggingFace con el conjunto de datos Imagenette de Fast AI en Cloud TPU v5e.
El modelo ViT fue el primero en entrenar con éxito un codificador de transformadores en ImageNet con excelentes resultados en comparación con las redes convolucionales. Para obtener más información, consulta los siguientes recursos:
Configurar
Crea variables de entorno:
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 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}
Podrás establecer una conexión SSH a tu VM de TPU una vez que el recurso en cola esté en el estado
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Cuando el recurso en cola está en el estado
ACTIVE
, el resultado será similar al siguiente:state: ACTIVE
Instala JAX y su biblioteca:
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'
Descarga el repositorio de Hugging Face y, luego, instala los requisitos:
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.16.1 && pip install -r examples/flax/vision/requirements.txt'
Descarga el conjunto de datos de 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'
Entrenar el modelo
Entrena el modelo con un búfer preasignado de 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'
Borra la TPU y el recurso en cola
Borra tu TPU y el recurso en cola al final de tu sesión.
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
Resultados de las comparativas de ViT
La secuencia de comandos de entrenamiento se ejecutó en v5litepod-4, v5litepod-16 y v5litepod-64. En la siguiente tabla, se muestran las capacidades de procesamiento con diferentes tipos de aceleradores.
Tipo de acelerador | v5litepod-4 | v5litepod-16 | v5litepod-64 |
Ciclo de entrenamiento | 3 | 3 | 3 |
Tamaño del lote global | 32 | 128 | 512 |
Capacidad de procesamiento (ejemplos/s) | 263,40 | 429.34 | 470,71 |
Entrena la difusión en Pokémon
En este instructivo, se muestra cómo entrenar el modelo de dispersión estable desde HuggingFace con el conjunto de datos de Pokémon en Cloud TPU v5e.
El modelo de difusión estable es un modelo de texto a imagen latente que genera imágenes fotorrealistas a partir de cualquier entrada de texto. Para obtener más información, consulta los siguientes recursos:
Configurar
Crea variables de entorno:
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 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}
Podrás establecer una conexión SSH a tu VM de TPU una vez que tu recurso en cola esté en el estado
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Cuando el recurso en cola está en el estado
ACTIVE
, el resultado será similar al siguiente:state: ACTIVE
Instala JAX y su biblioteca.
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'
Descarga el repositorio de HuggingFace y los requisitos de instalación.
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.16.1 clu && pip install -U -r examples/text_to_image/requirements_flax.txt'
Entrenar el modelo
Entrena el modelo con un búfer preasignado de 4 GB.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project=${PROJECT_ID} \
--zone=${ZONE} \
--worker=all \
--command='cd diffusers/examples/text_to_image && JAX_PLATFORMS=tpu,cpu python3 train_text_to_image_flax.py --pretrained_model_name_or_path=duongna/stable-diffusion-v1-4-flax --dataset_name=lambdalabs/pokemon-blip-captions --resolution=128 --center_crop --random_flip --train_batch_size=4 --mixed_precision=fp16 --max_train_steps=1500 --learning_rate=1e-05 --max_grad_norm=1 --output_dir=sd-pokemon-model'
Borra la TPU y el recurso en cola
Borra la TPU y el recurso en cola al final de la sesión.
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
Resultados de comparativas para la difusión
La secuencia de comandos de entrenamiento se ejecutó en v5litepod-4, v5litepod-16 y v5litepod-64. En la siguiente tabla, se muestran las capacidades de procesamiento.
Tipo de acelerador | v5litepod-4 | v5litepod-16 | v5litepod-64 |
Paso del entrenamiento | 1,500 | 1,500 | 1,500 |
Tamaño del lote global | 32 | 64 | 128 |
Capacidad de procesamiento (ejemplos/s) | 36,53 | 43,71 | 49:36 |
Entrena GPT2 en el conjunto de datos de OSCAR
En este instructivo, se muestra cómo entrenar el modelo GPT2 de HuggingFace usando el conjunto de datos OSCAR en Cloud TPU v5e.
GPT2 es un modelo Transformer previamente entrenado con texto sin procesar sin etiquetado manual. Se entrenó para predecir la siguiente palabra en oraciones. Para obtener más información, consulta los siguientes recursos:
Configurar
Crea variables de entorno:
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 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}
Podrás establecer una conexión SSH a tu VM de TPU una vez que el recurso en cola esté en el estado
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Cuando el recurso en cola está en el estado
ACTIVE
, el resultado será similar al siguiente:state: ACTIVE
Instala JAX y su biblioteca.
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'
Descarga el repositorio de HuggingFace y los requisitos de instalación.
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'
Descarga los parámetros de configuración para entrenar el modelo.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command='cd transformers/examples/flax/language-modeling && gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt .'
Entrenar el modelo
Entrena el modelo con un búfer preasignado de 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'
Borra la TPU y el recurso en cola
Borra la TPU y el recurso en cola al final de la sesión.
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
Resultados de comparativas de GPT2
La secuencia de comandos de entrenamiento se ejecutó en v5litepod-4, v5litepod-16 y v5litepod-64. En la siguiente tabla, se muestran las capacidades de procesamiento.
v5litepod-4 | v5litepod-16 | v5litepod-64 | |
Ciclo de entrenamiento | 3 | 3 | 3 |
Tamaño del lote global | 64 | 64 | 64 |
Capacidad de procesamiento (ejemplos/s) | 74,60 | USD 72.97 | 72,62 |
PyTorch/XLA
En las siguientes secciones, se describen ejemplos de cómo entrenar modelos de PyTorch/XLA en la TPU v5e.
Entrena ResNet con el entorno de ejecución de PJRT
PyTorch/XLA está migrando de XRT a PjRt desde PyTorch 2.0 y versiones posteriores. Estas son las instrucciones actualizadas a fin de configurar v5e para las cargas de trabajo de entrenamiento de PyTorch/XLA.
Configurar
Crea variables de entorno:
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 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}
Podrás establecer una conexión SSH a tu VM de TPU una vez que tu QueuedResource esté en estado
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Cuando el recurso en cola está en el estado
ACTIVE
, el resultado será similar al siguiente:state: ACTIVE
Instala las dependencias específicas de 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'
Entrena el modelo 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'
Borra la TPU y el recurso en cola
Borra la TPU y el recurso en cola al final de la sesión.
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
Resultado de la comparativa
En la siguiente tabla, se muestran las capacidades de procesamiento de las comparativas.
Tipo de acelerador | Capacidad de procesamiento (ejemplos por segundo) |
v5litepod-4 | 4,240 ex/s |
v5litepod-16 | 10,810 ex/s |
v5litepod-64 | 46,154 ex/s |
Entrenamiento de GPT2 en la versión 5e
En este instructivo, se explica cómo ejecutar GPT2 en v5e con el repositorio de HuggingFace en PyTorch/XLA mediante el conjunto de datos wikitext.
Configurar
Crea variables de entorno:
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 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}
Podrás establecer una conexión SSH a tu VM de TPU una vez que tu QueuedResource esté en el estado
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Cuando el recurso en cola esté en el estado
ACTIVE
, el resultado será similar al siguiente:state: ACTIVE
Instala las dependencias de 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'
Descarga el repositorio de HuggingFace y los requisitos de instalación.
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 '
Descarga los parámetros de configuración del modelo previamente entrenado.
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --worker=all \ --command=' gsutil cp -r gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/my_config_2.json transformers/examples/pytorch/language-modeling/ gsutil cp gs://cloud-tpu-tpuvm-artifacts/config/xl-ml-test/pytorch/gpt2/fsdp_config.json transformers/examples/pytorch/language-modeling/'
Entrenar el modelo
Entrena el modelo 2B con un tamaño de lote de 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'
Borra la TPU y el recurso en cola
Borra la TPU y el recurso en cola al final de la sesión.
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
Resultado de la comparativa
La secuencia de comandos de entrenamiento se ejecutó en v5litepod-4, v5litepod-16 y v5litepod-64. En la siguiente tabla, se muestran las capacidades de procesamiento de las comparativas para diferentes tipos de aceleradores.
v5litepod-4 | v5litepod-16 | v5litepod-64 | |
Ciclo de entrenamiento | 3 | 3 | 3 |
config | 600 millones | 2,000 MILLONES | 16,000 millones |
Tamaño del lote global | 64 | 128 | 256 |
Capacidad de procesamiento (ejemplos/s) | 66 | 77 | 31 |
Entrena ViT en v5e
En este instructivo, se explica cómo ejecutar VIT en v5e con el repositorio de HuggingFace en PyTorch/XLA en el conjunto de datos cifar10.
Configurar
Crea variables de entorno:
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 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}
Podrás establecer una conexión SSH a tu VM de TPU una vez que tu QueuedResource esté en el estado
ACTIVE
:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Cuando el recurso en cola esté en el estado
ACTIVE
, el resultado será similar al siguiente:state: ACTIVE
Instala dependencias de 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'
Descarga el repositorio de HuggingFace y los requisitos de instalación.
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"
Entrenar el modelo
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'
Borra la TPU y el recurso en cola
Borra la TPU y el recurso en cola al final de la sesión.
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
Resultado de la comparativa
En la siguiente tabla, se muestran las capacidades de procesamiento de las comparativas para diferentes tipos de aceleradores.
v5litepod-4 | v5litepod-16 | v5litepod-64 | |
Ciclo de entrenamiento | 3 | 3 | 3 |
Tamaño del lote global | 32 | 128 | 512 |
Capacidad de procesamiento (ejemplos/s) | 201 | 657 | 2,844 |
TensorFlow 2.x
En las siguientes secciones, se describen ejemplos de cómo entrenar modelos de TensorFlow 2.x en TPU v5e.
Entrena Resnet en un v5e de un solo host
En este instructivo, se describe cómo entrenar ImageNet en v5litepod-4
o v5litepod-8
con un conjunto de datos falso. Si deseas usar un conjunto de datos diferente, consulta Prepara el conjunto de datos.
Configurar
Crea variables de entorno:
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
puede serv5litepod-4
ov5litepod-8
.-
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}
Podrás establecer una conexión SSH a tu VM de TPU una vez que tu recurso en cola esté en el estado
ACTIVE
. Para verificar el estado del recurso en cola, usa el siguiente comando:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Conéctate a la TPU con SSH
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Configura algunas variables de entorno
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
Cambia al directorio del repositorio de modelos y los requisitos de instalación.
cd ${MODELS_REPO} && git checkout r2.15.0 pip install -r official/requirements.txt
Entrenar el modelo
Ejecuta la secuencia de comandos de entrenamiento.
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"
Borra la TPU y el recurso en cola
Borra tu TPU
gcloud compute tpus tpu-vm delete ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet
Borra tu solicitud de recurso en cola
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet
Entrena a Resnet en un v5e de varios hosts
En este instructivo, se describe cómo entrenar ImageNet en v5litepod-16
o en versiones más grandes con un conjunto de datos falso. Si deseas usar un conjunto de datos diferente, consulta Prepara el conjunto de datos.
Crea variables de entorno:
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
puede serv5litepod-16
o mayor.-
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}
Podrás establecer una conexión SSH a tu VM de TPU una vez que tu recurso en cola esté en el estado
ACTIVE
. Para verificar el estado de tu recurso en cola, usa el siguiente comando:gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Conéctate a tu TPU (trabajador cero) con SSH
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE}
Configura algunas variables de entorno
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
Cambia al directorio del repositorio de modelos y los requisitos de instalación.
cd $MODELS_REPO && git checkout r2.15.0 pip install -r official/requirements.txt
Entrenar el modelo
Ejecuta la secuencia de comandos de entrenamiento.
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*"
Borra la TPU y el recurso en cola
Borra tu TPU
gcloud compute tpus tpu-vm delete ${TPU_NAME} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet
Borra tu solicitud de recurso en cola
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \ --project=${PROJECT_ID} \ --zone=${ZONE} \ --quiet