Entrenamiento de Cloud TPU v5e

Cloud TPU v5e es el acelerador de IA de última generación de Google Cloud. Con un de 256 chips por Pod, se optimiza una v5e para que tenga el valor producto para transformador, texto a imagen y red neuronal convolucional (CNN) para el entrenamiento, el ajuste y la entrega. Para obtener más información sobre cómo usar Cloud TPU v5e para entrega, consulta Inferencia con v5e.

Para obtener más información sobre el hardware y la configuración 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 solicitud

Necesitas cuota si quieres usar la TPU v5e para el entrenamiento. Hay diferentes tipos de cuotas las TPU bajo demanda, las TPU reservadas y las VMs Spot de TPU. Existen diferentes cuotas obligatorias si usas tu TPU v5e para inferencia. Para obtener más información sobre las cuotas, consulta Cuotas Para solicitar una cuota de TPU v5e, comunícate con Cloud Ventas

Crea una cuenta y un proyecto de Google Cloud

Necesitas una cuenta y un proyecto de Google Cloud para usar Cloud TPU. Para ver más consulta Configura una Cloud TPU entorno.

Nota: Para usar TPU v5e, debes instalar el comando gcloud alpha componentes.

Crear una Cloud TPU

La práctica recomendada es aprovisionar Cloud TPU v5es en cola recursos con queued-resource create kubectl. Para obtener más información, consulta Administra elementos en cola recursos.

También puedes usar la API de Create Node (gcloud alpha compute tpus tpu-vm create) para aprovisionar Cloud TPU v5es. Para obtener más información, consulta Administra la TPU. recursos.

Para obtener más información configuraciones de v5e para el entrenamiento; consulta Tipos de Cloud TPU v5e para entrenamiento.

Configuración del framework

En esta sección, se describe el proceso de configuración general del entrenamiento de modelos personalizados con JAX o PyTorch con TPU v5e. La asistencia de TensorFlow está disponible en la TPU tpu-vm-tf-2.17.0-pjrt y tpu-vm-tf-2.17.0-pod-pjrt del entorno de ejecución.

Para obtener instrucciones sobre la configuración de la inferencia, consulta Inferencia v5e introducción.

Configuración para JAX

Si tienes cortes 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 VMs de TPU en un solo paso sin usar SSH para acceder a cada 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 se produjeron con un segmento v5litepod-16). Este código prueba que todo esté instalado correctamente verificando que JAX detecte los TensorCores de Cloud TPU y puede 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 un solo elemento. 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 Entrenamiento de la v5e con JAX.

Configuración para PyTorch

Ten en cuenta que la v5e solo es compatible con el entorno de ejecución PJRT. y PyTorch 2.1 y versiones posteriores usarán 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 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 se produce un error cuando instalas las ruedas de torch, torch_xla o torchvision me gusta pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222, cambiar a una versión inferior con este comando:

pip3 install setuptools==62.1.0

Ejecuta una secuencia de comandos con PJRT

Nota: Para los modelos que tienen asignaciones considerables y frecuentes, usa tcmalloc. puede mejorar significativamente el tiempo de entrenamiento en comparación con el valor predeterminado malloc por lo que el malloc predeterminado que se usa en una VM de TPU es tcmalloc Sin embargo, según la carga de trabajo (por ejemplo, con DLRM, que tiene asignaciones muy grandes para sus tablas de incorporación) tcmalloc puede causar una demora. En este caso, puedes intentar desactivar la siguiente variable usando el malloc predeterminado:

unset LD_PRELOAD

El siguiente es un ejemplo en el que se usa una secuencia de comandos de Python para hacer 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')

Para comenzar, prueba los instructivos de PyTorch de este documento Entrenamiento de v5e con 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 se encuentren el estado FAILED

Ejemplos de JAX/FLAX

En las siguientes secciones, se describen ejemplos de cómo entrenar modelos JAX y FLAX en TPU v5e

Entrena ImageNet en v5e

En este instructivo, se describe cómo entrenar ImageNet en v5e con datos de entrada falsos. Si si quieres usar datos reales, consulta el archivo README en GitHub.

Configurar

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

2. Crea un recurso de TPU:

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

   Nota: La marca QUOTA_TYPE puede ser reserved o best-effort. Si no se especifica ninguna de estas opciones, el valor predeterminado es on-demand Revisa las cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.

   Podrás conectarte mediante SSH a tu VM de TPU una vez que tu recurso en cola esté 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, la salida ser similar a lo siguiente:

    state: ACTIVE
   

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

   Nota: Si ves un error que sugiere que uses jax>=0.4.8, puedes ignorarlo. el mensaje.

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

5. Para generar datos falsos, el modelo necesita información sobre las dimensiones de 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 a partir 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 Transformador de Vision. (ViT) de HuggingFace usando la Imagenette de Fast AI conjunto de datos en Cloud TPU v5e.

El modelo ViT fue el primero en entrenar correctamente un codificador de transformadores en ImageNet con excelentes resultados comparados con las redes convolucionales. Para ver más consulta los siguientes recursos:

• Descripción general de ViT
• Recurso de código ViT

Configurar

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

2. Crea un recurso de TPU:

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

   Nota: La marca QUOTA_TYPE puede ser reserved o best-effort. Si no se especifica ninguna de estas opciones, el valor predeterminado es on-demand Revisa las cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.

   Podrás conectarte mediante 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 por lo siguiente:

    state: ACTIVE
   

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

4. Descargar el repositorio de Hugging Face y 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==2.17.0 && pip install -r examples/flax/vision/requirements.txt'
   

5. 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. El En la siguiente tabla, se muestran las capacidades de procesamiento con diferentes tipos de aceleradores.

Tipo de acelerador	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	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 con los HuggingFace usando el Pokémon conjunto de datos en Cloud TPU v5e.

El modelo de dispersión estable es un modelo de texto a imagen latente que genera y fotorrealistas a partir de cualquier entrada de texto. Para obtener más información, consulta los siguientes recursos:

• Descripción general de la dispersión estable
• Código fuente de dispersión estable

Configurar

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

2. Crea un recurso de TPU:

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

   Nota: La marca QUOTA_TYPE puede ser reserved o best-effort. Si no se especifica ninguna de estas opciones, el valor predeterminado es on-demand Revisa las cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.

   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
   

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

4. Descarga el repositorio de HuggingFace y 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.17.0 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	1500	1500	1500
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 etiquetado manual. Se entrenó para predecir la siguiente palabra en oraciones. Para obtener más información, consulta los siguientes recursos:

• Descripción general de GPT2
• Fuente del código GPT2

Configurar

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

2. Crea un recurso de TPU:

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

   Nota: La marca QUOTA_TYPE puede ser reserved o best-effort. Si no se especifica ninguna de estas opciones, el valor predeterminado es on-demand Revisa las cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.

   Podrás conectarte mediante 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 por lo siguiente:

    state: ACTIVE
   

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

4. Descargar el repositorio de HuggingFace y 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'
   

5. 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 && gcloud storage cp gs://cloud-tpu-tpuvm-artifacts/v5litepod-preview/jax/gpt . --recursive'
   

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
Epoch	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 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. Estos son los se actualizaron las instrucciones de configuración de v5e para las cargas de trabajo de entrenamiento de PyTorch/XLA.

Configurar

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

2. Crea un recurso de TPU:

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

   Nota: La marca QUOTA_TYPE puede ser reserved o best-effort. Si no se especifica ninguna de estas opciones, el valor predeterminado es on-demand Revisa las cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.

   Podrás establecer una conexión SSH a tu VM de TPU una vez que tu recurso 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
   

3. 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 la versión 5e con el repositorio de HuggingFace. en PyTorch/XLA mediante el conjunto de datos wikitext.

Configurar

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

2. Crea un recurso de TPU:

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

   Nota: La marca QUOTA_TYPE puede ser reserved o best-effort. Si no se especifica ninguna de estas opciones, el valor predeterminado es on-demand Revisa las cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.

   Podrás conectarte mediante SSH a tu VM de TPU una vez que QueuedResource se encuentre ACTIVE:

   gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Cuando el recurso en cola se encuentre en el estado ACTIVE, la salida ser similar a lo siguiente:

   state: ACTIVE
   

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

4. Descargar el repositorio de HuggingFace y 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
         '
   

5. 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='
         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/'
   

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. El En la siguiente tabla, se muestran las capacidades de procesamiento de las comparativas para diferentes tipos de aceleradores.

	v5litepod-4	v5litepod-16	v5litepod-64
Epoch	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 la versión 5e con el repositorio de HuggingFace. en PyTorch/XLA en el conjunto de datos cifar10.

Configurar

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

2. Crea un recurso de TPU:

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

   Nota: La marca QUOTA_TYPE puede ser reserved o best-effort. Si no se especifica ninguna de estas opciones, el valor predeterminado es on-demand Revisa las cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.

   Podrás conectarte mediante SSH a tu VM de TPU una vez que QueuedResource se encuentre el estado ACTIVE:

    gcloud compute tpus queued-resources describe ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE}
   

   Cuando el recurso en cola se encuentre en el estado ACTIVE, la salida ser similar a lo siguiente:

    state: ACTIVE
   

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

4. Descargar el repositorio de HuggingFace y 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
Epoch	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 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 quieres usar un conjunto de datos diferente, consulta Prepara el conjunto de datos.

Configurar

1. 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 ser v5litepod-4 o v5litepod-8.

2. Crea un recurso de TPU:

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --${QUOTA_TYPE}
   

   Nota: La marca QUOTA_TYPE puede ser reserved o best-effort. Si no se especifica ninguna de estas opciones, el valor predeterminado es on-demand Revisa las cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.

   Podrás conectarte mediante SSH a tu VM de TPU una vez que tu recurso en cola esté 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}
   

3. Conéctate a la TPU con SSH

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

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

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

1. Borra tu TPU

   gcloud compute tpus tpu-vm delete ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet
   

2. 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 más grande con un conjunto de datos falso. Si deseas usar un conjunto de datos diferente, consulta Prepara el conjunto de datos.

1. 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 ser v5litepod-16 o mayor.

2. Crea un recurso de TPU:

   gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
      --node-id=${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --accelerator-type=${ACCELERATOR_TYPE} \
      --runtime-version=${RUNTIME_VERSION} \
      --${QUOTA_TYPE}
   

   Nota: La marca QUOTA_TYPE puede ser reserved o best-effort. Si no se especifica ninguna de estas opciones, el valor predeterminado es on-demand Revisa las cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.

   Podrás conectarte mediante SSH a tu VM de TPU una vez que tu recurso en cola esté el estado ACTIVE Para comprobar 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}
   

3. Conéctate a tu TPU (trabajador cero) con SSH

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

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

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

1. Borra tu TPU

   gcloud compute tpus tpu-vm delete ${TPU_NAME} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet
   

2. Borra tu solicitud de recurso en cola

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
      --project=${PROJECT_ID} \
      --zone=${ZONE} \
      --quiet