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 fino supervisado (SFT) en una sola instancia de máquina virtual (VM) de unidad de procesamiento tensorial (TPU) en Google Cloud con MaxText, una pila de entrenamiento de alto rendimiento basada en JAX para modelos de lenguaje grandes (LLM).v6e-8

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.
• Convierte el modelo ajustado nuevamente 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 v6e-8 TPU. En este instructivo, se supone que ya tienes una reserva para estas TPUs en tu proyecto deGoogle Cloud para asegurarte de tener la capacidad disponible cuando la necesites. Para obtener más información sobre las 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 que tengas una cuenta, genera un token de acceso:

  1. En la página Welcome to Hugging Face, haz clic en el avatar de tu cuenta y selecciona Access tokens.
  2. En la página Tokens de acceso, haz clic en Crear token nuevo.
  3. Selecciona el tipo de token Read 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 del 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 a través de 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 NAME="YOUR_TPU_NAME"
export NETWORK="default"

Reemplaza lo siguiente:

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

Para autenticarte con Google Cloud , ejecuta 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, vinculada 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 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 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 restantes requeridas:

install_tpu_post_train_extra_deps

Convierte el modelo al formato de MaxText

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

Especifica tus variables de entorno, como tu 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 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 previamente.

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. Este proceso tarda alrededor de 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

Cómo volver 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 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. Como 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.

Realiza una limpieza

Para evitar que se generen cargos adicionales, borra los recursos que creaste 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.