Executar o ajuste supervisionado em uma VM de TPU usando o MaxText

Este tutorial oferece um guia detalhado para executar o ajuste refinado supervisionado (SFT) em uma única instância de v6e-8 máquina virtual (VM) da Unidade de Processamento de Tensor (TPU) no Google Cloud usando o MaxText, uma pilha de treinamento de alto desempenho baseada em JAX para modelos de linguagem grandes (LLMs).

Objetivos

• Configure uma instância de VM do Cloud TPU.
• Instale o MaxText e as dependências dele.
• 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:

• Cloud TPU

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.

Importante: para concluir este tutorial, é preciso ter acesso a TPUs v6e-8. Este tutorial pressupõe que você já tenha uma reserva para essas TPUs no seu projetoGoogle Cloud para garantir que a capacidade esteja disponível quando você precisar dela. 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. 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 insira 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 gemma3-4b.

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:

• Administrador da TPU (roles/tpu.admin)
• Usuário da conta de serviço (roles/iam.serviceAccountUser)
• Editor do Compute (roles/compute.editor)

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.

Configure o ambiente

Configure as variáveis de ambiente executando o seguinte script:

export PROJECT="YOUR_PROJECT_ID"
export ZONE="YOUR_ZONE"
export RESERVATION="YOUR_RESERVATION_NAME"
export NAME="YOUR_TPU_NAME"
export NETWORK="default"

Substitua:

• YOUR_PROJECT_ID: o ID do projeto do Google Cloud
• YOUR_ZONE: a zona que você quer usar
• YOUR_RESERVATION_NAME: sua reserva de capacidade
• YOUR_TPU_NAME: o nome da instância de VM da Cloud TPU

Execute o seguinte comando para fazer a autenticação com Google Cloud :

gcloud auth login

Criar uma VM do Cloud TPU

Crie uma instância de VM de TPU do Cloud TPU com oito chips de TPU v6e vinculados à sua reserva de capacidade.

gcloud compute instances create "${NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --network="${NETWORK}" \
    --tags="${NAME}" \
    --machine-type=ct6e-standard-8t \
    --image-project=ubuntu-os-accelerator-images \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --boot-disk-size=200GB \
    --maintenance-policy=TERMINATE \
    --instance-termination-action=DELETE \
    --provisioning-model=RESERVATION_BOUND \
    --reservation-affinity=specific \
    --reservation="${RESERVATION}"

Depois que a instância de VM for criada, conecte-se a ela usando SSH.

gcloud compute ssh "${NAME}" --zone "${ZONE}" --project "${PROJECT}"

Conclua as etapas a seguir na sua instância de VM de TPU.

Instalar o MaxText

Atualize os pacotes do sistema na instância de VM da TPU.

sudo apt update && sudo apt upgrade -y --fix-missing

Instale o Python 3.12, que é necessário para o MaxText, e o pacote de ambiente virtual dele.

sudo apt install -y build-essential cmake ninja-build

Use uv para acelerar a instalação do pacote Python.

curl -LsSf https://astral.sh/uv/install.sh | sh
source "${HOME}/.local/bin/env"

Crie e ative um ambiente virtual chamado maxtext_venv.

uv python install 3.12
uv venv --python 3.12 --seed maxtext_venv
source maxtext_venv/bin/activate

Instale o MaxText e as dependências necessárias para tarefas pós-treinamento.

UV_TORCH_BACKEND=cpu uv pip install "maxtext[tpu-post-train]==0.2.2" --resolution=lowest

Instale as dependências necessárias restantes executando o seguinte comando:

install_tpu_post_train_extra_deps

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.

Especifique as variáveis de ambiente, como o token de acesso do Hugging Face, o nome do modelo que você quer usar e o diretório em que quer salvar o modelo no formato MaxText.

export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME='gemma3-4b'
export MODEL_CHECKPOINT_DIRECTORY=/dev/shm/$MODEL_NAME/mt-format/
export USE_PATHWAYS=0 # Set to 1 for Pathways, 0 for McJAX
export LAZY_LOAD_TENSORS=False # True to use lazy load, False to use eager load.

Substitua YOUR_HF_TOKEN pelo token de acesso do Hugging Face que você criou anteriormente.

Para converter o modelo do formato Hugging Face para o formato MaxText, execute o script a seguir. Essa conversão leva cerca de cinco minutos para ser concluída.

python3 -m maxtext.checkpoint_conversion.to_maxtext \
    model_name="${MODEL_NAME?}" \
    hf_access_token="${HF_TOKEN?}" \
    base_output_directory="${MODEL_CHECKPOINT_DIRECTORY?}" \
    scan_layers=True \
    use_multimodal=False \
    hardware=cpu \
    skip_jax_distributed_system=true \
    checkpoint_storage_use_zarr3=$((1 - USE_PATHWAYS)) \
    checkpoint_storage_use_ocdbt=$((1 - USE_PATHWAYS)) \
    --lazy_load_tensors="${LAZY_LOAD_TENSORS?}"

Iniciar a carga de trabalho de treinamento

Depois que o processo de conversão for concluído, você poderá iniciar a carga de trabalho do SFT.

1. Configure os parâmetros de treinamento da carga de trabalho SFT.

   # -- MaxText configuration --
   export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
   RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
   export RUN_NAME
   export STEPS=1000
   export PER_DEVICE_BATCH_SIZE=1
   
   # -- Dataset configuration --
   export DATASET_NAME="HuggingFaceH4/ultrachat_200k"
   export TRAIN_SPLIT="train_sft"
   export TRAIN_DATA_COLUMNS="['messages']"
   
   export MAXTEXT_CKPT_PATH=$MODEL_CHECKPOINT_DIRECTORY/0/items
   export TPU_ACCELERATOR_TYPE=v6e-8
   export TPU_WORKER_ID=0
   TPU_NAME=$(hostname)
   export TPU_NAME
   export TPU_SKIP_MDS_QUERY=1
   export TPU_WORKER_HOSTNAMES=localhost
   export TPU_TOPOLOGY=2x4
   export TPU_CHIPS_PER_HOST_BOUNDS=2,4,1
   export TPU_HOST_BOUNDS=1,1,1

2. Inicie o job de treinamento. Isso leva cerca de 10 minutos em uma instância de VM v6e-8.

   python3 -m maxtext.trainers.post_train.sft.train_sft \
       run_name="${RUN_NAME?}" \
       base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \
       model_name="${MODEL_NAME?}" \
       load_parameters_path="${MAXTEXT_CKPT_PATH?}" \
       per_device_batch_size="${PER_DEVICE_BATCH_SIZE?}" \
       steps="${STEPS?}" \
       hf_path="${DATASET_NAME?}" \
       train_split="${TRAIN_SPLIT?}" \
       train_data_columns="${TRAIN_DATA_COLUMNS?}" \
       profiler=xplane

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.

1. Defina os caminhos para exportação e os parâmetros treinados.

   export HF_EXPORT=/dev/shm/$MODEL_NAME/hf-trained/
   export POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/$STEPS/model_params

2. Execute a conversão de volta para o formato do Hugging Face.

   python3 -m maxtext.checkpoint_conversion.to_huggingface \
       model_name="${MODEL_NAME}" \
       load_parameters_path="${POST_TRAIN_PATH}" \
       base_output_directory="${HF_EXPORT}" \
       scan_layers=True \
       use_multimodal=False \
       weight_dtype=bfloat16

Depois que a conversão for concluída, o modelo ajustado armazenado em /dev/shm/gemma3-4b/hf-trained estará pronto para uso. Como você perde o acesso ao conteúdo da pasta /dev/shm quando a VM é reinicializada, mova o modelo ajustado para um armazenamento permanente ou faça upload dele para o Hugging Face Hub.

Limpar

Para evitar cobranças extras, exclua os recursos criados durante este tutorial.

Excluir a instância de VM de TPU

Saia da instância de VM de TPU da Cloud TPU e exclua-a.

# 1. Delete TPU instance
echo "Deleting TPU instance: ${NAME}..."
gcloud compute instances delete "${NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --quiet || true

# 2. Delete IAP firewall rule
echo "Deleting firewall rule: ${FIREWALL_RULE_NAME}..."
gcloud compute firewall-rules delete "${FIREWALL_RULE_NAME}" \
    --project="${PROJECT}" \
    --quiet || true

A seguir

• Para mais informações sobre o Cloud TPU, consulte Introdução ao Cloud TPU.
• Para detalhes da arquitetura e da configuração da TPU v6e-8, consulte TPU v6e.