Execute o treinamento de aprendizado por reforço multi-host para Gemma 4 26B no TPU v6e

Este tutorial mostra como executar o treinamento de aprendizado por reforço (RL, na sigla em inglês) de vários hosts em um cluster de Unidade de Processamento de Tensor (TPU) v6e-64 usando o MaxText e o Cluster Toolkit. Você usa Cluster Toolkit para executar uma carga de trabalho de treinamento de vários hosts e exportar os resultados de volta para o formato Hugging Face para veiculação.

Objetivos

• Instalar o Cluster Toolkit e as dependências dele.
• Instalar o MaxText e as dependências dele.
• Implantar um cluster do Cluster Toolkit.
• Converter um modelo do Hugging Face para o formato MaxText.
• Executar uma carga de trabalho de treinamento de RL no cluster da TPU v6e.
• Converter o modelo ajustado de volta para o formato Hugging Face para veiculação.

Custos

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

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

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

Importante: para concluir este tutorial, você precisa ter acesso às TPUs v6e-64. Este tutorial pressupõe que você já tenha uma reserva para essas TPUs no seu Google Cloud projeto para garantir que tenha capacidade disponível quando você 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. É possível se inscrever para uma conta sem custo financeiro no Hugging Face. Depois de ter uma conta, gere um token de acesso:

1. Na página "Bem-vindo ao Hugging Face", clique no avatar da sua conta e selecione Tokens de acesso.
2. Na página Tokens de acesso, clique em Criar novo token.
3. Selecione o tipo de token Ler e insira um nome para o token.
4. Seu token de acesso será exibido. Salve o token em um local seguro.

• No site do Hugging Face, aceite o contrato de licença do modelo que você planeja treinar. Este tutorial usa o modelo gemma4-26b.

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

• Para concluir este tutorial:
  • Administrador da TPU (roles/tpu.admin)
  • Usuário da conta de serviço (roles/iam.serviceAccountUser)
  • Editor do Compute (roles/compute.editor)
• Para preparar a imagem do contêiner MaxText:
  • Editor do Cloud Build (roles/cloudbuild.builds.editor)
  • Administrador do Artifact Registry (roles/artifactregistry.admin)
  • Administrador de armazenamento (roles/storage.admin)
  • 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 com 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 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"

Substitua:

• YOUR_PROJECT_ID: o ID do Google Cloud projeto.
• YOUR_REGION: a região em que você quer implantar o cluster.
• YOUR_ZONE: a zona em que você quer implantar o cluster.
• YOUR_CLUSTER_NAME: o nome do cluster do Google Kubernetes Engine.
• YOUR_REPOSITORY_NAME: o nome do repositório do Artifact Registry para as imagens do MaxText.
• YOUR_BUCKET_NAME: um nome globalmente exclusivo para um bucket do Cloud Storage.
• YOUR_RESERVATION_NAME: o nome da sua reserva.
• YOUR_HF_TOKEN: seu token de acesso do Hugging Face.

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

Instale o pacote pré-criado do Cluster Toolkit seguindo as instruções em Instalar o Cluster Toolkit.

Preparar a imagem do contêiner MaxText

Para preparar a imagem do contêiner 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 ${REPOSITORY_NAME} \
       --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 seguinte conteúdo:

   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'

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

Criar o cluster do Cluster Toolkit

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

1. Crie um papel personalizado do Identity and Access Management (IAM), chamado 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. Crie um bucket do Cloud Storage:

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

3. Por padrão, a conta de serviço do pool de nós do cluster não tem as permissões necessárias para gravar no bucket do Cloud Storage. Para permitir que a conta de serviço do pool de nós grave no bucket do Cloud Storage, conceda a ela o papel Storage Admin. Para conceder esse papel, edite o arquivo gke-tpu-v6e-advanced.yaml atualizando o service-account módulo chamado 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. Aplique configurações personalizadas ao bloco gke-tpu-v6e-cluster que substituem as configurações padrão de IPv6 e tipo de máquina:

   - 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

5. Implante o cluster do Cluster Toolkit usando o blueprint gke-tpu-v6e-advanced.yaml e transmitindo as variáveis necessárias usando a 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=${CLUSTER_NODEPOOL_COUNT},tpu_topology=${TOPOLOGY},authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
       -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. Depois de terminar de criar o cluster do Cluster Toolkit, configure o Docker:

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

2. Para simplificar os comandos subsequentes, configure 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}

3. Converta o modelo do formato Hugging Face para o formato MaxText e armazene-o no bucket do 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"

4. Verifique o status do job de conversão:

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

Iniciar a carga de trabalho de treinamento

Depois que o processo de conversão for concluído, inicie a carga de trabalho de treinamento de 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'}}\""

Verifique o status do job de treinamento:

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

Converter o modelo treinado de volta para o formato Hugging Face

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="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"

Verifique o status do job de conversão:

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

Limpar

Para evitar cobranças adicionais, exclua os recursos criados durante este 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}

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-64 TPU, consulte TPU v6e.
• Para mais informações sobre o Cluster Toolkit, consulte Visão geral do Cluster Toolkit.