O Google usa tecnologia de IA na tradução de conteúdos para seu idioma de preferência. As traduções com IA podem ter erros.

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 da Unidade de Processamento de Tensor (TPUs) v6e de vários hosts 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 SFT na TPU.
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 em uma conta sem custo financeiro no Hugging Face. Depois de criar uma conta, gere um token de acesso:

Na página Bem-vindo ao Hugging Face, clique no avatar da sua conta e selecione Tokens de acesso.
Na página Tokens de acesso, clique em Criar novo token.
Selecione o tipo de token Leitura e insira um nome para ele.
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 gemma4-31b.

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 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)
- Administrador do Storage (roles/storage.admin)

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

export PROJECT="YOUR_PROJECT_ID"
export REGION="YOUR_REGION"
export ZONE="YOUR_ZONE"
export CLUSTER_NAME="gke-tpu-v6e"
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 MODEL_NAME="gemma4-31b"
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_RESERVATION_NAME: o nome da sua reserva.
YOUR_HF_TOKEN: seu token de acesso do Hugging Face.
YOUR_BUCKET_NAME: um nome exclusivo global 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.

Criar um cluster do Cluster Toolkit

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

Crie um bucket do Cloud Storage:

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

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 essa função, 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

Implante 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 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 \
    -l IGNORE \
    -w

Converter o modelo para o formato MaxText

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

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

Para converter o modelo do formato Hugging Face para o formato MaxText e armazená-lo no bucket do Cloud Storage, execute o seguinte script:

./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 seguinte comando:

# 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, inicie a carga de trabalho do SFT executando o seguinte comando:

./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 seguinte comando:

# 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 \
        --override_model_architecture"