Ejecuta el entrenamiento de aprendizaje por refuerzo en una VM de TPU con MaxText

En este instructivo, se proporciona una guía paso a paso para ejecutar el entrenamiento de aprendizaje por refuerzo (RL) en una sola instancia dev6e-8 máquina virtual (VM) de unidad de procesamiento tensorial (TPU) Google Cloud con MaxText, una pila de entrenamiento de alto rendimiento basada en JAX para modelos de lenguaje grandes (LLM).

Resumen de la carga de trabajo

La ejecución de esta carga de trabajo de posentrenamiento de RL en Cloud TPU incluye los siguientes componentes:

• Modelo: Llama 3.1 8B Instruct (llama3.1-8b-Instruct), un LLM de ponderaciones abiertas optimizado para tareas de diálogo y de seguimiento de instrucciones. En este instructivo, las ponderaciones del modelo se convierten del formato de Hugging Face al formato compatible con MaxText (mt-format) y se preparan para la ejecución distribuida en una TPU v6e-8 con JAX.

• Algoritmo de aprendizaje por refuerzo (RL): Optimización de políticas relativas de grupo (GRPO), en la que la carga de trabajo ejecuta una canalización de posentrenamiento escalable con el entrenador de RL integrado de MaxText. GRPO elimina la necesidad de un modelo de valor o crítico independiente mediante la evaluación de un grupo de resultados generados para la misma solicitud, con recompensas relativas para optimizar la política del actor de manera eficiente.

• Conjunto de datos y plantilla: Grade School Math 8K (GSM8K), un conjunto de datos de alta calidad que consta de problemas de palabras de matemáticas de la escuela primaria lingüísticamente diversos. Esta carga de trabajo lo usa para generar lanzamientos, calcular recompensas y mejorar las capacidades de razonamiento de varios pasos durante la fase de alineación de RL.

• Infraestructura y framework: Un solo nodo v6e-8 de Cloud TPU (ct6e-standard-8t) con MaxText (maxtext[tpu-post-train]), una pila de entrenamiento de alto rendimiento con tecnología de JAX.

Objetivos

• Configurar una instancia de VM de Cloud TPU
• Instalar MaxText y sus dependencias
• Convertir un modelo de Hugging Face al formato de MaxText
• Ejecutar una carga de trabajo de optimización de políticas relativas de grupo (GRPO) de RL en la TPU
• Volver a convertir el modelo entrenado al formato de Hugging Face para la entrega

Costos

En este documento, usarás los siguientes componentes facturables de Google Cloud:

• Cloud TPU

Para generar una estimación de costos en función del uso previsto, usa la calculadora de precios.

Es posible que los usuarios de Google Cloud nuevos cumplan con los requisitos para acceder a una prueba gratuita.

Cuando completes las tareas que se describen en este documento, podrás borrar los recursos que creaste para evitar que se te siga facturando. Para obtener más información, consulta Realiza una limpieza.

Importante: Para completar este instructivo, debes tener acceso a las TPUs v6e-8. En este instructivo, se supone que ya tienes una reserva para estas TPUs en tu Google Cloud proyecto para asegurarte de tener la capacidad disponible cuando la necesites. Para obtener más información sobre tus opciones para reservar capacidad de Cloud TPU para una fecha y hora futuras, consulta Acerca de las reservas de Cloud TPU.

Antes de comenzar

• Para usar este instructivo, necesitas un token de acceso de Hugging Face. Puedes registrarte para obtener una cuenta gratuita en Hugging Face. Después de tener una cuenta, genera un token de acceso:

  1. En la página de bienvenida a Hugging Face, haz clic en el avatar de tu cuenta y selecciona Tokens de acceso.
  2. En la página Tokens de acceso, haz clic en Crear token nuevo.
  3. Selecciona el tipo de token Leer y, luego, ingresa un nombre para tu token.
  4. Se muestra tu token de acceso. Guarda el token en un lugar seguro.

• En el sitio web de Hugging Face, acepta el acuerdo de licencia para el modelo que planeas entrenar. En este instructivo, se usa el modelo llama3.1-8b-Instruct.

Si quieres obtener los permisos que necesitas para completar este instructivo, pídele a tu administrador que te otorgue los siguientes roles de IAM en tu proyecto:

• Administrador de TPU (roles/tpu.admin)
• Usuario de la cuenta de servicio (roles/iam.serviceAccountUser)
• Editor de Compute (roles/compute.editor)

Para obtener más información sobre cómo otorgar roles, consulta Administra el acceso a proyectos, carpetas y organizaciones.

También puedes obtener los permisos necesarios mediante roles personalizados o cualquier otro rol predefinido.

Configura el entorno

Ejecuta la siguiente secuencia de comandos para configurar las variables de entorno:

export PROJECT="YOUR_PROJECT_ID"
export ZONE="YOUR_ZONE"
export RESERVATION="YOUR_RESERVATION_NAME"
export TPU_NAME="YOUR_TPU_NAME"

Reemplaza lo siguiente:

• YOUR_PROJECT_ID: Es el ID de tu Google Cloud proyecto.
• YOUR_ZONE: Es la zona que deseas usar.
• YOUR_RESERVATION_NAME: Es tu reserva de capacidad.
• YOUR_TPU_NAME: Es el nombre de tu instancia de VM de Cloud TPU .

Ejecuta el siguiente comando para realizar la autenticación con Google Cloud :

gcloud auth login

Crea tu VM de Cloud TPU

Crea una instancia de VM de Cloud TPU con 8 chips TPU v6e, vinculados a tu reserva de capacidad.

gcloud compute instances create "${TPU_NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --machine-type=ct6e-standard-8t \
    --image-project=ubuntu-os-accelerator-images \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --boot-disk-size=200GB \
    --maintenance-policy=TERMINATE \
    --instance-termination-action=DELETE \
    --provisioning-model=RESERVATION_BOUND \
    --reservation-affinity=specific \
    --reservation="${RESERVATION}"

Después de crear la instancia de VM, conéctate a ella con SSH.

gcloud compute ssh "${TPU_NAME}" --zone "${ZONE}" --project "${PROJECT}"

Completa los siguientes pasos en tu instancia de VM de TPU.

Instala MaxText

Actualiza los paquetes del sistema dentro de la instancia de VM de TPU.

sudo apt update && sudo apt upgrade -y --fix-missing

Instala Python 3.12, que requiere MaxText, y su paquete de entorno virtual.

sudo apt install -y build-essential cmake ninja-build

Usa uv para acelerar la instalación del paquete de Python.

curl -LsSf https://astral.sh/uv/install.sh | sh
source "${HOME}/.local/bin/env"

Crea un entorno virtual llamado maxtext_venv y actívalo.

uv python install 3.12
uv venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

Instala MaxText y las dependencias que requiere para las tareas de posentrenamiento.

UV_TORCH_BACKEND=cpu uv pip install "maxtext[tpu-post-train]==0.2.3" --resolution=lowest

Ejecuta el siguiente comando para instalar las dependencias requeridas restantes:

install_tpu_post_train_extra_deps

Convierte el modelo al formato de MaxText

Para entrenar el modelo en formato de MaxText, debes convertirlo del formato de Hugging Face al formato de MaxText.

Ingresa los siguientes valores:

• Tu token de acceso de Hugging Face
• El nombre del modelo que deseas usar
• El directorio en el que deseas guardar el modelo en formato de MaxText
• Opciones de carga y almacenamiento

export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME='llama3.1-8b-Instruct'
export MODEL_CHECKPOINT_DIRECTORY=/dev/shm/$MODEL_NAME/mt-format/
export USE_PATHWAYS=0 # Set to 1 for Pathways, 0 for McJAX
export LAZY_LOAD_TENSORS=False # True to use lazy load, False to use eager load.

Reemplaza YOUR_HF_TOKEN por el token de acceso de Hugging Face que creaste antes.

Para convertir el modelo del formato de Hugging Face al formato de MaxText, ejecuta la siguiente secuencia de comandos. Esta conversión tarda unos cinco minutos en completarse.

python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name="${MODEL_NAME?}" \
    hf_access_token="${HF_TOKEN?}" \
    base_output_directory="${MODEL_CHECKPOINT_DIRECTORY?}" \
    scan_layers=True \
    use_multimodal=False \
    hardware=cpu \
    skip_jax_distributed_system=true \
    checkpoint_storage_use_zarr3=$((1 - USE_PATHWAYS)) \
    checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS)) \
    --lazy_load_tensors="${LAZY_LOAD_TENSORS?}"

Inicia la carga de trabajo de entrenamiento

Una vez que se complete el proceso de conversión, puedes iniciar la carga de trabajo de RL.

1. Configura los parámetros de entrenamiento de la carga de trabajo de RL.

   # -- MaxText configuration --
   export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
   RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
   export RUN_NAME
   export CHIPS_PER_VM=8
   export NUM_BATCHES=50
   export MAXTEXT_CKPT_PATH=$MODEL_CHECKPOINT_DIRECTORY/0/items
   export TPU_ACCELERATOR_TYPE=v6e-8
   export TPU_WORKER_ID=0
   TPU_NAME=$(hostname)
   export TPU_NAME
   export TPU_SKIP_MDS_QUERY=1
   export TPU_WORKER_HOSTNAMES=localhost
   export TPU_TOPOLOGY=2x4
   export TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
   export TPU_HOST_BOUNDS=1,1,1

2. Inicia el trabajo de entrenamiento. Esto tarda unos 10 minutos en una instancia de VM v6e-8.

   python3 -m maxtext.trainers.post_train.rl.train_rl \
       model_name="${MODEL_NAME?}" \
       load_parameters_path="${MAXTEXT_CKPT_PATH?}" \
       run_name="${RUN_NAME?}" \
       base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \
       chips_per_vm="${CHIPS_PER_VM?}" \
       num_batches="${NUM_BATCHES?}" \
       num_test_batches=10 \
       rollout_data_parallelism=1 \
       rollout_tensor_parallelism=-1

Vuelve a convertir el modelo entrenado al formato de Hugging Face

Una vez que se complete la carga de trabajo de entrenamiento, vuelve a convertir el modelo al formato de Hugging Face.

1. Establece las rutas de exportación y los parámetros entrenados.

   export HF_EXPORT=/dev/shm/$MODEL_NAME/hf-trained/
   export HF_MODEL_NAME=llama3.1-8b
   export POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/actor/$NUM_BATCHES/model_params

2. Ejecuta la conversión de nuevo al formato de Hugging Face.

   python3 -m maxtext.checkpoint_conversion.to_huggingface \
       model_name="${HF_MODEL_NAME?}" \
       load_parameters_path="${POST_TRAIN_PATH?}" \
       base_output_directory="${HF_EXPORT?}" \
       scan_layers=True \
       use_multimodal=False \
       weight_dtype=bfloat16

Una vez que se complete la conversión, tu modelo ajustado que se almacena en /dev/shm/$MODEL_NAME/hf-trained estará listo para usarse. Debido a que pierdes el acceso al contenido de la carpeta /dev/shm cuando se reinicia la VM, debes mover el modelo ajustado al almacenamiento persistente o subirlo a Hugging Face Hub.

Limpia

Para evitar que se generen cargos adicionales, borra los recursos creados durante este instructivo.

Borra la instancia de VM de TPU

Borra la instancia de VM de Cloud TPU.

gcloud compute instances delete "${TPU_NAME}" --zone="${ZONE}" --project="${PROJECT}" --quiet

¿Qué sigue?

• Para obtener más información sobre Cloud TPU, consulta Introducción a Cloud TPU.
• Para obtener detalles sobre la arquitectura y la configuración de la TPU v6e-8, consulta TPU v6e.