Executar o ajuste supervisionado de vários hosts no modelo Qwen3-14b usando o MaxText

Este tutorial oferece um guia detalhado para executar o ajuste supervisionado (SFT, na sigla em inglês) no modelo Qwen3-14b usando o MaxText no Cloud TPU. Você vai aprender a criar uma imagem de contêiner especializada, provisionar um cluster do Google Kubernetes Engine (GKE) com o Pathways usando o Accelerated Processing Kit (XPK) e executar uma carga de trabalho de treinamento de vários hosts.

Objetivos

• Aprenda a criar uma imagem de contêiner MaxText personalizada otimizada para pós-treinamento.
• Provisione um cluster do GKE usando o XPK com o Pathways ativado.
• Converta o modelo Qwen3 14b do formato Hugging Face para o formato MaxText.
• Execute uma carga de trabalho de treinamento de SFT de vários hosts no Cloud TPU.
• Converta o modelo ajustado de volta para o formato Hugging Face para veiculação.

Custos

Neste documento, você vai usar os seguintes componentes faturáveis do Google Cloud:

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

Para gerar uma estimativa de custo baseada na projeção de uso, use a calculadora de preços.

Novos Google Cloud usuários podem estar qualificados para um teste sem custo financeiro.

Ao concluir as tarefas descritas neste documento, é possível evitar o faturamento contínuo excluindo os recursos criados. Para mais informações, consulte Limpeza.

Antes de começar

• Verifique se sua conta de usuário ou de serviço tem os seguintes papéis:
  • roles/compute.admin, para criar a VM de build
  • roles/artifactregistry.admin, para gerenciar o repositório do Docker
  • roles/storage.admin, para gerenciar o bucket de dados
  • roles/container.admin, para criar e gerenciar o cluster do Google Kubernetes Engine
  • roles/iam.serviceAccountAdmin, para criar a conta de serviço da carga de trabalho
  • roles/resourcemanager.projectIamAdmin, para definir políticas do Identity and Access Management (IAM)
  • roles/iam.serviceAccountUser, para agir como a conta de serviço
• Instale e inicialize a Google Cloud CLI.
• Verifique se você instalou o Python 3.12 ou mais recente na estação de trabalho.

• Você precisa de um token de acesso do Hugging Face para usar este tutorial. É possível se inscrever em uma conta sem custo financeiro no Hugging Face. Depois de ter uma conta, gere um token de acesso:

  1. Na página Welcome to Hugging Face, clique no avatar da sua conta e selecione Access tokens.
  2. Na página Access tokens, clique em Create new token.
  3. Selecione o tipo de token Read e insira um nome para ele.
  4. Seu token de acesso será exibido. Salve o token em um local seguro.

Configure o ambiente

Configure as variáveis de ambiente executando o script a seguir:

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"

Substitua:

• YOUR_PROJECT_ID: ID do Google Cloud projeto
• YOUR_REGION: a região que você quer usar
• YOUR_ZONE: a zona que você quer usar
• YOUR_CLUSTER_NAME: um nome para o cluster do Google Kubernetes Engine
• YOUR_GCS_BUCKET: um nome exclusivo para o bucket do Cloud Storage
• YOUR_RESERVATION_NAME: sua reserva de capacidade
• YOUR_HF_TOKEN: seu token de acesso do Hugging Face

Prepare a imagem do contêiner do MaxText

Para preparar a imagem de contêiner do MaxText, incluindo a instalação das dependências necessárias, siga estas etapas:

1. Crie um bucket do Cloud Storage:

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

2. Crie um repositório do 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. Crie um arquivo no diretório raiz do repositório com o nome cloudbuild.yaml e o conteúdo a seguir:

   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. Use o Cloud Build para criar a imagem Docker do MaxText:

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

Crie o cluster do Google Kubernetes Engine

Para executar o treinamento de SFT no modelo Qwen3 14b, você precisa de um cluster do Google Kubernetes Engine equipado com chips de TPU. Instale o Accelerated Processing Kit (XPK) e crie um cluster do GKE com suporte ao 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

Prepare o modelo para treinamento

Converta o modelo base para o formato MaxText usando uma carga de trabalho baseada em CPU. Não execute essa tarefa em várias máquinas em paralelo. O comando a seguir inclui uma verificação para garantir que a conversão seja executada em apenas um nó de 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"

Acompanhe o progresso da conversão do modelo

Para acompanhar o progresso da conversão, faça o seguinte:

1. Para listar os pods programados no cluster do GKE, execute o comando kubectl get pod.
2. Encontre o pod chamado qwen-hf-to-mt-slice-job-0-0-HASH.
3. Para inspecionar a saída do pod em tempo real, execute o comando kubectl logs -f POD_NAME.

Iniciar a carga de trabalho de treinamento

Depois que o processo de conversão for concluído, você poderá iniciar a carga de trabalho de ajuste fino de SFT usando o 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"

Monitorar a carga de trabalho de treinamento

Monitore o status da carga de trabalho usando a interface de linha de comando (CLI) do XPK.

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

Para conferir os registros e a utilização da TPU, use o Google Cloud console. Também é possível conferir os registros executando o comando a seguir:

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

Substitua HASH pelo hash numérico no nome do pod. Para verificar o valor desse hash, execute o comando kubectl get pod e examine a lista de pods retornada.

Converter o modelo treinado de volta para o formato Hugging Face

Depois que a carga de trabalho de treinamento for concluída, converta os pontos de verificação de volta para o 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"

Para acompanhar o progresso da conversão, execute o comando kubectl logs -f qwen-mt-to-hf-slice-job-0-0-HASH, substituindo HASH pelo hash numérico no nome do pod.

Depois que a conversão for concluída, o modelo ajustado armazenado em gs://$GCS_BUCKET/qwen-3-14b/hf-trained/ estará pronto para uso.

Limpar

Para evitar cobranças extras, exclua os recursos criados durante este tutorial, incluindo o cluster do Google Kubernetes Engine, o bucket do Cloud Storage e o repositório do Artifact Registry.

**Cuidado**: se você planeja continuar usando o modelo treinado neste tutorial, salve uma versão a que você possa acessar antes de excluir o bucket do Cloud Storage que o contém.

Para excluir os recursos criados para este tutorial, execute o comando a seguir:

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

A seguir

• Para mais informações sobre o Cloud TPU, consulte Introdução ao Cloud TPU.
• Para detalhes de arquitetura e configuração da v6e-32 TPU, consulte TPU v6e.