Ejecuta el ajuste supervisado en una VM de TPU con MaxText

En este instructivo, se proporciona una guía paso a paso para ejecutar el ajuste supervisado (SFT) en una sola v6e-8 instancia de 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 SFT en una VM de Cloud TPU incluye los siguientes componentes:

• Modelo: Gemma 3 4B (gemma3-4b), un modelo de lenguaje ligero de pesos abiertos. 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 en una sola VM de Cloud TPU v6e-8 con JAX.
• Método de entrenamiento: SFT con MaxText. La carga de trabajo ejecuta una canalización de ajuste en una sola instancia de Cloud TPU v6e-8 con JAX. Optimiza los parámetros del modelo en más de 1,000 pasos con el entrenador de SFT integrado de MaxText (maxtext.trainers.post_train.sft.train_sft) y los conjuntos de datos de ajuste de instrucciones enfocados en el diálogo.
• Conjunto de datos y plantilla: UltraChat 200k (HuggingFaceH4/ultrachat_200k), un conjunto de datos de diálogo de varias turnos de alta calidad y a gran escala (train_sft dividido con la columna messages). Se usa para ajustar el modelo y mejorar sus capacidades conversacionales y de seguimiento de instrucciones.
• Infraestructura y framework: una sola instancia de VM de Cloud TPU v6e-8 (ct6e-standard-8t). El framework es MaxText (maxtext[tpu-post-train]) con tecnología de JAX.

Objetivos

• Configura una instancia de VM de Cloud TPU.
• Instala MaxText y sus dependencias.
• Convierte un modelo de Hugging Face al formato de MaxText.
• Ejecuta una carga de trabajo de entrenamiento de SFT en la TPU.
• Vuelve a convertir el modelo ajustado 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 TPU v6e-8. En este instructivo, se supone que ya tienes una reserva para estas TPU 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 gemma3-4b.

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 NAME="YOUR_TPU_NAME"
export NETWORK="default"

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 "${NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --network="${NETWORK}" \
    --tags="${NAME}" \
    --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 "${NAME}" --zone "${ZONE}" --project "${PROJECT}"

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

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 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 posteriores al entrenamiento.

UV_TORCH_BACKEND=cpu 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 de MaxText

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

Especifica tus variables de entorno, como el token de acceso de Hugging Face, el nombre del modelo que deseas usar y el directorio en el que deseas guardar el modelo en formato de MaxText.

export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME='gemma3-4b'
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 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 SFT.

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

   # -- 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 STEPS=1000
   export PER_DEVICE_BATCH_SIZE=1
   
   # -- Dataset configuration --
   export DATASET_NAME="HuggingFaceH4/ultrachat_200k"
   export TRAIN_SPLIT="train_sft"
   export TRAIN_DATA_COLUMNS="['messages']"
   
   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.sft.train_sft \
       run_name="${RUN_NAME?}" \
       base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \
       model_name="${MODEL_NAME?}" \
       load_parameters_path="${MAXTEXT_CKPT_PATH?}" \
       per_device_batch_size="${PER_DEVICE_BATCH_SIZE?}" \
       steps="${STEPS?}" \
       hf_path="${DATASET_NAME?}" \
       train_split="${TRAIN_SPLIT?}" \
       train_data_columns="${TRAIN_DATA_COLUMNS?}" \
       profiler=xplane

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 acceso para la exportación y los parámetros entrenados.

   export HF_EXPORT=/dev/shm/$MODEL_NAME/hf-trained/
   export POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/$STEPS/model_params

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

   python3 -m maxtext.checkpoint_conversion.to_huggingface \
       model_name="${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 almacenado en /dev/shm/gemma3-4b/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

Sal de la instancia de VM de Cloud TPU y, luego, bórrala.

# 1. Delete TPU instance
echo "Deleting TPU instance: ${NAME}..."
gcloud compute instances delete "${NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --quiet || true

# 2. Delete IAP firewall rule
echo "Deleting firewall rule: ${FIREWALL_RULE_NAME}..."
gcloud compute firewall-rules delete "${FIREWALL_RULE_NAME}" \
    --project="${PROJECT}" \
    --quiet || true

¿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.