Esegui il fine-tuning supervisionato su Gemma 4 su TPU v6e

Questo tutorial mostra come eseguire il fine-tuning supervisionato (SFT) su un cluster di Tensor Processing Unit (TPU) v6e utilizzando MaxText e Cluster Toolkit. Utilizzerai Cluster Toolkit per eseguire un carico di lavoro di addestramento multi-host ed esportare i risultati nel formato Hugging Face per l'erogazione.

Obiettivi

  • Installare Cluster Toolkit e le relative dipendenze.
  • Installare MaxText e le relative dipendenze.
  • Eseguire il deployment di un cluster Cluster Toolkit.
  • Convertire un modello Hugging Face nel formato MaxText.
  • Eseguire un carico di lavoro di addestramento SFT sulla TPU.
  • Convertire di nuovo il modello ottimizzato nel formato Hugging Face per l'erogazione.

Costi

In questo documento vengono utilizzati i seguenti componenti fatturabili di Google Cloud:

Per generare una stima dei costi in base all'utilizzo previsto, utilizza il calcolatore prezzi.

I nuovi Google Cloud utenti potrebbero avere diritto a una prova senza costi.

Al termine delle attività descritte in questo documento, puoi evitare l'addebito di ulteriori costi eliminando le risorse che hai creato. Per saperne di più, consulta Liberare spazio.

Prima di iniziare

Per utilizzare questo tutorial, devi disporre di un token di accesso a Hugging Face. Puoi registrarti per un account senza costi su Hugging Face. Dopo aver creato un account, genera un token di accesso:

  1. Nella pagina Benvenuto in Hugging Face, fai clic sull'avatar del tuo account e seleziona Token di accesso.
  2. Nella pagina Token di accesso, fai clic su Crea nuovo token.
  3. Seleziona il tipo di token Lettura e inserisci un nome per il token.
  4. Viene visualizzato il token di accesso. Salva il token in un luogo sicuro.
  • Sul sito web di Hugging Face, accetta il contratto di licenza per il modello che intendi addestrare. Questo tutorial utilizza il modello gemma4-31b.

Per ottenere le autorizzazioni necessarie per completare questo tutorial, chiedi all'amministratore di concederti i seguenti ruoli IAM nel tuo progetto:

Per saperne di più sulla concessione dei ruoli, consulta Gestisci l'accesso a progetti, cartelle e organizzazioni.

Potresti anche riuscire a ottenere le autorizzazioni richieste tramite i ruoli personalizzati o altri ruoli predefiniti.

Impostare le variabili di ambiente

Imposta le variabili di ambiente eseguendo il seguente script:

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"

Sostituisci quanto segue:

  • YOUR_PROJECT_ID: l'ID del tuo Google Cloud progetto.
  • YOUR_REGION: la regione in cui vuoi eseguire il deployment del tuo cluster.
  • YOUR_ZONE: la zona in cui vuoi eseguire il deployment del cluster.
  • YOUR_REPOSITORY_NAME: il nome del repository Artifact Registry per le immagini MaxText.
  • YOUR_RESERVATION_NAME: il nome della prenotazione.
  • YOUR_HF_TOKEN: il token di accesso a Hugging Face.
  • YOUR_BUCKET_NAME: un nome univoco a livello globale per un bucket Cloud Storage.

Installare le dipendenze di Cluster Toolkit

Per completare questo tutorial da un client o una workstation Linux o macOS, segui i passaggi pertinenti in Installare le dipendenze nella documentazione di Cluster Toolkit.

Se utilizzi Cloud Shell, puoi saltare questa sezione.

Installare Cluster Toolkit

Installa il pacchetto predefinito per Cluster Toolkit seguendo le istruzioni riportate in Installare Cluster Toolkit.

Preparare l'immagine container MaxText

Per preparare l'immagine container MaxText, inclusa l'installazione delle dipendenze richieste, completa i seguenti passaggi:

  1. Crea un bucket Cloud Storage:

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

  2. Crea un repository 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. Crea un file nella directory principale del repository con il nome file cloudbuild.yaml e il seguente contenuto:

    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. Utilizza Cloud Build per creare l'immagine Docker MaxText:

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

Creare il cluster Cluster Toolkit

Per creare ed eseguire il deployment di un cluster Cluster Toolkit con 32 chip TPU v6e, completa i seguenti passaggi:

  1. Crea un ruolo IAM (Identity and Access Management) personalizzato, denominato 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. Crea un bucket Cloud Storage:

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

  3. Per impostazione predefinita, il account di servizio del pool di nodi del cluster non dispone delle autorizzazioni necessarie per scrivere nel bucket Cloud Storage. Per consentire al account di servizio del pool di nodi di scrivere nel bucket Cloud Storage, devi concedergli il ruolo Storage Admin. Per concedere questo ruolo, modifica il file gke-tpu-v6e-advanced.yaml aggiornando il modulo 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. Esegui il deployment del cluster Cluster Toolkit utilizzando il progetto iniziale gke-tpu-v6e-advanced.yaml e passando le variabili richieste utilizzando il flag --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

Convertire il modello nel formato MaxText

Per addestrare il modello in formato MaxText, devi convertirlo dal formato Hugging Face al formato MaxText.

  1. Dopo aver creato il cluster Cluster Toolkit, configura Docker:

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

  2. Per semplificare i comandi successivi, configura il progetto, il cluster e la località predefiniti:

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

  3. Per convertire il modello dal formato Hugging Face al formato MaxText e archiviarlo nel bucket Cloud Storage, esegui il seguente script:

    ./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"

Per controllare lo stato del job di conversione, esegui il comando seguente:

# 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}

Avviare il carico di lavoro di addestramento

Al termine del processo di conversione, puoi avviare il carico di lavoro SFT eseguendo il comando seguente:

./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"

Per controllare lo stato del job di addestramento, esegui il comando seguente:

# 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}

Convertire di nuovo il modello addestrato nel formato Hugging Face

Al termine del carico di lavoro di addestramento, converti di nuovo il modello nel formato 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"

Per controllare lo stato del job di conversione, esegui il comando seguente:

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

Libera spazio

Per evitare addebiti aggiuntivi, elimina le risorse create durante questo tutorial.

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}

Passaggi successivi

  • Per saperne di più su Cloud TPU, consulta Introduzione a Cloud TPU.
  • Per i dettagli sull'architettura e sulla configurazione della v6e-32 TPU, consulta TPU v6e.