Exécuter un entraînement par renforcement sur une VM TPU à l'aide de MaxText

Ce tutoriel fournit un guide pas à pas pour exécuter un entraînement par apprentissage par renforcement (RL) sur une seule v6e-8 instance de machine virtuelle (VM) Tensor Processing Unit (TPU) Google Cloud à l'aide de MaxText, une pile d'entraînement hautes performances basée sur JAX pour les grands modèles de langage (LLM).

Résumé de la charge de travail

L'exécution de cette charge de travail de post-entraînement RL sur Cloud TPU implique les composants suivants :

  • Modèle : Llama 3.1 8B Instruct (llama3.1-8b-Instruct), un LLM à poids ouverts optimisé pour les tâches de dialogue et de suivi des instructions. Dans ce tutoriel, les poids du modèle sont convertis du format Hugging Face au format compatible avec MaxText (mt-format) et préparés pour une exécution distribuée sur un TPU v6e-8 à l'aide de JAX.

  • Algorithme d'apprentissage par renforcement (RL) : Group Relative Policy Optimization (GRPO), dans lequel la charge de travail exécute un pipeline de post-entraînement évolutif à l'aide de l'entraîneur RL intégré de MaxText. GRPO élimine le besoin d'un modèle de valeur ou de critique distinct en évaluant un groupe de sorties générées pour la même invite, en utilisant des récompenses relatives pour optimiser efficacement la stratégie de l'acteur.

  • Ensemble de données et modèle : Grade School Math 8K (GSM8K), un ensemble de données de haute qualité composé de problèmes mathématiques linguistiquement diversifiés pour les élèves du primaire. Cette charge de travail l'utilise pour générer des déploiements, calculer des récompenses et améliorer les capacités de raisonnement en plusieurs étapes lors de la phase d'alignement RL.

  • Infrastructure et framework : un seul nœud Cloud TPU v6e-8 (ct6e-standard-8t) utilisant MaxText (maxtext[tpu-post-train]), une pile d'entraînement hautes performances basée sur JAX.

Objectifs

  • Configurer une instance de VM Cloud TPU.
  • Installer MaxText et ses dépendances.
  • Convertir un modèle Hugging Face au format MaxText.
  • Exécuter une charge de travail RL Group Relative Policy Optimization (GRPO) sur le TPU.
  • Reconvertir le modèle entraîné au format Hugging Face pour la diffusion.

Coûts

Dans ce tutoriel, vous utilisez les composants facturables de suivants Google Cloud:

Obtenez une estimation des coûts en fonction de votre utilisation prévue, utilisez le simulateur de coût.

Les nouveaux Google Cloud utilisateurs de peuvent bénéficier d'un essai sans frais.

Une fois que vous avez terminé les tâches décrites dans ce document, supprimez les ressources que vous avez créées pour éviter que des frais vous soient facturés. Pour en savoir plus, consultez la section Effectuer un nettoyage.

Avant de commencer

  • Vous avez besoin d'un jeton d'accès Hugging Face pour suivre ce tutoriel. Vous pouvez créer un compte sans frais sur Hugging Face. Une fois que vous avez un compte, générez un jeton d'accès :

    1. Sur la page "Welcome to Hugging Face" (Bienvenue sur Hugging Face), cliquez sur l'avatar de votre compte, puis sélectionnez Access tokens (Jetons d'accès).
    2. Sur la page Access tokens (Jetons d'accès), cliquez sur Create new token (Créer un jeton).
    3. Sélectionnez le type de jeton Read (Lecture), puis saisissez un nom pour votre jeton.
    4. Votre jeton d'accès s'affiche. Enregistrez le jeton dans un endroit sûr.

  • Sur le site Web Hugging Face, acceptez le contrat de licence du modèle que vous prévoyez d'entraîner. Ce tutoriel utilise le modèle llama3.1-8b-Instruct.

Pour obtenir les autorisations nécessaires pour suivre ce tutoriel, demandez à votre administrateur de vous accorder les rôles IAM suivants sur votre projet :

Pour en savoir plus sur l'attribution de rôles, consultez Gérer l'accès aux projets, aux dossiers et aux organisations.

Vous pouvez également obtenir les autorisations requises avec des rôles personnalisés ou d'autres rôles prédéfinis.

Configurer l'environnement

Configurez vos variables d'environnement en exécutant le script suivant :

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

Remplacez les éléments suivants :

  • YOUR_PROJECT_ID: ID de votre Google Cloud projet
  • YOUR_ZONE : zone que vous souhaitez utiliser
  • YOUR_RESERVATION_NAME : votre réservation de capacité
  • YOUR_TPU_NAME : nom de votre instance de VM Cloud TPU

Authentifiez-vous auprès de Google Cloud en exécutant la commande suivante :

gcloud auth login

Créer votre VM Cloud TPU

Créez une instance de VM Cloud TPU avec huit puces TPU v6e, liée à votre réservation de capacité.

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}"

Une fois l'instance de VM créée, connectez-vous à celle-ci à l'aide de SSH.

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

Suivez les étapes ci-dessous dans votre instance de VM TPU.

Installer MaxText

Mettez à jour les packages système dans l'instance de VM TPU.

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

Installez Python 3.12, qui est requis par MaxText, et son package d'environnement virtuel.

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

Utilisez uv pour accélérer l'installation du package Python.

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

Créez un environnement virtuel nommé maxtext_venv et activez-le.

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

Installez MaxText et les dépendances requises pour les tâches de post-entraînement.

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

Installez les dépendances requises restantes en exécutant la commande suivante :

install_tpu_post_train_extra_deps

Convertir le modèle au format MaxText

Pour entraîner le modèle au format MaxText, vous devez le convertir du format Hugging Face au format MaxText.

Indiquez les valeurs suivantes :

  • Votre jeton d'accès Hugging Face
  • Nom du modèle que vous souhaitez utiliser
  • Répertoire dans lequel vous souhaitez enregistrer le modèle au format MaxText
  • Options de chargement et de stockage
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.

Remplacez YOUR_HF_TOKEN par le jeton d'accès Hugging Face que vous avez créé précédemment.

Pour convertir le modèle du format Hugging Face au format MaxText, exécutez le script suivant. Cette conversion prend environ cinq minutes.

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?}"

Démarrer la charge de travail d'entraînement

Une fois le processus de conversion terminé, vous pouvez démarrer la charge de travail RL.

  1. Configurez les paramètres d'entraînement de la charge de travail 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. Démarrez le job d'entraînement. Cette opération prend environ 10 minutes sur une instance 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

Reconvertir le modèle entraîné au format Hugging Face

Une fois la charge de travail d'entraînement terminée, reconvertissez le modèle au format Hugging Face.

  1. Définissez les chemins d'exportation et les paramètres entraînés.

    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. Exécutez la reconversion au format 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

Une fois la conversion terminée, votre modèle ajusté stocké dans /dev/shm/$MODEL_NAME/hf-trained est prêt à être utilisé. Étant donné que vous perdez l'accès au contenu du dossier /dev/shm lorsque la VM redémarre, vous devez déplacer le modèle ajusté vers un stockage persistant ou l'importer dans le hub Hugging Face.

Effectuer un nettoyage

Pour éviter des frais supplémentaires, supprimez les ressources créées lors de ce tutoriel.

Supprimer votre instance de VM TPU

Supprimez votre instance de VM Cloud TPU.

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

Étape suivante