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, una v5e está optimizada para ser el producto de mayor valor para el entrenamiento, el ajuste y la entrega de redes neuronales convolucionales (CNN), de transformador y de texto a imagen.

Conceptos, arquitectura del sistema y configuración de Cloud TPU v5e

Si es la primera vez que usas Cloud TPU, consulta la página principal de la documentación de TPU.

Los conceptos generales de Cloud TPU (por ejemplo, porciones, hosts, chips y TensorCores) y la arquitectura del sistema de Cloud TPU se describen en la página Arquitectura del sistema de Cloud TPU.

Cada versión de Cloud TPU requiere tipos de aceleradores específicos para el entrenamiento y la inferencia. Estos tipos de aceleradores se describen en la sección Configuraciones de v5e.

Inferencia

La inferencia es el proceso de usar un modelo entrenado para hacer predicciones sobre datos nuevos. Lo usa el proceso de entrega.

Porciones

Una porción representa una colección de chips, todos ubicados dentro del mismo Pod conectado por interconexiones entre chips (ICI) de alta velocidad. La v5e tiene formas de porciones 2D. Consulta la tabla en la sección de configuración de v5e para conocer las formas de porción compatibles.

La forma de chip y la topología de chip también se refieren a las formas de porciones.

Entrega

La entrega es el proceso de implementación de un modelo de aprendizaje automático entrenado en un entorno de producción en el que se puede usar para realizar predicciones o tomar decisiones. La latencia y la disponibilidad a nivel de servicio son importantes para la entrega.

Comparación entre un host único y varios hosts

Un host es una computadora física (CPU) que ejecuta VMs. Un host puede ejecutar varias VMs a la vez.

Las Slices que usan menos de 8 chips usan como máximo un host. Las Slices de más de 8 chips tienen acceso a más de un solo host y pueden ejecutar entrenamiento distribuido mediante varios hosts. Consulta la página Arquitectura de sistemas de TPU para obtener detalles sobre las porciones y los chips.

v5e admite el entrenamiento y la inferencia de varios hosts (mediante SAX).

VM de TPU

Una máquina virtual que ejecuta Linux y que tiene acceso a las TPU subyacentes. Para las TPU v5e, cada VM de TPU tiene acceso directo a 1, 4 u 8 chips, según el tipo de acelerador que especifique el usuario. Una VM de TPU también se conoce como trabajador.

Trabajador

Consulta VM de TPU.

Comenzar

Para obtener información sobre el hardware de TPU v5e, consulta Arquitectura del sistema.

Protección de la capacidad

Comunícate con Ventas de Cloud para comenzar a usar Cloud TPU v5e para tus cargas de trabajo de IA.

Prepara un proyecto de Google Cloud

  1. Accede a tu Cuenta de Google. Si aún no lo has hecho, regístrate para obtener una cuenta nueva.
  2. En la consola de Google Cloud, selecciona o crea un proyecto de Google Cloud en la página del selector de proyectos.
  3. La configuración de facturación es obligatoria para todo uso de Google Cloud, así que asegúrate de que la facturación esté habilitada en tu proyecto.

  4. Instala gcloud alpha components.

  5. Si eres un usuario de TPU que reutiliza los componentes gcloud alpha existentes, actualízalos para asegurarte de que se admitan los comandos y las marcas relevantes:

    gcloud components update
    
  6. Habilita la API de TPU a través del siguiente comando gcloud en Cloud Shell. (También puedes habilitarlo en la consola de Google Cloud).

    gcloud services enable tpu.googleapis.com
    
  7. Habilita la cuenta de servicio de TPU.

    Las cuentas de servicio permiten que el servicio de Cloud TPU acceda a otros servicios de Google Cloud. Usar una cuenta de servicio administrada por el usuario es una práctica recomendada de Google Cloud. Sigue estas guías para crear y otorgar roles. Se necesitan los siguientes roles:

    • Administrador de TPU
    • Administrador de almacenamiento: Se necesita para acceder a Cloud Storage
    • Escritor de registros: Se necesita para escribir registros con la API de Logging
    • Escritor de métricas de Monitoring: Se necesita para escribir métricas en Cloud Monitoring.
  8. Configura el proyecto y la zona.

    El ID de tu proyecto es el nombre de tu proyecto que se muestra en la consola de Cloud.

    export PROJECT_ID=your-project-id
    export ZONE=us-west4-a
    
    gcloud compute tpus tpu-vm service-identity create --zone=${ZONE}
    
    gcloud auth login
    gcloud config set project ${PROJECT_ID}
    gcloud config set compute/zone ${ZONE}
    

Aprovisiona el entorno de Cloud TPU

La práctica recomendada es aprovisionar Cloud TPU v5 como recursos en cola con el comando queued-resource create. Sin embargo, también puedes usar la API de Crear nodo (gcloud alpha compute tpus tpu-vm create) para aprovisionar Cloud TPU v5es.

Crea variables de entorno

Configura las variables de entorno necesarias para la creación de TPU.

Reemplaza las variables (en rojo) en la siguiente lista con los valores que usarás para tu trabajo de entrenamiento o de inferencia.

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
Variable Descripción
PROJECT_ID Nombre del proyecto de Google Cloud
ACCELERATOR_TYPE Consulta la sección Tipos de aceleradores para conocer los tipos de aceleradores compatibles.
ZONA Toda la capacidad está en us-west4-a.
RUNTIME_VERSION Usa v2-alpha-tpuv5-lite.
SERVICE_ACCOUNT Esta es la dirección de tu cuenta de servicio, que puedes encontrar en la consola de Google Cloud -> IAM -> Cuentas de servicio. Por ejemplo: [email protected]
TPU_NAME El ID de texto asignado por el usuario de la TPU que se crea cuando se asigna la solicitud de recursos en cola.
QUEUED_RESOURCE_ID El ID de texto asignado por el usuario de la solicitud de recursos en cola. Consulta el documento sobre recursos en cola para obtener información acerca de los recursos en cola.
QUOTA_TYPE Puede ser reserved o best-effort. Si no se especifica ninguna de estas opciones, el valor predeterminado de RESOURCE_TYPE es on-demand. Consulta la sección quotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.
VALID_UNTIL_DURATION Es el tiempo durante el cual la solicitud es válida. Consulta los recursos en cola para obtener información sobre las diferentes duraciones válidas.

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}

Si el recurso en cola se crea correctamente, el estado dentro del campo response será WAITING_FOR_RESOURCES o FAILED. Si el recurso en cola está en estado WAITING_FOR_RESOURCES, significa que pasó una validación preliminar y espera capacidad. Una vez que la capacidad esté disponible, la solicitud pasará a PROVISIONING. Estar en estado WAITING_FOR_RESOURCES no significa que se garantiza que obtendrás la cuota asignada, y puede llevar un tiempo cambiar del estado WAITING_FOR_RESOURCES al estado ACTIVE. Si el recurso en cola se encuentra en el estado FAILED, el motivo de la falla aparecerá en el resultado. La solicitud vencerá si no se completa una solicitud en --valid-until-duration y el estado se vuelve "FAILED".

Podrás acceder a la VM de TPU con SSH una vez que el recurso en cola esté en el estado ACTIVE.

Usa los comandos [list](/tpu/docs/managing-tpus-tpu-vm) o [describe](/tpu/docs/managing-tpus-tpu-vm) para consultar el estado del recurso en cola.

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

El state representa el estado de un recurso en cola. Los estados se definen de la siguiente manera:

Estado Descripción
WAITING_FOR_RESOURCES Se recibió el comando create del recurso en cola y se comenzará a aprovisionar, en cuanto la capacidad esté disponible.
PROVISIONING Se están aprovisionando las porciones de TPU.
ACTIVE Todas las TPU se aprovisionan y están listas para usarse. Si se proporciona una secuencia de comandos de inicio, comenzará a ejecutarse en todas las TPU, cuando el estado del recurso en cola pase a ACTIVE.
FAILED No se pudieron aprovisionar los segmentos.
SUSPENDING Se están borrando una o más partes.
SUSPENDED Se borran todas las partes subyacentes, pero el recurso en cola permanece intacto hasta que se borra de forma explícita. En este momento, un recurso en cola suspendido no se puede reanudar y se debe borrar.
DELETING Se está borrando el recurso en cola.

Conéctate a la VM de TPU con SSH

En la siguiente sección, se describe cómo instalar objetos binarios en cada VM de TPU en tu porción de TPU y ejecutar código. En este contexto, una VM de TPU también se conoce como trabajador.

Consulta la sección Tipos de VM para determinar cuántas VM tendrá tu porción.

Para instalar los objetos binarios o ejecutar el código, conéctate a la VM de TPU con el comando tpu-vm ssh.

gcloud compute tpus tpu-vm ssh ${TPU_NAME}

Para acceder a una VM de TPU específica con SSH, usa la marca --worker que sigue un índice basado en 0:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} --worker=1

Si tienes formas de porción de más de 8 chips, tendrás varias VM en una porción. En este caso, usa la marca --worker=all para ejecutar la instalación en todas las VM de TPU sin tener que conectarte a cada una por separado. Por ejemplo:

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='pip install "jax[tpu]==0.4.16" -f http://storage.googleapis.com/jax-releases/libtpu_releases.html'

Administra

Todos los comandos que puedes usar para administrar las VMs de TPU se describen en Administra las TPU.

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.15.0-pjrt y tpu-vm-tf-2.15.0-pod-pjrt.

Para obtener instrucciones de configuración de inferencia, consulta Introducción a la inferencia v5e.

Configuración de JAX

Si tienes formas de porción de más de 8 chips, tendrás varias VM 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 http://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 una porción v5litepod-16). En este código, se comprueba que todo esté instalado correctamente. Para ello, verifica que JAX vea los TensorCores 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 especificada. 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 v5e con JAX.

Configuración de PyTorch

Ten en cuenta que v5e solo es compatible con el entorno de ejecución de PJRT y que 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 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 http://storage.googleapis.com/libtpu-releases/index.html
      pip3 install torch_xla[tpu] -f http://storage.googleapis.com/libtpu-releases/index.html'

Si no pudiste instalar las ruedas para torch/torch_xla/torchvision y ves un error como pkg_resources.extern.packaging.requirements.InvalidRequirement: Expected end or semicolon (after name and no valid version specifier) torch==nightly+20230222, cambia a una versión inferior con este comando:

pip3 install setuptools==62.1.0

Ejecuta una secuencia de comandos con PJRT:

unset LD_PRELOAD

El siguiente es un ejemplo que 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')

Prueba los instructivos de PyTorch en este documento para comenzar con el entrenamiento de la versión 5e 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 encuentran en el estado FAILED.

Supervisar y generar perfiles

Cloud TPU v5e admite la supervisión y la creación de perfiles con los mismos métodos que las generaciones anteriores de Cloud TPU. Puedes leer Cómo generar un perfil de tu modelo con las herramientas de Cloud TPU para obtener más información sobre la generación de perfiles y sobre la supervisión de VMs de Cloud TPU para obtener más información sobre la supervisión.

Ejemplos de JAX/FLAX

Entrena ImageNet en v5e

En este instructivo, se describe cómo entrenar ImageNet en v5e con datos de entrada falsos. Si deseas 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}
    

    Podrás establecer una conexión SSH a la 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 QueuedResource esté en el estado ACTIVE, el resultado será similar al siguiente:

     state: ACTIVE
    
  3. Instala la versión más reciente 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 http://storage.googleapis.com/jax-releases/libtpu_releases.html'
    
  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 http://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 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 http://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 completes 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 Hugging Face, que se implementan en FLAX, están listos para usarse 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 Vision Transformer (ViT) de HuggingFace con el conjunto de datos imagenette Fast AI en Cloud TPU v5e.

El modelo ViT fue el primero que entrenó correctamente un codificador Transformer en ImageNet con excelentes resultados en comparación con las redes convolucionales. Para obtener más información, consulta los siguientes recursos:

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}
    

    Podrás establecer una conexión SSH a la VM de TPU una vez que el recurso en cola se encuentre en 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, 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 http://storage.googleapis.com/jax-releases/libtpu_releases.html'
    
  4. Descarga el repositorio de Hugging Face y los requisitos de instalación:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
       --project=${PROJECT_ID} \
       --zone=${ZONE} \
       --worker=all \
       --command='git clone http://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'
    
  5. Descarga el conjunto de datos Imagenette:

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
       --project=${PROJECT_ID} \
       --zone=${ZONE} \
       --worker=all \
       --command='cd transformers && wget http://s3.amaz888onaws.com.fqhub.com/fast-ai-imageclas/imagenette2.tgz && tar -xvzf imagenette2.tgz'
    

Entrenar el modelo

Entrena el modelo con un búfer previamente asignado 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 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 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
Rendimiento (ejemplos/s) 263,40 429,34 470,71

Distribución de entrenamientos en Pokémon

En este instructivo, se muestra cómo entrenar el modelo de dispersión estable de HuggingFace con el conjunto de datos de Pokémon en Cloud TPU v5e.

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

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}
    

    Podrás establecer una conexión SSH a la 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 se encuentre 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 http://storage.googleapis.com/jax-releases/libtpu_releases.html'
    
  4. Descarga el repositorio de HuggingFace y, luego, instala los requisitos.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
       --project=${PROJECT_ID} \
       --zone=${ZONE} \
       --worker=all \
       --command='git clone http://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 previamente asignado 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 los recursos 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 de entrenamiento 1,500 1,500 1,500
Tamaño del lote global 32 64 128
Rendimiento (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 con el conjunto de datos de OSCAR en Cloud TPU v5e.

GPT2 es un modelo Transformer previamente entrenado con texto sin procesar y sin etiquetado manual. Se entrenó para predecir la siguiente palabra en las oraciones. Para obtener más información, consulta los siguientes recursos:

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}
    

    Podrás establecer una conexión SSH a la 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 se encuentre 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 http://storage.googleapis.com/jax-releases/libtpu_releases.html'
    
  4. Descarga el repositorio de HuggingFace y, luego, instala los requisitos.

    gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
       --project=${PROJECT_ID} \
       --zone=${ZONE} \
       --worker=all \
       --command='git clone http://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 archivos 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 previamente asignado 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 para 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
Rendimiento (ejemplos/s) 74,60 72,97 72,62

PyTorch y XLA

Entrena ResNet con el entorno de ejecución PJRT

PyTorch/XLA está migrando de XRT a PjRt desde PyTorch 2.0+. Estas son las instrucciones actualizadas para configurar la versión 5e de las cargas de trabajo de entrenamiento de PyTorch/XLA.

Configuración
  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}
    

    Podrás establecer una conexión SSH a la VM de TPU una vez que QueuedResource esté en 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, el resultado será similar al siguiente:

     state: ACTIVE
    
  3. Instala dependencias específicas 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 http://storage.googleapis.com/libtpu-releases/index.html
          pip3 install torch_xla[tpu] -f http://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 http://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 comparativas

En la siguiente tabla, se muestran las capacidades de procesamiento de comparativas.

Tipo de acelerador Capacidad de procesamiento (ejemplos/segundo)
v5litepod-4 4240 ex/s
v5litepod-16 10,810 ex/s
V5litepod-64 46,154 ex/s

Entrena GPT2 en v5e

En este instructivo, se explica cómo ejecutar GPT2 en v5e con el repositorio de HuggingFace en PyTorch/XLA con el conjunto de datos de wikitext.

Configuración

  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}
    

    Podrás establecer una conexión SSH a la VM de TPU una vez que QueuedResource esté en 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, el resultado será similar al siguiente:

    state: ACTIVE
    
  3. Instala las dependencias linterna/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 http://storage.googleapis.com/libtpu-releases/index.html
          pip3 install torch_xla[tpu] -f http://storage.googleapis.com/libtpu-releases/index.html'
    
  4. 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 http://github.com/pytorch/xla.git
          pip install --upgrade accelerate
          git clone http://github.com/huggingface/transformers.git
          cd transformers
          git checkout ebdb185befaa821304d461ed6aa20a17e4dc3aa2
          pip install .
          git log -1
          pip install datasets evaluate scikit-learn
          '
    
  5. Descarga los archivos 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 comparativas

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 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
Rendimiento (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.

Configuración

  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}
    

    Podrás establecer una conexión SSH a la VM de TPU una vez que 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 se encuentre en el estado ACTIVE, el resultado será similar al siguiente:

     state: ACTIVE
    
  3. Instala dependencias linterna/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 http://storage.googleapis.com/libtpu-releases/index.html
          pip3 install torch_xla[tpu] -f http://storage.googleapis.com/libtpu-releases/index.html'
    
  4. 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 http://github.com/suexu1025/transformers.git vittransformers; \
          cd vittransformers; \
          pip3 install .; \
          pip3 install datasets; \
          wget http://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 comparativas

En la siguiente tabla, se muestran las capacidades de procesamiento de 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
Rendimiento (ejemplos/s) 201 657 2.844

TensorFlow 2.x

Entrena Resnet en un solo host v5e

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

  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}
    

    Podrás establecer una conexión SSH a la VM de TPU una vez que el recurso en cola se encuentre 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}
    
  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 recursos en cola

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

Entrena Resnet en una versión 5e de varios hosts

En este instructivo, se describe cómo entrenar ImageNet en v5litepod-16 o más grandes mediante 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}
    

    Podrás establecer una conexión SSH a la VM de TPU una vez que el recurso en cola se encuentre 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}
    
  3. Conéctate a la TPU (worker 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 recursos en cola

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