Entrena en una TPU de host único con Pax


En este documento, se proporciona una breve introducción al trabajo con Pax en una TPU de host único (v2-8, v3-8 y v4-8).

Pax es un framework para configurar. y ejecutar experimentos de aprendizaje automático sobre JAX. Pax se enfoca en simplificar el AA a gran escala compartiendo componentes de infraestructura con frameworks de AA existentes y utilizando la biblioteca de modelado Praxis para la modularidad.

Objetivos

  • Configura recursos de TPU para el entrenamiento
  • Instala Pax en una TPU de host único
  • Entrena un modelo SPMD basado en transformadores con Pax

Antes de comenzar

Ejecuta los siguientes comandos para configurar gcloud para que use tu proyecto de Cloud TPU y, luego, instala los componentes necesarios para entrenar un modelo que ejecute Pax en una TPU de un solo host.

Instala Google Cloud CLI

Google Cloud CLI contiene herramientas y bibliotecas para interactuar con los productos y servicios de Google Cloud CLI. Si aún no lo has instalado instálala ahora con las instrucciones que se indican en Cómo instalar Google Cloud CLI.

Configura el comando gcloud

Ejecuta gcloud auth list para ver las cuentas disponibles.

$ gcloud config set account account

$ gcloud config set project project-id

Habilita la API de Cloud TPU

Habilita la API de Cloud TPU con el siguiente comando de gcloud en Cloud Shell. (también puedes habilitarlo desde la consola de Google Cloud).

$ gcloud services enable tpu.googleapis.com

Ejecuta el siguiente comando para crear una identidad de servicio (una cuenta de servicio).

$ gcloud beta services identity create --service tpu.googleapis.com

Crea una VM de TPU

Con las VMs de Cloud TPU, tu modelo y código se ejecutan directamente en la VM de TPU. Establece una conexión SSH directamente en la VM de TPU. Puedes ejecutar código arbitrario, instalar paquetes, ver registros y depurar código directamente en la VM de TPU.

Para crear la VM de TPU, ejecuta el siguiente comando desde un Cloud Shell o la terminal de tu computadora donde Google Cloud CLI está instalado.

Configura el zone según la disponibilidad en tu contrato. hacer referencia a regiones y zonas de TPU si es necesario.

Establece la variable accelerator-type en v2-8, v3-8 o v4-8.

Establece la variable version en tpu-vm-base para las versiones de TPU v2 y v3, o en tpu-vm-v4-base para las TPU v4.

$ gcloud compute tpus tpu-vm create tpu-name \
--zone zone \
--accelerator-type accelerator-type \
--version version

Conéctate a tu VM de Google Cloud TPU

Establece una conexión SSH a la VM de TPU con el siguiente comando:

$ gcloud compute tpus tpu-vm ssh tpu-name --zone zone

Cuando accedas a la VM, el indicador de shell cambiará de username@projectname a username@vm-name:

Instala Pax en la VM de Google Cloud TPU

Instala Pax, JAX y libtpu en tu VM de TPU con los siguientes comandos:

(vm)$ python3 -m pip install -U pip \
python3 -m pip install paxml jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

Verificación del sistema

Para probar que todo esté instalado correctamente, verifica que JAX vea los núcleos de TPU:

(vm)$ python3 -c "import jax; print(jax.device_count())"

Se muestra la cantidad de núcleos de TPU, que debería ser 8 si usas un v2-8 o v3-8, o 4 si usas una v4-8.

Ejecuta código Pax en una VM de TPU

Ahora puedes ejecutar cualquier código Pax que desees. Los ejemplos de lm_cloud son un excelente punto de partida para comenzar a ejecutar modelos en Pax. Por ejemplo, el siguientes comandos entrenan un Modelo de lenguaje SPMD basado en parámetros 2B basado en un transformador en datos sintéticos.

Los siguientes comandos muestran el resultado del entrenamiento para un modelo de lenguaje SPMD. Entrena para 300 pasos en aproximadamente 20 minutos.

(vm)$ python3 .local/lib/python3.10/site-packages/paxml/main.py  --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=job_log_dir

En la porción v4-8, el resultado debe incluir lo siguiente:

Pérdidas y tiempos de pasos

tensor de resumen en el paso=step_# loss = loss
tensor de resumen en step=step_# Pasos por segundo x

Limpia

Para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos usados en este instructivo, borra el proyecto que contiene los recursos o conserva el proyecto y borra los recursos individuales.

Cuando termines de usar la VM de TPU, sigue estos pasos para limpiar los recursos.

Desconéctate de la instancia de Compute Engine, si aún no lo hiciste:

(vm)$ exit

Borra tu Cloud TPU.

$ gcloud compute tpus tpu-vm delete tpu-name  --zone zone

¿Qué sigue?

Para obtener más información sobre Cloud TPU, consulta los siguientes vínculos: