Google utilise la technologie IA pour traduire le contenu dans votre langue préférée. Les traductions générées par IA peuvent contenir des erreurs.

Exécuter un entraînement d'apprentissage par renforcement multi-hôte pour Gemma 4 26B sur TPU v6e

Ce tutoriel vous explique comment exécuter un entraînement d'apprentissage par renforcement (RL) multihôte sur un cluster Tensor Processing Unit (TPU) v6e-64 à 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 diffusion.

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 RL sur le cluster TPU v6e.
Reconvertir le modèle affiné au format Hugging Face pour la diffusion.

Coûts

Dans ce document, 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 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 :

Sur la page Bienvenue sur Hugging Face, cliquez sur l'avatar de votre compte, puis sélectionnez Access tokens (Jetons d'accès).
Sur la page Access tokens (Jetons d'accès), cliquez sur Create new token (Créer un jeton).
Sélectionnez le type de jeton Read (Lecture), puis saisissez un nom pour votre jeton.
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-26b.

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 TPU (roles/tpu.admin)
- Utilisateur du compte de service (roles/iam.serviceAccountUser)
- Éditeur Compute (roles/compute.editor)
Pour préparer votre image de conteneur MaxText :
- Éditeur Cloud Build (roles/cloudbuild.builds.editor)
- Administrateur Artifact Registry (roles/artifactregistry.admin)
- Administrateur de l'espace de stockage (roles/storage.admin)
- Administrateur Service Usage (roles/serviceusage.serviceUsageAdmin)

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 :

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export REPOSITORY_NAME="YOUR_REPOSITORY_NAME"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export CLOUD_IMAGE_NAME="${REGION}-docker.pkg.dev/${PROJECT}/${REPOSITORY_NAME}/maxtext_base:latest"
export COMPUTE_TYPE="ct6e-standard-4t"
export TPU_TYPE="v6e-64"
export TOPOLOGY="8x8"
export CLUSTER_NODEPOOL_COUNT=1
export PW_CPU_MACHINE_TYPE="c4d-standard-96"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-26b"
export HF_TOKEN="YOUR_HF_TOKEN"

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_CLUSTER_NAME : nom de votre cluster Google Kubernetes Engine.
YOUR_REPOSITORY_NAME : nom du dépôt Artifact Registry pour vos images MaxText.
YOUR_BUCKET_NAME : nom globalement unique pour un bucket Cloud Storage.
YOUR_RESERVATION_NAME : nom de votre réservation.
YOUR_HF_TOKEN : jeton d'accès Hugging Face.

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écompilé 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 installer les dépendances requises, procédez comme suit :

Créez un bucket Cloud Storage :

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

Créez 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

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.4 --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:
  machineType: 'E2_HIGHCPU_32'

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 64 puces TPU v6e, procédez comme suit :

Créez un rôle IAM (Identity and Access Management) 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"

Créez un bucket Cloud Storage :

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

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 accorder le rôle Storage Admin. Pour accorder ce rôle, modifiez le fichier gke-tpu-v6e-advanced.yaml en mettant à jour le service-account module nommé 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
```

Appliquez des paramètres personnalisés au bloc gke-tpu-v6e-cluster qui remplacent les paramètres IPv6 et de type de machine par défaut :

- id: gke-tpu-v6e-cluster
  source: modules/scheduler/gke-cluster
  use: [gke-tpu-v6e-net-0, workload_service_account]
  settings:
    enable_private_ipv6_google_access: false
    system_node_pool_disk_size_gb: $(vars.system_node_pool_disk_size_gb)
    system_node_pool_taints: []
    enable_private_endpoint: false # Allows access from authorized public IPs
    enable_pathways_for_tpus: $(vars.enable_pathways_for_tpus)
    enable_dataplane_v2: true
    configure_workload_identity_sa: true

Déployez votre cluster Cluster Toolkit à l'aide du plan 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=${CLUSTER_NODEPOOL_COUNT},tpu_topology=${TOPOLOGY},authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    -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.

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

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}

Convertissez le modèle du format Hugging Face au format MaxText et stockez-le dans votre bucket Cloud Storage :

./gcluster job submit \
  --name="gemma4-hf-to-mt" \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT}" \
  --location="${REGION}" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME}" \
  --compute-type="${COMPUTE_TYPE}" \
  --topology="${TOPOLOGY}" \
  --await-job-completion \
  --command="[ \"\$JOB_COMPLETION_INDEX\" != \"0\" ] || \
    python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name=${MODEL_NAME} \
    hf_access_token=${HF_TOKEN} \
    --hf_model_path='google/gemma-4-26b-a4b-it' \
    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"

Vérifiez l'état de la tâche de conversion :

# 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 gemma4-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é, démarrez la charge de travail d'entraînement RL :

./gcluster job submit \
  --name="gemma4-training" \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT}" \
  --location="${REGION}" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME}" \
  --compute-type="${COMPUTE_TYPE}" \
  --topology="${TOPOLOGY}" \
  --pathways \
  --pathways-gcs-location="gs://${GCS_BUCKET}/pathways/" \
  --env="GRPC_DNS_RESOLVER=native" \
  --pathways-proxy-env="GRPC_DNS_RESOLVER=native" \
  --pathways-server-env="GRPC_DNS_RESOLVER=native" \
  --pathways-worker-env="GRPC_DNS_RESOLVER=native" \
  --command="export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}'); \
      python3 -c \"import pathlib, tpu_inference.layers.common.fused_moe_gmm as f; p = pathlib.Path(f.__file__); p.write_text(p.read_text().replace('onehot_moe_permute_threshold: int = 0,', 'onehot_moe_permute_threshold: int = 100000,'))\"; \
      JAX_PLATFORMS=proxy,cpu ENABLE_PATHWAYS_PERSISTENCE=1 \
      python3 -m maxtext.trainers.post_train.rl.train_rl \
      run_name=rl \
      base_output_directory=gs://${GCS_BUCKET}/${MODEL_NAME}/trained/ \
      model_name=${MODEL_NAME} \
      scan_layers=False \
      load_parameters_path=gs://${GCS_BUCKET}/${MODEL_NAME}/max-text-format/0/items/ \
      hf_access_token=${HF_TOKEN} \
      num_batches=50 \
      batch_size=8 \
      rollout_tensor_parallelism=2 \
      rollout_expert_parallelism=4 \
      trainer_devices_fraction=0.5 \
      sampler_devices_fraction=0.5 \
      tokenizer_path='google/gemma-4-26b-a4b-it' \
      ici_tensor_parallelism=2 \
      ici_expert_parallelism=4 \
      hbm_utilization_vllm=0.55 \
      remat_policy=full \
      async_scheduling=False \
      allow_split_physical_axes=true \
      ragged_gather_reduce_fallback=True \
      vllm_hf_overrides='{architectures: [\"MaxTextForCausalLM\"]}' \
      vllm_additional_config=\"{'maxtext_config': {'model_name': '${MODEL_NAME}', 'allow_split_physical_axes': 'true', 'use_ragged_sort': 'false', 'ragged_gather_reduce_fallback': 'true', 'prefuse_moe_weights': 'true', 'weight_dtype': 'bfloat16'}}\""

Vérifiez l'état de la tâche d'entraînement :

# 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 gemma4-training --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="gemma4-mt-to-hf" \
  --cluster="${CLUSTER_NAME}" \
  --project="${PROJECT}" \
  --location="${REGION}" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME}" \
  --compute-type="${COMPUTE_TYPE}" \
  --topology="${TOPOLOGY}" \
  --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/rl/checkpoints/actor/50/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"