Entrenamiento de Cloud TPU v5p

Cloud TPU v5p es la quinta generación de Cloud TPU de Google Cloud y el sucesor de la TPU v4. v5p está optimizada para el entrenamiento a gran escala y para ser una plataforma líder en el desarrollo de LLM fundamentales, modelos de difusión y IA generativa. En un nivel alto, v5p proporciona hasta el doble del rendimiento de v4 y, al mismo tiempo, incluye el doble de TPU en un pod (6,000 de la porción más grande en comparación con 3,000 en v4), lo que genera hasta 4 veces el rendimiento a nivel del pod. También se ejecuta a una frecuencia de reloj más alta (1.75 GHz en comparación con 1.05 GHz), agrega SparseCore para incorporaciones a gran escala y triplica la capacidad de la memoria de gran ancho de banda (HBM).

Conceptos de Cloud TPU v5p

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

En la página Arquitectura del sistema de Cloud TPU, se describen los conceptos de Cloud TPU (por ejemplo, slices, hosts y TensorCores) y la arquitectura del sistema de Cloud TPU para todas las versiones de Cloud TPU.

Cada versión de Cloud TPU requiere tipos de aceleradores específicos para la capacitación o la inferencia. Estos tipos de aceleradores se describen en las configuraciones de v5p.

Administra recursos de TPU

En todos los comandos de este documento, se da por sentado que estás creando una VM de TPU v5p. Para obtener más información sobre los comandos para crear VMs de TPU, consulta Administra TPUs o la guía del usuario de recursos en cola para administrar recursos en cola. Para que los comandos sean más fáciles de ejecutar, las muestras de código de este documento usan las siguientes variables de entorno:

export PROJECT_ID=your-project
export ACCELERATOR_TYPE=v5p-8
export ZONE=us-east5-a
export RUNTIME_VERSION=v2-alpha-tpuv5
export TPU_NAME=your-tpu-name

Descripciones de las variables de entorno

PROJECT_ID

Es el proyecto de Google Cloud en el que creas tu TPU.

ACCELERATOR_TYPE

El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.

ZONE

Es la zona en la que deseas crear la Cloud TPU.

RUNTIME_VERSION

La versión del software de la TPU.

TPU_NAME

Es el nombre definido por el usuario de la TPU que usas.

Configuración del framework

En esta sección, se describe el proceso de configuración general para el entrenamiento de modelos con JAX o PyTorch con TPU v5p.

Configuración para JAX

Si tienes formas de rebanada superiores a 4 chips, tendrás varias VMs en una rebanada. En este caso, debes usar la marca --worker=all para ejecutar la instalación en todas las VMs de TPU con un solo comando:

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

Puedes ejecutar el siguiente comando para verificar la cantidad de dispositivos (los resultados que se muestran aquí se produjeron con una porción v5p-32). Este código verifica que todo esté instalado correctamente comprobando 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 determinada. jax.local_device_count() indica la cantidad 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.]

Usa --node=all para ejecutar el comando en todos los trabajadores de Multislice.

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='python3 -c "import jax; print(jax.device_count()); print(jax.local_device_count())"'

Prueba los instructivos de JAX en este documento para comenzar a entrenar con v5p usando JAX.

Configuración para PyTorch

El entorno de ejecución de PJRT es el único entorno de ejecución compatible con v5p, y PyTorch 2.1 y versiones posteriores usan 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 Pods v5p con PyTorch/XLA 2.2.0 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
sudo apt-get install libopenblas-dev -y
pip install numpy
pip install torch torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html
pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu
'

Usa una secuencia de comandos de Python con PJRT para validar la instalación. La secuencia de comandos muestra los dispositivos TPU disponibles (los resultados que se muestran aquí se produjeron con una porción v5p-32).

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
--project ${PROJECT_ID} --zone ${ZONE} --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

SSH: Attempting to connect to worker 0...
SSH: Attempting to connect to worker 1...
SSH: Attempting to connect to worker 2...
SSH: Attempting to connect to worker 3...
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']
['xla:0', 'xla:1', 'xla:2', 'xla:3']

Usa --node=all para ejecutar el comando en todos los trabajadores de Multislice.

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
--project ${PROJECT_ID} --zone ${ZONE} --node=all --worker=all \
--command='
PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"
'

Prueba los instructivos de PyTorch que se incluyen en este documento para comenzar a usar el entrenamiento de v5p con PyTorch.

Supervisa y crea perfiles

Cloud TPU v5p admite la supervisión y la generación de perfiles con los mismos métodos que las generaciones anteriores de Cloud TPU. Puedes consultar Perfila tu modelo con las herramientas de Cloud TPU para obtener más información sobre el perfilado y Supervisa las VMs de Cloud TPU para obtener más información sobre la supervisión.

Instructivos de capacitación

En esta sección, se incluyen instructivos de capacitación para un solo segmento. Para adaptar estos instructivos al entrenamiento de Multislice, puedes agregar la marca --node=all a los comandos SSH. Para obtener detalles y prácticas recomendadas, consulta la introducción a Multislice.

Nota: Antes de ejecutar cualquiera de estos instructivos, debes configurar tu entorno de Cloud TPU. Una vez que hayas configurado tu entorno de TPU, estará todo listo para ejecutar los siguientes instructivos.

• Difusión 2.1
• MaxText
• ResNet en un solo host v5p
• ResNet en un v5p de varios hosts

Instructivos de JAX

Train Diffusion 2.1

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

El modelo de difusió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:

• Descripción general de la Difusión estable
• Fuente de código estable de Diffusion

Nota: Antes de ejecutar este instructivo, debes configurar tu entorno de Cloud TPU.

Configurar

1. Crea variables de entorno:

   export GCS_BUCKET_NAME=your-bucket
   export PROJECT_ID=your-project-ID
   export ACCELERATOR_TYPE=v5p-32
   export ZONE=europe-west4-b
   export LOCATION=europe-west4
   export RUNTIME_VERSION=v2-alpha-tpuv5
   export SERVICE_ACCOUNT=your-service-account
   export TPU_NAME=your-tpu-name
   export QUEUED_RESOURCE_ID=your-qr-name
   export QUOTA_TYPE=spot
   export VALID_UNTIL_DURATION=1d

   Descripciones de las marcas de comandos

   Variable	Descripción
   ID DEL PROYECTO	Nombre del proyecto de Google Cloud
   ACCELERATOR_TYPE	Consulta la página Versiones de TPU para conocer la versión de TPU.
   ZONA	Consulta el documento Regiones y zonas de TPU para conocer las zonas compatibles.
   UBICACIÓN	Es la región de Google Cloud en la que se creará el bucket de almacenamiento de Cloud Storage.
   RUNTIME_VERSION	Para v5p, usa v2-alpha-tpuv5 para RUNTIME_VERSION.
   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: [e-mail]
   TPU_NAME	El ID de texto asignado por el usuario de la TPU que se crea cuando se asigna la solicitud de recurso en cola.
   QUEUED_RESOURCE_ID	El ID de texto asignado por el usuario de la solicitud de recursos en cola. Consulta el documento Recursos en cola para obtener información sobre los recursos en cola.
   QUOTA_TYPE	Puede ser reserved o spot. Si no se especifica ninguno de estos, el valor predeterminado de QUOTA_TYPE es on-demand. Consulta cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.
   VALID_UNTIL_DURATION	Es la duración de la validez de la solicitud. Consulta Recursos en fila para obtener información sobre las diferentes duraciones válidas.

2. Configura un bucket de almacenamiento para el resultado de tu modelo.

   gcloud storage buckets create gs://$GCS_BUCKET_NAME \
    --project=$PROJECT_ID \
    --location=$LOCATION

3. Crea un recurso de TPU:

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

   Nota: La marca QUOTA_TYPE puede ser reserved o spot. Si no se especifica ninguno de estos, el valor predeterminado es on-demand. Consulta 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 el recurso en cola esté en el estado ACTIVE. Ejecuta el siguiente comando para verificar el estado de tu recurso en cola:

   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
   

4. Entrena el modelo

   gcloud compute tpus tpu-vm ssh $TPU_NAME --zone=$ZONE --project $PROJECT_ID --worker=all --command="
   git clone https://github.com/google/maxdiffusion
   cd maxdiffusion
   git reset --hard 57629bcf4fa32fe5a57096b60b09f41f2fa5c35d # This identifies the GitHub commit to use.
   pip3 install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html # Install the latest version of JAX
   pip3 install -r requirements.txt
   pip3 install .
   export LIBTPU_INIT_ARGS=""
   python -m src.maxdiffusion.models.train src/maxdiffusion/configs/base_2_base.yml run_name=my_run base_output_directory=gs://$GCS_BUCKET_NAME enable_profiler=False"

Limpia

Borra tu TPU y la solicitud de recursos en cola al final de la sesión o para quitar las solicitudes de recursos en cola que estén en el estado "FAILED". Para borrar un recurso en cola, borra las porciones y, luego, la solicitud de 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

O bien, usa --force para borrar las porciones y la solicitud de recursos en cola en un solo paso:

# With --force
gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID}
--project ${PROJECT_ID} --zone ${ZONE} --quiet --force

Resultados de las comparativas

La secuencia de comandos de entrenamiento de Stable Diffusion se ejecutó en v5p-8, v5p-32 y v5p-128. En la siguiente tabla, se muestra la capacidad de procesamiento.

	v5p-8	v5p-32	v5p-128
Paso de entrenamiento	150	150	150
Tamaño del lote global	32	64	64
Capacidad de procesamiento (ejemplos/s)	12.10	18.08	19.10

MaxText

En este instructivo, se muestra cómo entrenar el modelo MaxText con un conjunto de datos sintético en Cloud TPU.

MaxText es un LLM de alto rendimiento, escalable de forma arbitraria, de código abierto y bien probado, escrito en Python/JAX puro y orientado a Cloud TPU. MaxText brinda a los investigadores y desarrolladores una herramienta accesible y adaptable para avanzar en las fronteras de la investigación y el desarrollo del procesamiento de lenguaje natural (PLN).

Antes de ejecutar este instructivo, debes configurar tu entorno de Cloud TPU.

1. Configura variables de entorno

   export PROJECT_ID=your_project_ID
   export TPU_NAME=your_tpu_name # user defined TPU name
   export ACCELERATOR_TYPE=v5p-256
   export ZONE=us-east5-a
   export RUNTIME_VERSION=v2-alpha-tpuv5
   export RUN_NAME=your_experiment_run_name # user defined name for this run
   export GCS_BUCKET_NAME=your_bucket_name # Output cloud folder. Should start with gs://
   export MAXTEXT_OUTPUT_PATH=${GCS_BUCKET_NAME}/your_experiment_output_path
   export NUM_SLICES=1 # Update the value to a number >1 for Multislice.

   Descripciones de las marcas de comandos

   Variable	Descripción
   ID DEL PROYECTO	Nombre del proyecto de Google Cloud
   TPU_NAME	Un nombre definido por el usuario para tu TPU.
   ACCELERATOR_TYPE	Consulta la página Versiones de TPU para conocer la versión de TPU.
   ZONA	Consulta el documento Regiones y zonas de TPU para conocer las zonas compatibles.
   RUNTIME_VERSION	Para v5p, usa v2-alpha-tpuv5 para la versión del entorno de ejecución.
   RUN_NAME	Es el nombre de la ejecución del experimento que proporcionó el usuario.

   Configuración opcional recomendada para Multislice:

   export NETWORK_NAME=your_network_name
   export FIREWALL_RULE_NAME=your_firewall_rule_name

   Si ejecutas cargas de trabajo de Multislice y deseas un rendimiento óptimo de la red, considera crear una red dedicada con una unidad de transmisión máxima (MTU) de 8,896 bytes y configurar las reglas de firewall adecuadas. Si bien es opcional, este paso puede mejorar significativamente el rendimiento, en especial cuando se aumenta la cantidad de rebanadas en la red del centro de datos (DCN). Ten en cuenta que crear una red requiere el permiso compute.networks.create en el proyecto. En los siguientes ejemplos, se muestra cómo crear una red y una regla de firewall dedicadas.

   Crea una red dedicada:

   gcloud compute networks create ${NETWORK_NAME} \
   --mtu=8896 \
   --project=${PROJECT_ID} \
   --subnet-mode=auto \
   --bgp-routing-mode=regional

   Crea una regla de firewall:

   gcloud compute firewall-rules create ${FIREWALL_RULE_NAME} \
   --network ${NETWORK_NAME} --allow tcp,icmp,udp --project=${PROJECT_ID}

2. Clona el repositorio de MaxText

   git clone https://github.com/google/maxtext.git

3. Entrena el modelo

   En las siguientes secciones, se describen dos opciones para entrenar MaxText.

   Opción 1

   Si deseas que una secuencia de comandos administre todo el flujo de trabajo, desde el aprovisionamiento de Cloud TPU y la instalación de dependencias hasta la ejecución de tu modelo y la eliminación de recursos, puedes usar multihost_job.py.

   cd maxtext && python3 multihost_job.py --PROJECT=${PROJECT_ID} --ZONE=${ZONE} \
   --NUM_SLICES=${NUM_SLICES} --TPU_TYPE=${ACCELERATOR_TYPE} \
   --VERSION=${RUNTIME_VERSION} --RUN_NAME=${RUN_NAME} #user defined run name \
   --BUCKET_NAME=${GCS_BUCKET_NAME} \ #used to store logs and configs
   --COMMAND="bash setup.sh && bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME} OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

   Nota: Agrega --CQR_EXTRA_ARGS="--network=${NETWORK_NAME}" si usas una red personalizada.

   Después de iniciar la secuencia de comandos, deberías ver un mensaje similar al siguiente en el registro. Se hace referencia a la ubicación del registro en el mensaje de salida. Haz clic en el primer vínculo para acceder a los registros de todos los trabajadores una vez que se complete el aprovisionamiento de TPU.

   ------------------------------------
   
   multihost_job finished running, TPUs are starting up to run your job remotely.
   
   Logs for your job are displayed here:
   https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22_log%22%2529;?project=PROJECT_ID
   
   To see the output of a single host, you may edit the slice and worker
   number in the `log_file_path` property here:
   
   https://console.cloud.google.com/logs/query;query=resource.type%3D%22gce_instance%22%20AND%0Alog_id%2528%22RUN_NAME_log%22%2529%20AND%0Alabels.%22agent.googleapis.com%2Flog_file_path%22%3D%20%22%2FRUN_NAME%2Fmain_command_log_slice_0_worker_0%22;?project=PROJECT_ID
   
   When your job is finished, the main command log is in your Cloud Storage
   bucket:
   https://console.cloud.google.com/storage/browser/YOUR_BUCKET_NAME/RUN_NAME?project=PROJECT_ID
   
   View the status of the created TPUs using:
   gcloud compute tpus queued-resources list --filter=RUN_NAME
   --zone=ZONE --project=PROJECT_ID
   

Opción 2

Para ejecutar la secuencia de comandos de entrenamiento varias veces en una Cloud TPU aprovisionada, usa la secuencia de comandos multihost_runner.py para usar el recurso.

1. Configura variables para crear una TPU.

   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export VALID_DURATION=1d
   export QUOTA_TYPE=quota_type

   Nota: Las siguientes marcas solo son necesarias si usas la función Multislice:

   --node-count ${NODE_COUNT} \
   --node-prefix ${NODE_PREFIX} # optional, the default is QUEUED_RESOURCE_ID

   Nota: Si especificas --node-count y --node-prefix, no especifiques --node-id cuando crees el recurso de TPU en el siguiente paso. Se generará un error si se especifican las tres marcas.

2. Crea un recurso de TPU.

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

   Nota: La marca QUOTA_TYPE puede ser reserved o spot. Si no se especifica ninguno de estos, el valor predeterminado es on-demand. Consulta cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.

   Podrás conectarte a tus VMs de TPU con SSH una vez que tu QueuedResource esté en el estado ACTIVE:

   Usa el comando describe para consultar el estado de tu recurso en cola.

   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. Conéctate a tu TPU con SSH

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

4. Instala dependencias

   Nota: Una vez que hayas usado ssh para conectarte a tu VM de TPU, debes configurar las variables ${TPU_NAME} y ${MAXTEXT_OUTPUT_PATH} para usarlas en la VM de TPU:

   export TPU_NAME=your_tpu_name
   export MAXTEXT_OUTPUT_PATH=output-path

   cd maxtext && python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
   --COMMAND='bash setup.sh'

5. Ejecuta el modelo con varias secuencias de comandos de configuración, como 32b.sh y 64b.sh. Si ejecutas la secuencia de comandos desde una VM de TPU, debes agregar la marca --INTERNAL_IP=true.

   python3 multihost_runner.py --TPU_PREFIX=${TPU_NAME} \
   --COMMAND="bash MaxText/configs/experimental/64b.sh RUN_NAME=${RUN_NAME}
   OUTPUT_PATH=${MAXTEXT_OUTPUT_PATH} PLATFORM=gce"

   Nota: Establece --TPU_PREFIX=${NODE_PREFIX} para Multislice en lugar de --TPU_PREFIX=${TPU_NAME}.

Limpia

Borra tus TPU y recursos en cola.

Resultados de las comparativas

La secuencia de comandos de entrenamiento de MaxText se ejecutó de 32 B a 1160 B con precisión bf16. Los resultados de estas ejecuciones se muestran en la siguiente tabla.

Cantidad de parámetros	Tipo de acelerador	TFLOP/chip/s	Utilización de flops del modelo (MFU)
32 B	v5p-128	3.28E+02	71.47%
64 B	v5p-128	3.23E+02	70.31%
128 B	v5p-256	3.15E+02	68.68%
128 B	v5p-512	3.15E+02	68.53%
256 B	v5p-1024	3.16E+02	68.82%
512 B	v5p-1024	2.94E+02	63.99%
1,024 B	v5p-2048	2.49E+02	64.05%
1,024 B	v5p-4096	2.97E+02	64.80%
1,160 B	v5p-7680	2.95E+02	64.27%
1,160 B	v5p-12288	3.04E+02	66.23%

El modelo de parámetros de 256B se probó en v5p-512 y v5p-1024 con ponderaciones bf16 e int8. En la siguiente tabla, se muestran los resultados de estas pruebas.

	v5p-512	v5p-512	v5p-1024	v5p-1024
Tamaño del lote global (tokens)	5.24E+05	5.24E+05	1.05E+06	1.05E+06
Precisión	bf16	int8	bf16	int8
TFLOP/chip/s	307	408	308	414
Utilización de flops del modelo (MFU)	66.98%	88.85%	67.09%	90.23%

Instructivos de TensorFlow

Entrena ResNet en un solo host v5p

En este instructivo, se describe cómo entrenar ImageNet en una TPU v5p-8 con un conjunto de datos falso. Si quieres usar un conjunto de datos diferente, consulta Cómo preparar el conjunto de datos.

Nota: Antes de ejecutar este instructivo, debes configurar tu entorno de Cloud TPU.

Configurar

1. Crea variables de entorno:

   export PROJECT_ID=your-project-ID
   export ACCELERATOR_TYPE=v5p-32
   export ZONE=us-east1-c
   export RUNTIME_VERSION=tpu-vm-tf-2.18.0-pjrt
   export TPU_NAME=your-tpu-name
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type

   En este instructivo, usa v5p-8 como ACCELERATOR_TYPE.

2. Crea un recurso de TPU:

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

   Nota: La marca QUOTA_TYPE puede ser reserved o spot. Si no se especifica ninguno de estos, el valor predeterminado es on-demand. Consulta cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.

   Podrás conectarte a tu VM de TPU con SSH una vez que el recurso en fila esté en el estado ACTIVE. Para verificar el estado de tu recurso en cola, usa el siguiente comando:

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

3. Conéctate a tu 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, luego, instala los requisitos.

   cd ${MODELS_REPO} && git checkout r2.15.0
   pip install -r official/requirements.txt

Entrena el modelo

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

Limpia

Borra tus TPU y recursos en cola.

Entrena ResNet en un v5p de varios hosts

En este instructivo, se describe cómo entrenar ImageNet en v5p-16 o más con un conjunto de datos falso. Si quieres usar un conjunto de datos diferente, consulta Prepara el conjunto de datos.

Nota: Antes de ejecutar este instructivo, debes configurar tu entorno de Cloud TPU.

1. Crea variables de entorno:

   export PROJECT_ID=your_project_ID
   export TPU_NAME=your_tpu_name
   export ZONE=us-east1-c
   export ACCELERATOR_TYPE=v5p-16
   export RUNTIME_VERSION=tpu-vm-tf-2.18.0-pod-pjrt
   export QUEUED_RESOURCE_ID=your-queued-resource-id
   export QUOTA_TYPE=quota-type

   ACCELERATOR_TYPE puede ser v5p-16 o mayor.

2. Crea un recurso de TPU:

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

   Nota: La marca QUOTA_TYPE puede ser reserved o spot. Si no se especifica ninguno de estos, el valor predeterminado es on-demand. Consulta cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.

   Podrás conectarte a tu VM de TPU con SSH una vez que el recurso en fila esté en el estado ACTIVE.

   Usa el comando describe para consultar el estado de tu recurso en cola:

   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 TPU_NAME=your_tpu_name
   export MODELS_REPO=/usr/share/tpu/models
   export PYTHONPATH="${MODELS_REPO}:${PYTHONPATH}"
   export MODEL_DIR=gcp-directory-to-store-model
   export DATA_DIR=gs://cloud-tpu-test-datasets/fake_imagenet
   export TPU_LOAD_LIBRARY=0

5. Cambia al directorio del repositorio de modelos y, luego, instala los requisitos.

   cd $MODELS_REPO && git checkout r2.15.0
   pip install -r official/requirements.txt

Entrena el modelo

1. Ejecuta la secuencia de comandos de entrenamiento.

   python3 official/vision/train.py \
     --tpu=${TPU_NAME} \
     --experiment=resnet_imagenet \
     --mode=train_and_eval \
     --config_file=official/vision/configs/experiments/image_classification/imagenet_resnet50_tpu.yaml \
     --model_dir=${MODEL_DIR} \
     --params_override="runtime.distribution_strategy=tpu,task.train_data.input_path=${DATA_DIR}/train*,task.validation_data.input_path=${DATA_DIR}/validation*,task.train_data.global_batch_size=2048,task.validation_data.global_batch_size=2048,trainer.train_steps=100"

Limpia

Borra tus TPU y recursos en cola.

PyTorch/XLA

Llama 2

En este instructivo, se explica cómo entrenar el modelo Llama 2 7B en v5p con una bifurcación del repositorio de HuggingFace en PyTorch/XLA con paralelización general y escalable para gráficos de procesamiento de AA (GSPMD).

Configuración

1. Crea variables de entorno.

   export PROJECT_ID=your_project_ID
   export ACCELERATOR_TYPE=v5p-8
   export ZONE=us-east5-a
   export RUNTIME_VERSION=v2-alpha-tpuv5
   export SERVICE_ACCOUNT=your_service_account
   export TPU_NAME=your_tpu_name
   export QUEUED_RESOURCE_ID=your_queued_resource_id
   export QUOTA_TYPE=quota_type
   export VALID_DURATION=1d

2. Crea un recurso de TPU

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

   Nota: La marca QUOTA_TYPE puede ser reserved o spot. Si no se especifica ninguno de estos, el valor predeterminado es on-demand. Consulta cuotas para obtener información sobre los diferentes tipos de cuotas que admite Cloud TPU.

   Podrás conectarte a tu VM de TPU con SSH una vez que tu QueuedResource esté en el estado ACTIVE:

   Usa el comando describe para consultar el estado de tu recurso en cola.

   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
   

   Nota: Si ejecutas en un entorno de varias porciones, usa gcloud compute tpus queued-resources ssh --node=all ... en lugar de gcloud compute tpus tpu-vm ssh ... para tus comandos. Esto garantiza que todos los chips de TPU se usen correctamente.

3. Instala Pytorch/XLA y las dependencias necesarias.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME}  \
   --project ${PROJECT_ID} \
   --zone  ${ZONE} \
   --worker=all \
   --command='
   sudo apt-get update
   sudo apt-get install libopenblas-dev -y
   pip3 install numpy
   pip3 install typing-extensions
   pip install torch torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html
   '

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 -b llama2-google-next-training https://github.com/pytorch-tpu/transformers.git
   cd transformers
   pip3 install git+file://$PWD
   pip3 install datasets accelerate evaluate scikit-learn'

5. Descarga la configuración del modelo 7B.

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command="curl https://huggingface.co/TheBloke/Llama-2-7B-fp16/raw/main/config.json --output ~/config.json"

6. Entrena el modelo

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --project=${PROJECT_ID} \
   --zone=${ZONE} \
   --worker=all \
   --command='
   export PJRT_DEVICE=TPU
   export XLA_USE_BF16=1
   export XLA_IR_DEBUG=1
   export XLA_HLO_DEBUG=1
   
   export LIBTPU_INIT_ARGS="--xla_enable_async_collective_permute=true \
   --xla_tpu_enable_async_collective_fusion_multiple_steps=true \
   --xla_tpu_enable_async_collective_fusion=true \
   --xla_tpu_overlap_compute_collective_tc=true \
   --xla_enable_async_all_gather=true \
   --xla_jf_spmd_threshold_for_windowed_einsum_mib=0"
   
   export PROFILE_EPOCH=0
   export PROFILE_STEP=3
   export PROFILE_DURATION_MS=20000
   export PROFILE_LOGDIR=/tmp/home/
   
   cd transformers
   python examples/pytorch/language-modeling/run_clm.py \
    --tokenizer_name hf-internal-testing/llama-tokenizer \
    --dataset_name wikitext \
    --dataset_config_name wikitext-2-raw-v1 \
    --per_device_train_batch_size 96 \
    --per_device_eval_batch_size 8 \
    --num_train_epochs 1 \
    --do_train \
    --output_dir /tmp/output \
    --overwrite_output_dir \
    --config_name ~/config.json \
    --save_strategy no \
    --logging_strategy no \
    --remove_unused_columns no \
    --optim adafactor \
    --torch_dtype bfloat16 \
    --dataloader_drop_last yes \
    --block_size 2048 \
    --spmd_2d_sharding 1 \
    --spmd_grad_chkpt
   '

Si ejecutas en un entorno de varias porciones, debes configurar la marca --spmd_dcn_parallelism en la cantidad de porciones.

SPMD_USER_GUIDE proporciona una guía del usuario más detallada que explica todas las diferentes variables de entorno y los botones de activación de la secuencia de comandos de HF. Ten en cuenta que el elemento LIBTPU_INIT_ARGS se incorporará a PyTorch/XLA y se activará de forma predeterminada en versiones futuras.

Limpia

Borra tus TPU y recursos en cola.

Resultados de las comparativas

La capacidad de procesamiento de los tres tamaños de modelos Llama 2 se incluye en la siguiente tabla.

	v5p-8	v5p-128	v5p-128
Tamaño del modelo	7B	13B	70,000 millones
Tamaño del lote global	96	1024	128
Forma de la malla de fragmentación	(4, 1)	(64, 1)	(16, 4)
Utilización de flops del modelo (MFU)	56.67%	55.80%	51.85%

Asistencia y comentarios

Todos los comentarios son bienvenidos. Para compartir comentarios o solicitar asistencia, completa el formulario de comentarios o asistencia de Cloud TPU.