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

Neste tutorial, mostramos como executar o treinamento de aprendizado por reforço (RL) de vários hosts em um cluster de Unidade de Processamento de Tensor (TPU) v6e-32 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 pós-treinamento de RL no Cloud TPU e no GKE envolve os seguintes componentes:

  • Modelo:Gemma 4 26B A4B Instruct (google/gemma-4-26b-a4b-it), um modelo avançado de pesos abertos ajustado por instruções. 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.
  • Algoritmo de aprendizado por reforço (RL):otimização de política relativa de grupo (GRPO, na sigla em inglês) usando MaxText e Pathways. A carga de trabalho executa um pipeline de pós-treinamento escalonável usando o treinador de RL integrado do MaxText (maxtext.trainers.post_train.rl.train_rl) pareado com um back-end de amostrador vLLM. O GRPO otimiza a política gerando um grupo de saídas por comando e avaliando as recompensas relativas sem precisar de um modelo crítico separado.
  • Conjunto de dados e modelo:OpenMathInstruct-2 (openmathinstruct2_rl.json), um conjunto de dados de ajuste de instruções matemáticas em grande escala. A carga de trabalho usa modelos de chat estruturados durante a fase de lançamento do RL para aprimorar os recursos de raciocínio matemático em várias etapas.
  • Infraestrutura e framework:GKE, gerenciado usando o Cluster Toolkit, na Cloud TPU v6e (fração ct6e-standard-4t, v6e-32 com uma topologia 4x8). O framework é o MaxText (tpu_post_training:0.2.4) com Pathways e 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 RL no cluster da TPU v6e.
  • 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:

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.

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.

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 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 COMPUTE_TYPE="ct6e-standard-4t"
export TPU_TYPE="v6e-32"
export TOPOLOGY="4x8"
export CLUSTER_NODEPOOL_COUNT=1
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-26b"
export CLUSTER_TOOLKIT_VERSION="v1.103.0"
export HF_TOKEN="YOUR_HF_TOKEN"

Substitua:

  • YOUR_PROJECT_ID: o 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 até 20 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.

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

Autenticar seu ambiente

Antes de executar os scripts de implantação, autentique a Google Cloud CLI e defina o Application Default Credentials (ADC). Isso garante que o Cluster Toolkit possa provisionar recursos e enviar imagens para o Artifact Registry em seu nome.

gcloud auth login
gcloud auth application-default login --no-launch-browser
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. Use o comando gcluster deploy para implantar o cluster do Cluster Toolkit usando o blueprint gke-tpu-v6e-advanced.yaml e transmitindo as variáveis necessárias com a flag --vars:

    TMP_DIR=$(mktemp -d)
BLUEPRINT_FILE="${TMP_DIR}/gke-tpu-v6e-advanced.yaml"

cp ./examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml "${BLUEPRINT_FILE}"

# Change n2-standard-8 to e2-standard-8 to avoid GCE_STOCKOUT in some regions
sed -i "s/n2-standard-8/e2-standard-8/" "${BLUEPRINT_FILE}"
# 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)]' "${BLUEPRINT_FILE}"
# Grant the GKE Node Pool Service Account storage.admin access for checkpoint writes
sed -i "s/- storage.objectViewer/- storage.admin/" "${BLUEPRINT_FILE}"

./gcluster deploy "${BLUEPRINT_FILE}" \
    --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. 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="gemma4-hf-to-mt" \
  --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=False \
    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 gemma4-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 treinamento de RL:

    ./gcluster job submit \
  --name="gemma4-training" \
  --num-slices=1 \
  --image="${CLOUD_IMAGE_NAME:?}" \
  --compute-type="${COMPUTE_TYPE:?}" \
  --topology="${TOPOLOGY:?}" \
  --pathways \
  --pathways-gcs-location="gs://${GCS_BUCKET:?}/pathways/" \
  --gke-ttl-after-finished="24h" \
  --restarts=0 \
  --env="GRPC_DNS_RESOLVER=native" \
  --env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-proxy-env="GRPC_DNS_RESOLVER=native" \
  --pathways-proxy-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-server-env="GRPC_DNS_RESOLVER=native" \
  --pathways-server-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --pathways-worker-env="GRPC_DNS_RESOLVER=native" \
  --pathways-worker-env="FLAGS_pathways_enforce_subset_devices_form_subslice=false" \
  --command="export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}') && \
    export VLLM_ENABLE_V1_MULTIPROCESSING=0 && \
    python3 -c \"import pathlib, re; p = pathlib.Path('/deps/src/maxtext/trainers/post_train/rl/utils_rl.py'); p.write_text(re.sub('optax[.]adamw[(][^)]+[)]', 'optax.adafactor(learning_rate=learning_rate)', p.read_text()))\" && \
    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:?} \
    data_template_path=maxtext/examples/chat_templates/openmathinstruct2_rl.json \
    num_batches=50 \
    batch_size=4 \
    train_micro_batch_size=2 \
    max_target_length=640 \
    max_prefill_predict_length=256 \
    mu_dtype=bfloat16 \
    grad_dtype=bfloat16 \
    rollout_tensor_parallelism=4 \
    rollout_expert_parallelism=1 \
    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 \
    ici_fsdp_parallelism=-1 \
    hbm_utilization_vllm=0.7 \
    remat_policy=full \
    async_scheduling=False \
    allow_split_physical_axes=true \
    ragged_gather_reduce_fallback=True \
    enable_dp_attention=False \
    decode_sampling_temperature=0.8 \
    decode_sampling_top_k=50 \
    decode_sampling_nucleus_p=0.95 \
    learning_rate=2e-5 \
    learning_rate_schedule_steps=100 \
    rl.num_generations=4 \
    rl.reshard_chunk_size=4 \
    debug=True \
    vllm_hf_overrides='{\"architectures\": [\"MaxTextForCausalLM\"]}' \
    vllm_additional_config=\"{'maxtext_config': {'model_name': '${MODEL_NAME}', 'model_call_mode': 'inference', 'enable_dp_attention': false, 'allow_split_physical_axes': true, 'use_ragged_sort': false, 'ragged_gather_reduce_fallback': true, 'prefuse_moe_weights': true, 'weight_dtype': 'bfloat16'}}\""

  2. Verifique o status, aguarde até que o pod coordenador esteja pronto e transmita os registros do job de treinamento:

    # Use the list command to check status (initially shows "Suspended" briefly while the CPU node pool scales up)
./gcluster job list

# Ensure kubectl is configured
gcloud container clusters get-credentials "${CLUSTER_NAME:?}" \
  --location="${REGION:?}" \
  --project="${PROJECT:?}"

# Wait for the Pathways head pod to be ready before streaming logs
kubectl wait --for=condition=Ready pod \
    -l job-name=gemma4-training-pathways-head-0 \
    --timeout=600s

# Stream training logs (runs for 50 steps; wait for "Post RL Training: ... mean_reward=..." and job status Succeeded before converting to Hugging Face)
kubectl logs -f \
    -l job-name=gemma4-training-pathways-head-0 \
    -c workload-container

    Imediatamente após o envio, a saída do comando ./gcluster job list pode mostrar brevemente Suspended enquanto o GKE aumenta o pool de nós de CPU do coordenador do Pathways (cpu-np). Quando o job começa, ele gera uma avaliação inicial de pré-treinamento na etapa 0 (Pre RL Training: ... mean_reward=...), pausa a saída de registro por vários minutos durante a compilação inicial do XLA e, em seguida, treina por 50 etapas, mostrando Reshard finished e gerações de lançamento em cada etapa sem números de etapa.

    Antes de prosseguir para a próxima etapa, aguarde até que os registros mostrem Post RL Training: ... mean_reward=... e a saída do comando ./gcluster job list mostre o status de gemma4-training como Succeeded.

  3. Verifique se os checkpoints de treinamento foram gerados no bucket do Cloud Storage:

    gcloud storage ls "gs://${GCS_BUCKET:?}/${MODEL_NAME:?}/trained/rl/checkpoints/actor/"

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="gemma4-mt-to-hf" \
  --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=False \
      use_multimodal=False \
      weight_dtype=bfloat16 \
      --override_model_architecture"

  2. Opcional: para monitorar os registros de conversão em tempo real enquanto o comando --await-job-completion anterior aguarda no primeiro terminal, abra um terminal separado no diretório de trabalho em que você extraiu gcluster e use 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 gemma4-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 do modelo treinado do Hugging Face e os arquivos de configuração estão presentes 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:?}" --auto-approve
gcloud storage rm -r "gs://${GCS_BUCKET:?}"

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

A seguir