Exécuter l'affinage supervisé sur Gemma 4 sur TPU v6e

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

Objectifs

  • Installer Cluster Toolkit et ses dépendances.
  • Installer MaxText et ses dépendances.
  • Déployer un cluster Cluster Toolkit.
  • 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 affiné au format Hugging Face pour la mise en service.

Coûts

Dans ce document, vous utilisez les composants facturables suivants de 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 Libérer de l'espace.

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 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 gemma4-31b.

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 via des rôles personnalisés ou d'autres rôles prédéfinis.

Configurer vos variables d'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="gke-tpu-v6e"
export REPOSITORY_NAME="YOUR_REPOSITORY_NAME"
export CLOUD_IMAGE_NAME="${REGION}-docker.pkg.dev/${PROJECT}/${REPOSITORY_NAME}/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-31b"
export HF_TOKEN="YOUR_HF_TOKEN"
export GCS_BUCKET="YOUR_BUCKET_NAME"

Remplacez les éléments suivants :

  • YOUR_PROJECT_ID : ID de votre Google Cloud projet.
  • YOUR_REGION : région dans laquelle vous souhaitez déployer votre cluster.
  • YOUR_ZONE : zone dans laquelle vous souhaitez déployer votre cluster.
  • YOUR_REPOSITORY_NAME : nom du dépôt Artifact Registry pour vos images MaxText.
  • YOUR_RESERVATION_NAME : nom de votre réservation.
  • YOUR_HF_TOKEN : jeton d'accès Hugging Face.
  • YOUR_BUCKET_NAME : nom unique d'un bucket Cloud Storage.

Installer les dépendances de Cluster Toolkit

Pour suivre ce tutoriel à partir d'un client ou d'une station de travail Linux ou macOS, suivez les étapes appropriées dans Installer les dépendances dans la documentation de Cluster Toolkit.

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

Installer Cluster Toolkit

Installez le bundle prédéfini pour Cluster Toolkit en suivant les instructions de la section Installer Cluster Toolkit.

Préparer votre image de conteneur MaxText

Pour préparer votre image de conteneur MaxText, y compris l'installation des 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 ${REPOSITORY_NAME} \
    --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.3 --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 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 rôle Identity and Access Management (IAM) personnalisé nommé gke.gcsfuse.profileUser :

    # The GKE TPU v6e blueprint uses GCS Fuse CSI Storage Profiles which requires a custom IAM role.
# If this role is not already created in your project, you must create it before deploying.
gcloud iam roles create gke.gcsfuse.profileUser \
  --project=${PROJECT} \
  --title="GKE GCSFuse Profile User" \
  --description="Allows scanning GCS buckets for objects, retrieving bucket metadata, and creating Anywhere Caches." \
  --permissions="storage.objects.list,storage.buckets.get,storage.anywhereCaches.create,storage.anywhereCaches.get,storage.anywhereCaches.list,storage.anywhereCaches.update"

  2. Créez un bucket Cloud Storage :

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

  3. Par défaut, le compte de service du pool de nœuds de votre cluster ne dispose pas des autorisations requises pour écrire dans votre bucket Cloud Storage. Pour autoriser le compte de service du pool de nœuds à écrire dans votre bucket Cloud Storage, vous devez lui attribuer le rôle Storage Admin. Pour attribuer ce rôle, modifiez le fichier gke-tpu-v6e-advanced.yaml en mettant à jour le module node_pool_service_account :

    - id: node_pool_service_account
  source: modules/project/service-account
  settings:
    name: gke-np-sa
    project_roles:
    - logging.logWriter
    - monitoring.metricWriter
    - monitoring.viewer
    - stackdriver.resourceMetadata.writer
    - storage.admin            # Change from storage.objectViewer
    - artifactregistry.reader

  4. Déployez votre cluster Cluster Toolkit à l'aide du blueprint gke-tpu-v6e-advanced.yaml et transmettez les variables requises à l'aide de l'option --vars :

    ./gcluster deploy examples/gke-tpu-v6e/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 \
    -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. Une fois que vous avez terminé de créer votre cluster Cluster Toolkit, configurez Docker :

    # Configure docker for pulling images
gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker ${REGION}-docker.pkg.dev --quiet

  2. Pour simplifier les commandes suivantes, configurez 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}

  3. Pour convertir le modèle du format Hugging Face au format MaxText et le stocker dans votre bucket Cloud Storage, exécutez le script suivant :

    ./gcluster job submit --name hf-to-mt \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --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"

Pour vérifier l'état de la tâche de conversion, exécutez la commande suivante :

# Use the list command to check status
./gcluster job list \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# 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 \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

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 en exécutant la commande suivante :

./gcluster job submit --name sft \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --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"

Pour vérifier l'état de la tâche d'entraînement, exécutez la commande suivante :

# Use the list command to check status
./gcluster job list \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# 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 \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

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 :

./gcluster job submit --name mt-to-hf \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --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"

Pour vérifier l'état de la tâche de conversion, exécutez la commande suivante :

# Use the list command to check status
./gcluster job list \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# 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 \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/ - though again, it's ~2x the size of the original...

Libérer de l'espace

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

gcluster destroy ${CLUSTER_NAME} --robust
gcloud storage rm -r gs://${GCS_BUCKET}
gcloud artifacts repositories delete ${REPOSITORY_NAME} --location=${REGION} --project=${PROJECT} --quiet

# To delete the local deployment folder
rm -rf .ghpc ${CLUSTER_NAME}

Étape suivante