Google utiliza tecnología de IA para traducir contenido a tu idioma preferido. Las traducciones realizadas con IA pueden contener errores.

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) con Google Cloud MaxText, una pila de entrenamiento de alto rendimiento basada en JAX para modelos de lenguaje grandes (LLM).

Objetivos

Configurar una instancia de VM de Cloud TPU
Instalar MaxText y sus dependencias
Convertir un modelo de Hugging Face al formato 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 Libera espacio.

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 el token.
4. Se mostrará 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 tus 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 del Google Cloud proyecto.
ZONE_NAME: Es la zona que deseas usar.
RESERVATION_NAME: Es tu reserva de capacidad.
TPU_MACHINE_NAME: Es el nombre de tu instancia de VM de Cloud TPU .

Realiza la autenticación con Google Cloud ejecutando el siguiente comando:

gcloud auth login

Crea tu VM de Cloud TPU

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

gcloud alpha compute tpus tpu-vm create $TPU_NAME \
    --zone=$ZONE \
    --project=$PROJECT \
    --accelerator-type=v6e-8 \
    --version=v2-alpha-tpuv6e \
    --provisioning-model=reservation-bound \
    --reservation=$RESERVATION

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

gcloud compute tpus tpu-vm ssh $TPU_NAME --zone $ZONE --project $PROJECT

Completa los siguientes pasos en tu instancia de TPU VM.

Instala MaxText

Actualiza los paquetes del sistema en 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 python3.12 python3.12-venv

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 venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

Instala MaxText y las dependencias que requiere para las tareas posteriores al entrenamiento.

uv pip install maxtext[tpu-post-train]==0.2.2 --resolution=lowest

Ejecuta el siguiente comando para instalar las dependencias requeridas restantes:

install_tpu_post_train_extra_deps

Convierte el modelo al formato MaxText

Para entrenar el modelo en formato MaxText, debes convertirlo del formato Hugging Face al formato 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 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 anteriormente.

Para convertir el modelo del formato Hugging Face al formato 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.

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/
export RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
export CHIPS_PER_VM=8
export NUM_BATCHES=50
export MAXTEXT_CKPT_PATH=$MODEL_CHECKPOINT_DIRECTORY/0/items

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.

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

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 alpha compute tpus tpu-vm 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.