Executar o ajuste supervisionado no Gemma 4 na TPU v6e

Neste tutorial, mostramos como executar o ajuste supervisionado (SFT) em um cluster da Unidade de Processamento de Tensor (TPU) v6e usando o MaxText e o Cluster Toolkit. Você usa o Cluster Toolkit para executar uma carga de trabalho de treinamento multihost e exportar os resultados de volta para o formato do Hugging Face para disponibilizaçã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 SFT na TPU.
  • Converter o modelo ajustado de volta para o formato do Hugging Face para disponibilização.

Custos

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

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.

Antes de começar

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 "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 ele.
  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-31b.

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 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 executando o script a seguir:

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="gke-tpu-v6e"
export REPOSITORY_NAME="YOUR_REPOSITORY_NAME"
export CLOUD_IMAGE_NAME="${REGION}-docker.pkg.dev/${PROJECT}/${REPOSITORY_NAME}/maxtext_base:latest"
export TPU_TYPE="v6e-32"
export RESERVATION="YOUR_RESERVATION_NAME"
export MODEL_NAME="gemma4-31b"
export HF_TOKEN="YOUR_HF_TOKEN"
export GCS_BUCKET="YOUR_BUCKET_NAME"

Substitua:

  • YOUR_PROJECT_ID: o ID do seu 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_REPOSITORY_NAME: o nome do repositório do Artifact Registry para suas imagens do MaxText.
  • YOUR_RESERVATION_NAME: o nome da sua reserva.
  • YOUR_HF_TOKEN: seu token de acesso do Hugging Face.
  • YOUR_BUCKET_NAME: um nome globalmente exclusivo para um bucket do Cloud Storage.

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 do MaxText

Para preparar a imagem do 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 ${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.3 --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}"

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 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 módulo 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. 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=1,tpu_topology=4x8,authorized_cidr=0.0.0.0/0,reservation=${RESERVATION:-}" \
    --download-dependencies \
    -w

Converter o modelo para o formato MaxText

Para treinar o modelo no formato MaxText, é necessário convertê-lo do formato do 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. Para converter o modelo do formato do Hugging Face para o formato MaxText e armazená-lo no bucket do Cloud Storage, execute o script a seguir:

    ./gcluster job submit --name hf-to-mt \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --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"

Para verificar o status do job de conversão, execute o comando a seguir:

# 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 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, você poderá iniciar a carga de trabalho de SFT executando o comando a seguir:

./gcluster job submit --name sft \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --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"

Para verificar o status do job de treinamento, execute o comando a seguir:

# 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 sft --main-only -f \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION}

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 mt-to-hf \
    --cluster ${CLUSTER_NAME} \
    --project ${PROJECT} \
    --location ${REGION} \
    --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"

Para verificar o status do job de conversão, execute o comando a seguir:

# 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 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/ - though again, it's ~2x the size of the original...

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