Executar SFT para Gemma 4 31B em TPU v6e de vários hosts

Neste tutorial, mostramos como executar o ajuste supervisionado (SFT) em um cluster de Unidade de Processamento de Tensor (TPU) v6e de vários hosts usando o MaxText e o Cluster Toolkit. Use o Cluster Toolkit para executar uma carga de trabalho de treinamento de vários hosts e exportar os resultados de volta para o formato do Hugging Face para disponibilização.

Resumo da carga de trabalho

A execução dessa carga de trabalho de SFT no Cloud TPU e no GKE envolve os seguintes componentes:

• Modelo:Gemma 4 31B (google/gemma-4-31b-it), um modelo avançado de pesos abertos. Neste tutorial, os pesos do modelo são baixados do Hugging Face, convertidos em um formato compatível com o MaxText e implantados em um cluster do GKE com infraestrutura do Cloud TPU v6e (v6e-32).
• Método de treinamento:SFT usando o MaxText. A carga de trabalho executa um pipeline de ajuste fino escalonável na Cloud TPU v6e usando o JAX. Ele otimiza os parâmetros do modelo em 1.000 etapas usando o treinador de SFT integrado do MaxText (maxtext.trainers.post_train.sft.train_sft) e conjuntos de dados de ajuste de instruções.
• Conjunto de dados e modelo:UltraChat 200k (HuggingFaceH4/ultrachat_200k), um conjunto de dados de diálogo de vários turnos de alta qualidade e em grande escala. A carga de trabalho usa esse conjunto de dados para ajustar o modelo e melhorar os recursos de conversação e seguimento de instruções.
• Infraestrutura e framework:GKE, gerenciado usando o Cluster Toolkit, no Cloud TPU v6e (v6e-32, com uma topologia 4x8). O framework é o MaxText (tpu_post_training:0.2.4) com um back-end do JAX.

Objetivos

• Instale o Cluster Toolkit e as dependências dele.
• Implantar um cluster do Cluster Toolkit.
• Converter um modelo do Hugging Face para o formato MaxText.
• Execute uma carga de trabalho de treinamento de SFT na TPU.
• Converta o modelo ajustado de volta para o formato do Hugging Face para disponibilização.

Custos

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

• Google Kubernetes Engine
• Cloud TPU
• Cloud Storage
• Compute Engine

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

Novos usuários do Google Cloud 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 Limpar.

Importante: para concluir este tutorial, você precisa ter acesso às TPUs v6e-32. Este tutorial pressupõe que você já tenha uma reserva para essas TPUs no seu projeto doGoogle Cloud para garantir que tenha a capacidade disponível quando precisar. Para saber mais sobre as opções de reserva de capacidade do Cloud TPU para uma data e hora futuras, consulte Sobre as reservas do Cloud TPU.

Antes de começar

• Você precisa de um token de acesso do Hugging Face para usar este tutorial. Você pode 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 token de acesso.
  2. Na página Tokens de acesso, clique em Criar novo token.
  3. Selecione o tipo de token Read e insira um nome para ele.
  4. O token de acesso será exibido. Salve o token em um local seguro.

• No card do modelo google/gemma-4-31b-it no Hugging Face, aceite o contrato de licença do modelo.

Para conseguir as permissões necessárias para concluir o tutorial, peça ao administrador para conceder a você os seguintes papéis do IAM no seu projeto:

• Para concluir este tutorial:
  • Administrador de TPU (roles/tpu.admin)
  • Usuário da conta de serviço (roles/iam.serviceAccountUser)
  • Administrador do Compute (roles/compute.admin)
  • Administrador de armazenamento (roles/storage.admin)
  • Administrador do Kubernetes Engine (roles/container.admin)
  • Administrador da conta de serviço (roles/iam.serviceAccountAdmin)
  • Administrador de projetos do IAM (roles/resourcemanager.projectIamAdmin)
  • Administrador do Service Usage (roles/serviceusage.serviceUsageAdmin)

Para mais informações sobre a concessão de papéis, consulte Gerenciar o acesso a projetos, pastas e organizações.

Também é possível conseguir as permissões necessárias usando papéis personalizados ou outros papéis predefinidos.

Configurar as variáveis de ambiente

Configure as variáveis de ambiente:

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="YOUR_CLUSTER_NAME"
export GCS_BUCKET="YOUR_BUCKET_NAME"
export CLOUD_IMAGE_NAME="us-docker.pkg.dev/cloud-tpu-images/maxtext-images/tpu_post_training:0.2.4"
export TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME="gemma4-31b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"

Substitua:

• YOUR_PROJECT_ID: ID do seu projeto do Google Cloud.
• YOUR_REGION: a região em que a reserva de TPU está localizada e em que você implanta o cluster.
• YOUR_ZONE: a zona em que a reserva de TPU está localizada e em que você implanta o cluster.
• YOUR_CLUSTER_NAME: o nome do cluster do GKE. O nome precisa começar com uma letra minúscula, conter apenas letras minúsculas, números e hifens e ter 20 caracteres ou menos, para que os IDs de conta de serviço baseados nele permaneçam dentro do limite de 30 caracteres.
• YOUR_BUCKET_NAME: um nome globalmente exclusivo para um bucket do Cloud Storage.
• YOUR_RESERVATION_NAME: o nome da reserva de TPU v6e.
• YOUR_HF_TOKEN: seu token de acesso do Hugging Face.

Por padrão, o script define MODEL_NAME como gemma4-31b (o nome do modelo MaxText para google/gemma-4-31b-it).

Instalar dependências do Cluster Toolkit

Para concluir este tutorial em um cliente ou estação de trabalho Linux ou macOS, siga as etapas relevantes em Instalar dependências na documentação do Cluster Toolkit.

Se você estiver usando o Cloud Shell, pule esta seção.

Instalar o Cluster Toolkit

Faça o download e extraia o pacote pré-criado para a versão específica do Cluster Toolkit definida em CLUSTER_TOOLKIT_VERSION (v1.103.0) no diretório de trabalho atual:

case "$(uname -s)-$(uname -m)" in
  Linux-x86_64)              BUNDLE="gcluster_bundle_linux_amd64.tgz" ;;
  Linux-aarch64|Linux-arm64) BUNDLE="gcluster_bundle_linux_arm64.tgz" ;;
  Darwin-x86_64)             BUNDLE="gcluster_bundle_mac_amd64.tgz" ;;
  Darwin-arm64)              BUNDLE="gcluster_bundle_mac_arm64.tgz" ;;
  *) BUNDLE=""; echo "Unsupported platform: $(uname -s)-$(uname -m)" >&2 ;;
esac
mkdir -p cluster-toolkit
curl -fsSL "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_VERSION:?}/${BUNDLE:?}" \
  | tar -xz -C cluster-toolkit
cd cluster-toolkit

Observação: a documentação Instalar o Cluster Toolkit descreve como baixar a versão mais recente em um diretório cluster-toolkit. No entanto, como as variáveis de blueprint podem ser diferentes entre as versões do Cluster Toolkit, para este tutorial, é necessário instalar a versão exata especificada em CLUSTER_TOOLKIT_VERSION (v1.103.0), em vez da versão mais recente. A extração do pacote fornece o binário gcluster e os blueprints examples/ que você usa nas etapas a seguir. Se você extraiu o pacote em um diretório cluster-toolkit, mude para esse diretório (cd cluster-toolkit) antes de executar os comandos subsequentes.

Autenticar seu ambiente

Para garantir que o Cluster Toolkit possa provisionar recursos e enviar imagens para o Artifact Registry em seu nome, autentique a Google Cloud CLI e defina as Application Default Credentials:

gcloud auth login
gcloud auth application-default login --no-launch-browser
gcloud config set project "${PROJECT:?}"
gcloud auth application-default set-quota-project "${PROJECT:?}"

gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker "${REGION:?}-docker.pkg.dev" --quiet

Criar o cluster do Cluster Toolkit

Para criar e implantar um cluster do Cluster Toolkit com 32 chips de TPU v6e, siga estas etapas:

1. Crie um bucket do Cloud Storage:

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

2. Prepare a configuração do blueprint e copie o exemplo de blueprint de implantação para um diretório de trabalho isolado:

   mkdir -p tmp
   cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml tmp/

3. Atualize tmp/gke-tpu-v6e-advanced.yaml para conceder à conta de serviço do pool de nós do cluster (gke-np-sa) o papel de Administrador do Storage (roles/storage.admin) para gravar checkpoints no bucket do Cloud Storage, mude o tipo de máquina do pool de nós do sistema para e2-standard-8 e fixe o pool de nós do sistema na zona de destino para evitar falta de estoque de tipos de máquina regionais:

   # Grant the GKE Node Pool Service Account storage.admin access for checkpoint writes
   sed -i "s/- storage.objectViewer/- storage.admin/" tmp/gke-tpu-v6e-advanced.yaml
   # Change n2-standard-8 to e2-standard-8 to avoid GCE_STOCKOUT in some regions
   sed -i "s/n2-standard-8/e2-standard-8/" tmp/gke-tpu-v6e-advanced.yaml
   # Pin system node pool to a single zone to avoid cross-zone stockouts
   sed -i '/system_node_pool_machine_type/a \      system_node_pool_zones: [$(vars.zone)]' tmp/gke-tpu-v6e-advanced.yaml

4. Use o comando gcluster deploy para implantar o cluster do Cluster Toolkit usando o blueprint tmp/gke-tpu-v6e-advanced.yaml e transmitindo as variáveis necessárias com a flag --vars:

   ./gcluster deploy tmp/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 \
       -l IGNORE \
       --auto-approve -w

Converter o modelo para o formato MaxText

Para treinar o modelo no formato MaxText, é necessário convertê-lo do formato Hugging Face para o formato MaxText.

1. Para simplificar os comandos subsequentes, use o comando gcluster job config para configurar o projeto, o cluster e o local padrão:

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

2. Use o comando gcluster job submit para converter o modelo do formato Hugging Face para o formato MaxText e armazená-lo no bucket do Cloud Storage:

   ./gcluster job submit --name hf-to-mt \
       --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"

   O comando inclui a flag --await-job-completion, o que significa que ele aguarda no terminal atual até que a conversão seja concluída.

3. Opcional: para transmitir os registros de conversão enquanto o job está em execução, abra um terminal separado no diretório de trabalho em que você extraiu gcluster e execute o comando gcluster job logs:

   # Use the list command to check status
   ./gcluster job list
   
   # 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

4. No terminal original, verifique se os arquivos de modelo convertidos estão disponíveis no bucket do Cloud Storage:

   gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/max-text-format/"

Iniciar a carga de trabalho de treinamento

1. Depois que o processo de conversão for concluído, inicie a carga de trabalho de SFT:

   ./gcluster job submit --name sft \
       --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"

   O comando inclui a flag --await-job-completion, o que significa que ele aguarda no terminal atual até que o job de treinamento seja concluído.

2. Opcional: para transmitir os registros de treinamento em tempo real enquanto o job está em execução, abra um terminal separado no diretório de trabalho em que você extraiu gcluster e execute o comando gcluster job logs:

   # Use the list command to check status
   ./gcluster job list
   
   # 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

   Depois que a compilação inicial do XLA é concluída, o coordenador registra o progresso da saída para cada etapa na execução de 1.000 etapas no formato completed step: <step>, seconds: ..., TFLOP/s/device: ..., loss: ..., com a perda de treinamento diminuindo de aproximadamente 1.57 na etapa 5 para aproximadamente 0.82 na etapa 1.000. Antes de prosseguir para a próxima seção, aguarde até que a etapa 1000 seja concluída e a saída do comando ./gcluster job list mostre o status de sft como Succeeded.

   Observação :os JobSets concluídos no GKE são limpos automaticamente depois que o time to live (TTL) configurado expira, que é de uma hora por padrão (--gke-ttl-after-finished=1h). As cargas de trabalho concluídas não são listadas na saída do comando ./gcluster job list depois que o TTL expira. No entanto, todos os checkpoints e pesos do modelo são armazenados no bucket do Cloud Storage.

3. No terminal original, verifique se os checkpoints de treinamento de SFT foram gerados no bucket do Cloud Storage:

   gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/sft/checkpoints/1000/model_params/"

Converter o modelo treinado de volta para o formato do Hugging Face

1. Depois que a carga de trabalho de treinamento for concluída, converta o modelo de volta para o formato Hugging Face:

   ./gcluster job submit --name mt-to-hf \
       --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 \
           --override_model_architecture"

   O comando inclui a flag --await-job-completion, o que significa que ele aguarda no terminal atual até que a conversão seja concluída.

2. Opcional: para transmitir os registros de conversão em tempo real enquanto o job está em execução, abra um terminal separado no diretório de trabalho em que você extraiu gcluster e execute o comando gcluster job logs:

   # Use the list command to check status
   ./gcluster job list
   
   # 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
   
   # The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/ 

3. No terminal original, verifique se os pesos finais do Hugging Face Safetensors foram gerados no bucket do Cloud Storage:

   gcloud storage ls -l --readable-sizes "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/hf-trained/"

Limpar

Para evitar cobranças extras, use o comando gcluster destroy para excluir os recursos criados durante este tutorial:

./gcluster destroy "${CLUSTER_NAME:?}" --robust --auto-approve
gcloud storage rm -r "gs://${GCS_BUCKET:?}"

# To delete the local deployment folder and staged blueprint
rm -rf .ghpc "${CLUSTER_NAME:?}" tmp

A seguir

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