Exécuter SFT pour Gemma 4 31B sur un TPU v6e multi-hôtes

Ce tutoriel vous explique comment exécuter un affinage supervisé (SFT) sur un cluster de Tensor Processing Unit (TPU) v6e multi-hôte à l'aide de MaxText et de Cluster Toolkit. Vous utilisez Cluster Toolkit pour exécuter une charge de travail d'entraînement multi-hôte et exporter les résultats au format Hugging Face pour la diffusion.

Résumé de la charge de travail

L'exécution de cette charge de travail SFT sur Cloud TPU et GKE implique les composants suivants :

• Modèle : Gemma 4 31B (google/gemma-4-31b-it), un modèle avancé à pondérations ouvertes. Dans ce tutoriel, les pondérations du modèle sont téléchargées depuis Hugging Face, converties dans un format compatible avec MaxText et déployées sur un cluster GKE avec une infrastructure Cloud TPU v6e (v6e-32).
• Méthode d'entraînement : SFT à l'aide de MaxText. La charge de travail exécute un pipeline de réglage évolutif sur Cloud TPU v6e à l'aide de JAX. Il optimise les paramètres du modèle sur 1 000 étapes à l'aide de l'outil d'entraînement SFT intégré de MaxText (maxtext.trainers.post_train.sft.train_sft) et des ensembles de données de réglage des instructions.
• Ensemble de données et modèle : UltraChat 200k (HuggingFaceH4/ultrachat_200k), un ensemble de données de dialogue multitour à grande échelle et de haute qualité. La charge de travail utilise cet ensemble de données pour affiner le modèle et améliorer ses capacités de conversation et de suivi des instructions.
• Infrastructure et framework : GKE, géré à l'aide de Cluster Toolkit, sur Cloud TPU v6e (v6e-32, avec une topologie 4x8). Le framework est MaxText (tpu_post_training:0.2.4) avec un backend JAX.

Objectifs

• Installez Cluster Toolkit et ses dépendances.
• Déployer un cluster Cluster Toolkit
• Convertissez un modèle Hugging Face au format MaxText.
• Exécutez une charge de travail d'entraînement SFT sur le TPU.
• Convertissez le modèle affiné au format Hugging Face pour la diffusion.

Coûts

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

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Compute Engine

Obtenez une estimation des coûts en fonction de votre utilisation prévue à l'aide du simulateur de coût.

Les nouveaux utilisateurs de Google Cloud 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-32. Ce tutoriel suppose que vous disposez déjà d'une réservation pour ces TPU dans votre projetGoogle Cloud 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 créer un compte sans frais sur Hugging Face. Une fois que vous disposez d'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) et saisissez un nom pour votre jeton.
  4. Votre jeton d'accès s'affiche. Enregistrez le jeton dans un endroit sûr.

• Sur la fiche du modèle google/gemma-4-31b-it sur Hugging Face, acceptez le contrat de licence du modèle.

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 suivre ce tutoriel :
  • Administrateur de TPU (roles/tpu.admin)
  • Utilisateur du compte de service (roles/iam.serviceAccountUser)
  • Administrateur de Compute (roles/compute.admin)
  • Administrateur de l'espace de stockage (roles/storage.admin)
  • Administrateur Kubernetes Engine (roles/container.admin)
  • Administrateur de compte de service (roles/iam.serviceAccountAdmin)
  • Administrateur de projet IAM (roles/resourcemanager.projectIamAdmin)
  • Administrateur Service Usage (roles/serviceusage.serviceUsageAdmin)

Pour en savoir plus sur l'attribution de rôles, consultez la page 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 vos variables d'environnement

Configurez vos variables d'environnement :

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export CLOUD_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/maxtext-images/tpu_post_training:0.2.4"
export TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME="gemma4-31b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"

Remplacez les éléments suivants :

• YOUR_PROJECT_ID : ID de votre projet Google Cloud.
• YOUR_REGION : région dans laquelle se trouve votre réservation de TPU et où vous déployez votre cluster.
• YOUR_ZONE : zone où se trouve votre réservation de TPU et où vous déployez votre cluster.
• YOUR_BUCKET_NAME : nom unique d'un bucket Cloud Storage.
• YOUR_RESERVATION_NAME : nom de votre réservation de TPU v6e.
• YOUR_HF_TOKEN : votre jeton d'accès Hugging Face.

Par défaut, le script définit MODEL_NAME sur gemma4-31b (le nom du modèle MaxText pour google/gemma-4-31b-it) et CLUSTER_NAME sur gke-tpu-v6e. Si vous personnalisez CLUSTER_NAME, le nom doit commencer par une lettre minuscule, ne contenir que des lettres minuscules, des chiffres et des traits d'union, et comporter 20 caractères ou moins afin que les ID de compte de service générés ne dépassent pas la limite de 30 caractères.

Installer les dépendances de Cluster Toolkit

Pour suivre ce tutoriel à partir d'un client ou d'un poste de travail Linux ou macOS, suivez les étapes pertinentes de la section Installer les dépendances de la documentation de Cluster Toolkit.

Si vous utilisez Cloud Shell, vous pouvez ignorer cette section.

Installer Cluster Toolkit

Téléchargez et extrayez le bundle prédéfini pour la version spécifique de Cluster Toolkit définie dans CLUSTER_TOOLKIT_VERSION (v1.103.0) dans votre répertoire de travail actuel :

case "$(uname -s)-$(uname -m)" in
  Linux-x86_64)              BUNDLE="gcluster_bundle_linux_amd64.tgz" ;;
  Linux-aarch64|Linux-arm64) BUNDLE="gcluster_bundle_linux_arm64.tgz" ;;
  Darwin-x86_64)             BUNDLE="gcluster_bundle_mac_amd64.tgz" ;;
  Darwin-arm64)              BUNDLE="gcluster_bundle_mac_arm64.tgz" ;;
  *) BUNDLE=""; echo "Unsupported platform: $(uname -s)-$(uname -m)" >&2 ;;
esac
mkdir -p cluster-toolkit
curl -fsSL "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_VERSION:?}/${BUNDLE:?}" \
  | tar -xz -C cluster-toolkit
cd cluster-toolkit

Remarque : La documentation Installer Cluster Toolkit explique comment télécharger la dernière version dans un répertoire cluster-toolkit. Toutefois, comme les variables de plan peuvent différer d'une version de Cluster Toolkit à l'autre, vous devez installer la version exacte spécifiée dans CLUSTER_TOOLKIT_VERSION (v1.103.0) plutôt que la dernière version pour ce tutoriel. L'extraction du bundle fournit le binaire gcluster et les plans examples/ que vous utiliserez dans les étapes suivantes. Si vous avez extrait le bundle dans un répertoire cluster-toolkit, veillez à accéder à ce répertoire (cd cluster-toolkit) avant d'exécuter les commandes suivantes.

Authentifier votre environnement

Pour vous assurer que Cluster Toolkit peut provisionner des ressources et transférer des images vers Artifact Registry en votre nom, authentifiez la Google Cloud CLI et définissez vos Identifiants par défaut de l'application :

gcloud auth login
gcloud auth application-default login --no-launch-browser
gcloud config set project "${PROJECT:?}"
gcloud auth application-default set-quota-project "${PROJECT:?}"

gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker "${REGION:?}-docker.pkg.dev" --quiet

Créer votre cluster Cluster Toolkit

Pour créer et déployer un cluster Cluster Toolkit avec 32 puces TPU v6e, procédez comme suit :

1. Créez un bucket Cloud Storage :

   gcloud storage buckets create "gs://${GCS_BUCKET:?}" --project="${PROJECT:?}" --location="${REGION:?}" || true

2. Préparez la configuration du plan et copiez l'exemple de plan de déploiement dans un répertoire de travail isolé :

   mkdir -p tmp
   cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml tmp/

3. Mettez à jour tmp/gke-tpu-v6e-advanced.yaml pour accorder au compte de service du pool de nœuds de votre cluster (gke-np-sa) le rôle d'administrateur de l'espace de stockage (roles/storage.admin) afin d'écrire des points de contrôle dans votre bucket Cloud Storage, remplacez le type de machine du pool de nœuds système par e2-standard-8 et épinglez le pool de nœuds système à votre zone cible pour éviter les ruptures de stock de types de machines régionales :

   # Grant the GKE Node Pool Service Account storage.admin access for checkpoint writes
   sed -i "s/- storage.objectViewer/- storage.admin/" tmp/gke-tpu-v6e-advanced.yaml
   # Change n2-standard-8 to e2-standard-8 to avoid GCE_STOCKOUT in some regions
   sed -i "s/n2-standard-8/e2-standard-8/" tmp/gke-tpu-v6e-advanced.yaml
   # Pin system node pool to a single zone to avoid cross-zone stockouts
   sed -i '/system_node_pool_machine_type/a \      system_node_pool_zones: [$(vars.zone)]' tmp/gke-tpu-v6e-advanced.yaml

4. Utilisez la commande gcluster deploy pour déployer votre cluster Cluster Toolkit à l'aide du plan tmp/gke-tpu-v6e-advanced.yaml et en transmettant les variables requises à l'aide de l'option --vars :

   ./gcluster deploy tmp/gke-tpu-v6e-advanced.yaml \
       --vars "project_id=${PROJECT:?},deployment_name=${CLUSTER_NAME:?},region=${REGION:?},zone=${ZONE:?},num_slices=1,tpu_topology=4x8,authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
       --download-dependencies \
       -l IGNORE \
       --auto-approve -w

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.

1. Pour simplifier les commandes suivantes, utilisez la commande gcluster job config pour configurer votre projet, votre cluster et votre emplacement par défaut :

   # Configure gcluster Defaults
   ./gcluster job config set project "${PROJECT:?}"
   ./gcluster job config set cluster "${CLUSTER_NAME:?}"
   ./gcluster job config set location "${REGION:?}"

2. Utilisez la commande gcluster job submit pour convertir le modèle du format Hugging Face au format MaxText et le stocker dans votre bucket Cloud Storage :

   ./gcluster job submit --name hf-to-mt \
       --compute-type "${TPU_TYPE:?}" \
       --num-slices 1 \
       --image "${CLOUD_IMAGE_NAME:?}" \
       --await-job-completion \
       --command "[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
         python3 -m maxtext.checkpoint_conversion.to_maxtext \
           model_name=${MODEL_NAME:?} \
           hf_access_token=${HF_TOKEN:?} \
           base_output_directory=gs://${GCS_BUCKET:?}/${MODEL_NAME}/max-text-format/ \
           scan_layers=True \
           use_multimodal=False \
           skip_jax_distributed_system=true \
           checkpoint_storage_use_zarr3=0 \
           checkpoint_storage_use_ocdbt=0 \
           hardware=cpu \
           --lazy_load_tensors=True"

   La commande inclut l'option --await-job-completion, ce qui signifie qu'elle attend dans votre terminal actuel que la conversion se termine.

3. Facultatif : Pour diffuser les journaux de conversion pendant l'exécution du job, ouvrez un terminal distinct dans le répertoire de travail où vous avez extrait gcluster, puis exécutez la commande gcluster job logs :

   # Use the list command to check status
   ./gcluster job list
   
   # Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
   ./gcluster job logs hf-to-mt --main-only -f

4. Dans votre terminal d'origine, vérifiez que les fichiers du modèle converti sont disponibles dans votre bucket Cloud Storage :

   gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/max-text-format/"

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

1. Une fois le processus de conversion terminé, démarrez la charge de travail SFT :

   ./gcluster job submit --name sft \
       --compute-type "${TPU_TYPE:?}" \
       --num-slices 1 \
       --image "${CLOUD_IMAGE_NAME:?}" \
       --await-job-completion \
       --command "JAX_PLATFORMS=tpu,cpu ENABLE_PJRT_COMPATIBILITY=true JAX_TRACEBACK_FILTERING=off LIBTPU_INIT_ARGS=' --xla_tpu_scoped_vmem_limit_kib=61440 --xla_tpu_bf16_emission_mode=NATIVE_EMISSION --xla_tpu_enable_sparse_core_collective_offload_all_reduce=true --xla_tpu_use_single_sparse_core_for_all_gather_offload=true ' \
         python3 -m maxtext.trainers.post_train.sft.train_sft \
         run_name=sft \
         base_output_directory=gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/ \
         model_name=${MODEL_NAME} \
         load_parameters_path=gs://${GCS_BUCKET}/${MODEL_NAME}/max-text-format/0/items/ \
         hf_access_token=${HF_TOKEN:?} \
         dataset_type=hf \
         hf_path=HuggingFaceH4/ultrachat_200k \
         per_device_batch_size=1 steps=1000 \
         profiler=xplane \
         checkpoint_storage_use_zarr3=0 \
         checkpoint_storage_use_ocdbt=0 \
         skip_jax_distributed_system=False"

   La commande inclut l'option --await-job-completion, ce qui signifie qu'elle attend dans votre terminal actuel que le job d'entraînement se termine.

2. Facultatif : Pour diffuser les journaux d'entraînement en direct pendant l'exécution du job, ouvrez un terminal distinct dans le répertoire de travail où vous avez extrait gcluster, puis exécutez la commande gcluster job logs :

   # Use the list command to check status
   ./gcluster job list
   
   # Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
   ./gcluster job logs sft --main-only -f

   Une fois la compilation XLA initiale terminée, le coordinateur consigne la progression de la sortie pour chaque étape de l'exécution en 1 000 étapes au format completed step: <step>, seconds: ..., TFLOP/s/device: ..., loss: ..., la perte d'entraînement diminuant d'environ 1.57 à l'étape 5 à environ 0.82 à l'étape 1 000. Avant de passer à la section suivante, attendez que l'étape 1000 se termine et que le résultat de la commande ./gcluster job list affiche l'état de sft comme Succeeded.

   Remarque : Les JobSets terminés sur GKE sont automatiquement nettoyés après l'expiration de leur valeur TTL (Time To Live) configurée, qui est d'une heure par défaut (--gke-ttl-after-finished=1h). Les charges de travail terminées ne sont pas listées dans le résultat de la commande ./gcluster job list après l'expiration de leur TTL. Toutefois, tous les points de contrôle et pondérations de modèle sont stockés dans votre bucket Cloud Storage.

3. Dans votre terminal d'origine, vérifiez que les points de contrôle d'entraînement SFT sont générés dans votre bucket Cloud Storage :

   gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/sft/checkpoints/1000/model_params/"

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

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

   ./gcluster job submit --name mt-to-hf \
       --compute-type "${TPU_TYPE:?}" \
       --num-slices 1 \
       --image "${CLOUD_IMAGE_NAME:?}" \
       --await-job-completion \
       --command "[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
         python3 -m maxtext.checkpoint_conversion.to_huggingface \
           model_name=${MODEL_NAME:?} \
           hf_access_token=${HF_TOKEN:?} \
           load_parameters_path=gs://${GCS_BUCKET:?}/${MODEL_NAME}/trained/sft/checkpoints/1000/model_params/ \
           base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/ \
           skip_jax_distributed_system=true \
           hardware=cpu \
           scan_layers=True \
           use_multimodal=False \
           weight_dtype=bfloat16 \
           --override_model_architecture"

   La commande inclut l'option --await-job-completion, ce qui signifie qu'elle attend dans votre terminal actuel que la conversion se termine.

2. Facultatif : Pour diffuser les journaux de conversion en direct pendant l'exécution du job, ouvrez un terminal distinct dans le répertoire de travail où vous avez extrait gcluster, puis exécutez la commande gcluster job logs :

   # Use the list command to check status
   ./gcluster job list
   
   # Check progress of the job (--main-only targets the coordinator pod (Job Index 0, Pod Index 0) to avoid duplicate logs from other workers)
   ./gcluster job logs mt-to-hf --main-only -f
   
   # The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/ 

3. Dans votre terminal d'origine, vérifiez que les pondérations Hugging Face Safetensors finales sont générées dans votre bucket Cloud Storage :

   gcloud storage ls -l --readable-sizes "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/hf-trained/"

Effectuer un nettoyage

Pour éviter des frais supplémentaires, utilisez la commande gcluster destroy pour supprimer les ressources créées au cours de ce tutoriel :

./gcluster destroy "${CLUSTER_NAME:?}" --robust --auto-approve
gcloud storage rm -r "gs://${GCS_BUCKET:?}"

# To delete the local deployment folder and staged blueprint
rm -rf .ghpc "${CLUSTER_NAME:?}" tmp

Étapes suivantes

• 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-32, consultez TPU v6e.