Exécuter l'affinage supervisé multi-hôtes sur le modèle Qwen3-14b à l'aide de MaxText

Ce tutoriel fournit un guide par étapes pour exécuter un affinage supervisé (SFT) sur le modèle Qwen3-14b à l'aide de MaxText sur Cloud TPU. Vous apprendrez à créer une image de conteneur spécialisée, à provisionner un cluster Google Kubernetes Engine (GKE) avec Pathways à l'aide d'Accelerated Processing Kit (XPK) et à exécuter une charge de travail d'entraînement multi-hôte.

Objectifs

• Apprendre à créer une image de conteneur MaxText personnalisée optimisée pour le post-entraînement.
• Provisionner un cluster GKE à l'aide de XPK avec Pathways activé.
• Convertir le modèle Qwen3 14b du format Hugging Face au format MaxText.
• Exécuter une charge de travail d'entraînement SFT multi-hôte sur Cloud TPU.
• Reconvertir le modèle réglé au format Hugging Face pour la mise en service.

Coûts

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

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Artifact Registry
• Compute Engine, for the temporary build virtual machine (VM)

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

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

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

• Vérifiez que votre compte utilisateur ou votre compte de service dispose des rôles suivants :
  • roles/compute.admin, pour créer la VM de compilation
  • roles/artifactregistry.admin, pour gérer le dépôt Docker
  • roles/storage.admin, pour gérer le bucket de données
  • roles/container.admin, pour créer et gérer le cluster Google Kubernetes Engine
  • roles/iam.serviceAccountAdmin, pour créer le compte de service de la charge de travail
  • roles/resourcemanager.projectIamAdmin, pour définir des stratégies Identity and Access Management (IAM)
  • roles/iam.serviceAccountUser, pour agir en tant que compte de service
• Installez et initialisez Google Cloud CLI.
• Vérifiez que vous avez installé Python 3.12 ou une version ultérieure sur votre station de travail.

• 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 dans un endroit sûr.

Configurer l'environnement

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

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export GCS_BUCKET="YOUR_GCS_BUCKET"
export CLOUD_IMAGE_NAME="$REGION-docker.pkg.dev/$PROJECT/maxtext-images/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export CLUSTER_NODEPOOL_COUNT=1
export PW_CPU_MACHINE_TYPE="c4d-standard-96"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="qwen3-14b"
export HF_TOKEN="YOUR_HF_TOKEN"

Remplacez les éléments suivants :

• YOUR_PROJECT_ID: ID de votre Google Cloud projet
• YOUR_REGION : région que vous souhaitez utiliser
• YOUR_ZONE : zone que vous souhaitez utiliser
• YOUR_CLUSTER_NAME : nom de votre cluster Google Kubernetes Engine
• YOUR_GCS_BUCKET : nom unique de votre bucket Cloud Storage
• YOUR_RESERVATION_NAME : réservation de capacité
• YOUR_HF_TOKEN : jeton d'accès Hugging Face

Préparer votre image de conteneur MaxText

Pour préparer votre image de conteneur MaxText, y compris installer les dépendances requises, procédez comme suit :

1. Créez un bucket Cloud Storage :

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

2. Créer un dépôt Artifact Registry :

   gcloud artifacts repositories create maxtext-images \
       --repository-format=docker \
       --location=$REGION \
       --project=$PROJECT \
       --description="Docker repository for MaxText images in $REGION" || true

3. Créez un fichier dans le répertoire racine de votre dépôt avec le nom de fichier cloudbuild.yaml et le contenu suivant :

   steps:
     - name: 'gcr.io/cloud-builders/docker'
       entrypoint: 'bash'
       args:
         - '-c'
         - |
           set -euo pipefail
   
           # 0. Install prerequisites (if needed)
           apt-get update && apt-get install -y curl || apk add curl || true
   
           # 1. Install uv
           curl -LsSf https://astral.sh/uv/install.sh | sh
           source $$HOME/.local/bin/env
   
           # 2. Setup Python environment and install MaxText runner
           uv venv --python 3.12 --seed maxtext_venv
           source maxtext_venv/bin/activate
           uv pip install maxtext[runner]==0.2.1 --resolution=lowest
   
           # 3. Build the Docker image (Cloud Build has Docker pre-configured)
           build_maxtext_docker_image WORKFLOW=post-training
   
           # 4. Tag the image properly
           docker tag maxtext_base_image ${_CLOUD_IMAGE_NAME}
   
   # Cloud Build automatically pushes images listed here
   images:
     - '${_CLOUD_IMAGE_NAME}'
   
   options:
     # We use a high-CPU machine to match the n4-standard-16 from the VM tutorial
     machineType: 'E2_HIGHCPU_32'

4. Utilisez Cloud Build pour créer votre image Docker MaxText :

   gcloud builds submit . \
       --project=$PROJECT \
       --region=$REGION \
       --substitutions=_CLOUD_IMAGE_NAME="${CLOUD_IMAGE_NAME}"

Créer votre cluster Google Kubernetes Engine

Pour exécuter l'entraînement SFT sur le modèle Qwen3 14b, vous avez besoin d'un cluster Google Kubernetes Engine équipé de puces TPU. Installez Accelerated Processing Kit (XPK) et créez un cluster GKE compatible avec Pathways.

# Start with creating a new virtual environment to install XPK in.
VENV_DIR=venvp3
python3 -m venv $VENV_DIR
source $VENV_DIR/bin/activate
pip install xpk==1.14.0

xpk cluster create-pathways \
  --num-slices=${CLUSTER_NODEPOOL_COUNT} \
  --tpu-type=${TPU_TYPE} \
  --pathways-gce-machine-type=${PW_CPU_MACHINE_TYPE} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --cluster=${CLUSTER_NAME} \
  --custom-cluster-arguments="--enable-ip-alias" \
  --reservation=$RESERVATION \
  --default-pool-cpu-machine-type=n4-standard-16

gcloud container clusters get-credentials $CLUSTER_NAME \
  --location=$REGION \
  --project $PROJECT

Préparer le modèle pour l'entraînement

Convertissez le modèle de base au format MaxText à l'aide d'une charge de travail basée sur le processeur. N'exécutez pas cette tâche sur plusieurs machines en parallèle. La commande suivante inclut une vérification pour s'assurer que la conversion ne s'exécute que sur un seul nœud TPU.

xpk workload create \
  --workload "qwen-hf-to-mt" \
  --docker-image $CLOUD_IMAGE_NAME \
  --cluster ${CLUSTER_NAME} \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --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}/qwen-3-14b/max-text-format/ \
  scan_layers=True \
  use_multimodal=False \
  skip_jax_distributed_system=true \
  hardware=cpu \
  --lazy_load_tensors=True"

Suivre la progression de la conversion du modèle

Pour suivre la progression de la conversion, procédez comme suit :

1. Pour afficher la liste des pods planifiés sur votre cluster GKE, exécutez la commande kubectl get pod.
2. Recherchez le pod nommé qwen-hf-to-mt-slice-job-0-0-HASH.
3. Pour inspecter la sortie du pod en temps réel, exécutez la commande kubectl logs -f POD_NAME.

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

Une fois le processus de conversion terminé, vous pouvez démarrer la charge de travail de réglage fin SFT à l'aide de XPK.

xpk workload create-pathways \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-training" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --command="JAX_PLATFORMS=proxy JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE=1 \
  python3 -m maxtext.trainers.post_train.sft.train_sft \
  run_name=sft \
  base_output_directory=gs://${GCS_BUCKET}/qwen-3-14b/trained/ \
  model_name=${MODEL_NAME} \
  load_parameters_path=gs://${GCS_BUCKET}/qwen-3-14b/max-text-format/0/items/ \
  hf_access_token=${HF_TOKEN} \
  per_device_batch_size=1 \
  steps=1000 \
  profiler=xplane \
  checkpoint_storage_use_zarr3=0 \
  checkpoint_storage_use_ocdbt=0 \
  enable_single_controller=True"

Surveiller la charge de travail d'entraînement

Surveillez l'état de votre charge de travail à l'aide de l'interface de ligne de commande (CLI) XPK.

xpk workload list --cluster ${CLUSTER_NAME} --project ${PROJECT} --zone ${ZONE}

Pour afficher les journaux et l'utilisation des TPU, utilisez la Google Cloud console. Vous pouvez également afficher les journaux en exécutant la commande suivante :

kubectl logs -f qwen-training-pathways-head-0-0-HASH

Remplacez HASH par le hachage numérique dans le nom de votre pod. Pour vérifier la valeur de ce hachage, exécutez la commande kubectl get pod et examinez la liste des pods renvoyée.

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

Une fois votre charge de travail d'entraînement terminée, reconvertissez les points de contrôle au format Hugging Face.

xpk workload create \
  --cluster=${CLUSTER_NAME} \
  --project=${PROJECT} \
  --zone=${ZONE} \
  --docker-image=$CLOUD_IMAGE_NAME \
  --workload="qwen-mt-to-hf" \
  --tpu-type=${TPU_TYPE} \
  --num-slices=1 \
  --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}/qwen-3-14b/trained/sft/checkpoints/1000/model_params/ \
  base_output_directory=gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ \
  skip_jax_distributed_system=true \
  hardware=cpu \
  scan_layers=True \
  use_multimodal=False \
  weight_dtype=bfloat16"

Pour suivre la progression de la conversion, exécutez la commande kubectl logs -f qwen-mt-to-hf-slice-job-0-0-HASH, en remplaçant HASH par le hachage numérique dans le nom de votre pod.

Une fois la conversion terminée, votre modèle réglé stocké dans gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ est prêt à être utilisé.

Libérer de l'espace

Pour éviter que des frais supplémentaires ne vous soient facturés, supprimez les ressources créées au cours de ce tutoriel, y compris votre cluster Google Kubernetes Engine, votre bucket Cloud Storage et votre dépôt Artifact Registry.

Attention : Si vous prévoyez de continuer à utiliser le modèle que vous avez entraîné dans ce tutoriel, enregistrez une version à laquelle vous pourrez accéder avant de supprimer le bucket Cloud Storage qui le contient.

Pour supprimer les ressources que vous avez créées pour ce tutoriel, exécutez la commande suivante :

xpk cluster delete --cluster $CLUSTER_NAME --project $PROJECT --zone $ZONE --force

gcloud storage rm --recursive gs://$GCS_BUCKET

gcloud artifacts repositories delete maxtext-images --location=$REGION --project=$PROJECT --quiet

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