Executar treinamento de RL de vários hosts para Qwen3-30b-a3b na TPU v6e

Neste tutorial, mostramos como executar o treinamento de aprendizado por reforço (RL) com vários hosts em um cluster de Unidade de Processamento de Tensor (TPU) v6e-32 usando o MaxText e o Cluster Toolkit. Você usa 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 veiculação.

Objetivos

  • Instale o Cluster Toolkit e as dependências dele.
  • Implante 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 refinado de volta para o formato do Hugging Face para veiculação.

Custos

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

Para gerar uma estimativa de custo baseada na projeção de uso deste tutorial, 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 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 Leitura e digite um nome para ele.
    4. Seu token de acesso vai aparecer. 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 qwen3-30b-a3b.

Para conseguir as permissões necessárias a fim de concluir o tutorial, peça ao administrador para conceder a você os seguintes papéis do IAM no 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 TOPOLOGY="4x8"
export CLUSTER_NODEPOOL_COUNT=1
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="qwen3-30b-a3b"
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 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 (até 20 caracteres).
  • 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 no diretório de trabalho atual seguindo as instruções em Instalar o Cluster Toolkit.

Por exemplo, é possível baixar e extrair o pacote para o diretório de trabalho atual da seguinte forma:

wget -qO- "https://github.com/GoogleCloudPlatform/cluster-toolkit/releases/download/${CLUSTER_TOOLKIT_VERSION:-v1.103.0}/gcluster_bundle_linux_amd64.tgz" | tar -xz

Extrair o pacote para o diretório de trabalho atual fornece o binário gcluster e os blueprints examples/ que você usa nas etapas a seguir.

Criar um 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. Copie o blueprint do Cluster Toolkit para seu diretório de trabalho atual:

    cp examples/gke-tpu-v6e/gke-tpu-v6e-advanced.yaml .

  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 a função Storage Admin. Para conceder essa função, edite o arquivo gke-tpu-v6e-advanced.yaml atualizando o service-account module 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
    - artifactregistry.reader

  4. 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:

    ./gcluster deploy gke-tpu-v6e-advanced.yaml \
    --vars project_id="${PROJECT}" \
    --vars deployment_name="${CLUSTER_NAME}" \
    --vars region="${REGION}" \
    --vars zone="${ZONE}" \
    --vars num_slices="${CLUSTER_NODEPOOL_COUNT}" \
    --vars tpu_topology="${TOPOLOGY}" \
    --vars authorized_cidr="0.0.0.0/0" \
    --vars reservation="${RESERVATION:-}" \
    -l IGNORE --auto-approve -w

  5. Configure a autenticação do Container Registry e conceda o papel de administrador de armazenamento (roles/storage.admin) às suas contas de serviço do GKE:

    gcloud auth configure-docker gcr.io --quiet
gcloud auth configure-docker "${REGION}-docker.pkg.dev" --quiet
gcloud projects add-iam-policy-binding "${PROJECT}" --member="serviceAccount:${CLUSTER_NAME}-gke-wl-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet
gcloud projects add-iam-policy-binding "${PROJECT}" --member="serviceAccount:${CLUSTER_NAME}-gke-np-sa@${PROJECT}.iam.gserviceaccount.com" --role="roles/storage.admin" --quiet

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 qwen-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='Qwen/Qwen3-30B-A3B-Instruct-2507' \
  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"

  3. Use o comando gcluster job logs para verificar o status do job de conversão:

    # 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 qwen-hf-to-mt --main-only -f

  4. Verifique se os arquivos de modelo convertidos estão disponíveis no seu bucket do Cloud Storage:

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

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="qwen-rl" \
  --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="(echo 190G > /sys/fs/cgroup/memory.max || echo 190G > /sys/fs/cgroup/memory/memory.limit_in_bytes) 2>/dev/null || true && \
    export VLLM_HOST_IP=\$(hostname -I | awk '{print \$1}') && \
    export VLLM_ENABLE_V1_MULTIPROCESSING=0 && \
    python3 -c \"import pathlib, tunix.generate.vllm_sampler as vs; p = pathlib.Path(vs.__file__); p.write_text(p.read_text().replace('reshard_chunk_size: Optional[int] = None', 'reshard_chunk_size: Optional[int] = 4').replace('reshard_chunk_size=self.config.reshard_chunk_size', 'reshard_chunk_size=4'))\" && \
    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} \
    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 \
    num_test_batches=0 \
    batch_size=8 \
    train_micro_batch_size=8 \
    max_target_length=512 \
    max_prefill_predict_length=256 \
    mu_dtype=bfloat16 \
    grad_dtype=bfloat16 \
    rollout_tensor_parallelism=1 \
    rollout_expert_parallelism=4 \
    trainer_devices_fraction=0.5 \
    sampler_devices_fraction=0.5 \
    tokenizer_path='Qwen/Qwen3-30B-A3B-Instruct-2507' \
    ici_tensor_parallelism=2 \
    ici_expert_parallelism=4 \
    ici_fsdp_parallelism=-1 \
    hbm_utilization_vllm=0.55 \
    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'}}\""

  • Para verificar o status do job de treinamento:

    # Use the list command to check status
./gcluster job list

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

# Check progress of the job
kubectl logs -f \
    -l jobset.sigs.k8s.io/replicatedjob-name=pathways-head \
    -c workload-container

  • Para verificar 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

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

./gcluster job submit \
  --name="qwen-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=True \
  use_multimodal=False \
  weight_dtype=bfloat16 \
  --override_model_architecture"

  • Para verificar o status do job de conversão:

    # 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 qwen-mt-to-hf --main-only -f
# The trained model is now available in gs://${GCS_BUCKET}/${MODEL_NAME}/hf-trained/

  • Para verificar 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}"
gcloud storage rm -r "gs://${GCS_BUCKET}"

# To delete the local deployment folder and copied blueprint
rm -rf .ghpc "${CLUSTER_NAME}" gke-tpu-v6e-advanced.yaml

A seguir