Esegui il perfezionamento supervisionato multihost sul modello Qwen3-14b utilizzando MaxText

Questo tutorial fornisce una guida passo passo per l'esecuzione del fine-tuning supervisionato (SFT) sul modello Qwen3-14b utilizzando MaxText su Cloud TPU. Scoprirai come creare un'immagine container specializzata, eseguire il provisioning di un cluster Google Kubernetes Engine (GKE) con Pathways utilizzando Accelerated Processing Kit (XPK) ed eseguire un workload di addestramento multi-host.

Obiettivi

• Scopri come creare un'immagine container MaxText personalizzata ottimizzata per il post-addestramento.
• Esegui il provisioning di un cluster GKE utilizzando XPK con Pathways abilitato.
• Converti il modello Qwen3 14b dal formato Hugging Face al formato MaxText.
• Esegui un workload di addestramento SFT multi-host su Cloud TPU.
• Converti 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:

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

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 Esegui la pulizia.

Prima di iniziare

• Verifica che il tuo account utente o account di servizio disponga dei seguenti ruoli:
  • roles/compute.admin, per creare la VM di build
  • roles/artifactregistry.admin, per gestire il repository Docker
  • roles/storage.admin, per gestire il bucket di dati
  • roles/container.admin, per creare e gestire il cluster Google Kubernetes Engine
  • roles/iam.serviceAccountAdmin, per creare il account di servizio del workload
  • roles/resourcemanager.projectIamAdmin, per impostare le policy di Identity and Access Management (IAM)
  • roles/iam.serviceAccountUser, per agire come account di servizio
• Installa e inizializza Google Cloud CLI.
• Verifica di aver installato Python 3.12 o versioni successive sulla workstation.

• Per utilizzare questo tutorial, devi disporre di un token di accesso a Hugging Face. Puoi registrarti per un account senza costi su Hugging Face. Una volta 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.

Configura l'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="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"

Sostituisci quanto segue:

• YOUR_PROJECT_ID: l' Google Cloud ID progetto
• YOUR_REGION: la regione che vuoi utilizzare
• YOUR_ZONE: la zona che vuoi utilizzare
• YOUR_CLUSTER_NAME: un nome per il cluster Google Kubernetes Engine
• YOUR_GCS_BUCKET: un nome univoco per il tuo bucket Cloud Storage
• YOUR_RESERVATION_NAME: la prenotazione della capacità
• YOUR_HF_TOKEN: il token di accesso a Hugging Face

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

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

Crea il cluster Google Kubernetes Engine

Per eseguire l'addestramento SFT sul modello Qwen3 14b, devi disporre di un cluster Google Kubernetes Engine dotato di chip TPU. Installa Accelerated Processing Kit (XPK) e crea un cluster GKE con il supporto di 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

Prepara il modello per l'addestramento

Converti il modello di base nel formato MaxText utilizzando un workload basato su CPU. Non eseguire questa attività su più macchine in parallelo. Il comando seguente include un controllo per assicurarsi che la conversione venga eseguita su un solo nodo 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"

Monitora l'avanzamento della conversione del modello

Per monitorare l'avanzamento della conversione:

1. Per elencare i pod pianificati nel cluster GKE, esegui il comando kubectl get pod.
2. Trova il pod denominato qwen-hf-to-mt-slice-job-0-0-HASH.
3. Per controllare l'output del pod in tempo reale, esegui il comando kubectl logs -f POD_NAME.

Avvia il workload di addestramento

Al termine del processo di conversione, puoi avviare il workload di ottimizzazione SFT utilizzando 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"

Monitora il workload di addestramento

Monitora lo stato del workload utilizzando l'interfaccia a riga di comando (CLI) XPK.

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

Per visualizzare i log e l'utilizzo delle TPU, utilizza la Google Cloud console. Puoi anche visualizzare i log eseguendo il comando seguente:

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

Sostituisci HASH con l'hash numerico nel nome del pod. Per verificare il valore di questo hash, esegui il comando kubectl get pod ed esamina l'elenco dei pod restituito.

Converti di nuovo il modello addestrato nel formato Hugging Face

Al termine del workload di addestramento, converti di nuovo i checkpoint nel formato 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"

Per monitorare l'avanzamento della conversione, esegui il comando kubectl logs -f qwen-mt-to-hf-slice-job-0-0-HASH, sostituendo HASH con l'hash numerico nel nome del pod.

Al termine della conversione, il modello ottimizzato archiviato in gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ è pronto per essere utilizzato.

Libera spazio

Per evitare addebiti aggiuntivi, elimina le risorse create durante questo tutorial, inclusi il cluster Google Kubernetes Engine, il bucket Cloud Storage e il repository Artifact Registry.

Attenzione: se prevedi di continuare a utilizzare il modello che hai addestrato in questo tutorial, salva una versione a cui puoi accedere prima di eliminare il bucket Cloud Storage che lo contiene.

Per eliminare le risorse create per questo tutorial, esegui il comando seguente:

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

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.