Exécuter l'affinage supervisé sur une VM TPU à l'aide de MaxText

Ce tutoriel fournit un guide par étapes pour exécuter l'affinage supervisé (SFT) sur une seule v6e-8 instance de machine virtuelle (VM) Tensor Processing Unit (TPU) su Google Cloud en utilisant MaxText, une pile d'entraînement haute performance 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 SFT sur une VM Cloud TPU implique les composants suivants :

• Modèle : Gemma 3 4B (gemma3-4b), un modèle de langage léger à poids ouverts. 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 l'exécution sur une seule VM Cloud TPU v6e-8 à l'aide de JAX.
• Méthode d'entraînement : SFT à l'aide de MaxText. La charge de travail exécute un pipeline de réglage fin sur une seule instance Cloud TPU v6e-8 à l'aide de JAX. Elle optimise les paramètres du modèle en plus de 1 000 étapes à l'aide de l'entraîneur SFT intégré de MaxText (maxtext.trainers.post_train.sft.train_sft) et d'ensembles de données de réglage des instructions axés sur le dialogue.
• Ensemble de données et modèle : UltraChat 200k (HuggingFaceH4/ultrachat_200k), un ensemble de données de dialogue multi-tours à grande échelle et de haute qualité (division train_sft avec colonne messages). Il est utilisé pour régler le modèle et améliorer ses capacités de conversation et de suivi des instructions.
• Infrastructure et framework : une seule instance de VM Cloud TPU v6e-8 (ct6e-standard-8t). Le framework est MaxText (maxtext[tpu-post-train]) alimenté par 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 d'entraînement SFT sur le TPU.
• Reconvertir le modèle réglé au format Hugging Face pour la diffusion.

Coûts

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

• Cloud TPU

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.

Important : Pour suivre ce tutoriel, vous devez avoir accès aux TPU v6e-8. Ce tutoriel suppose que vous disposez déjà d'une réservation pour ces TPU dans votre Google Cloud projet afin de vous assurer que vous disposez de la capacité disponible lorsque vous en avez besoin. Pour en savoir plus sur les options de réservation de capacité Cloud TPU pour une date et une heure ultérieures, consultez À propos des réservations Cloud TPU.

Avant de commencer

• Vous avez besoin d'un jeton d'accès Hugging Face pour suivre ce tutoriel. Vous pouvez vous inscrire pour obtenir 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 gemma3-4b.

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

• Administrateur TPU (roles/tpu.admin)
• Utilisateur du compte de service (roles/iam.serviceAccountUser)
• Éditeur Compute (roles/compute.editor)

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

Remplacez les éléments suivants :

• YOUR_PROJECT_ID: ID de votre Google Cloud projet
• YOUR_ZONE : zone que vous souhaitez utiliser
• YOUR_RESERVATION_NAME : 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 "${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}"

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

gcloud compute ssh "${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 post-entraînement.

UV_TORCH_BACKEND=cpu uv pip install "maxtext[tpu-post-train]==0.2.2" --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.

Spécifiez vos variables d'environnement, telles que votre jeton d'accès Hugging Face, le nom du modèle que vous souhaitez utiliser et le répertoire dans lequel vous souhaitez enregistrer le modèle au format 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.

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

1. Configurez les paramètres d'entraînement de la charge de travail 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. 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.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

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 POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/$STEPS/model_params

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

Une fois la conversion terminée, votre modèle réglé stocké dans /dev/shm/gemma3-4b/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 réglé vers un stockage persistant ou le importer dans le Hugging Face Hub.

Effectuer un nettoyage

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

Supprimer votre instance de VM TPU

Quittez votre instance de VM Cloud TPU, puis supprimez-la.

# 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

Étape suivante

• Pour en savoir plus sur Cloud TPU, consultez la Présentation de Cloud TPU.
• Pour en savoir plus sur l'architecture et la configuration du TPU v6e-8, consultez TPU v6e.