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

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

Resumo da carga de trabalho

A execução dessa carga de trabalho de SFT em uma VM do Cloud TPU envolve os seguintes componentes:

• Modelo:Gemma 3 4B (gemma3-4b), um modelo de linguagem leve de pesos abertos. Neste tutorial, os pesos do modelo são convertidos do formato Hugging Face para o formato compatível com MaxText (mt-format) e preparados para execução em uma única VM do Cloud TPU v6e-8 usando o JAX.
• Método de treinamento:SFT usando MaxText. A carga de trabalho executa um pipeline de ajuste fino em uma única instância do Cloud TPU v6e-8 usando o JAX. Ele otimiza os parâmetros do modelo em mais de 1.000 etapas usando o treinador de SFT integrado do MaxText (maxtext.trainers.post_train.sft.train_sft) e conjuntos de dados de ajuste de instruções focados em diálogos.
• Conjunto de dados e modelo:UltraChat 200k (HuggingFaceH4/ultrachat_200k), um conjunto de dados de diálogo de várias rodadas de alta qualidade e grande escala (train_sft dividido com a coluna messages). Ele é usado para ajustar o modelo e melhorar as capacidades de conversação e de seguir instruções.
• Infraestrutura e framework:uma única instância de VM do Cloud TPU v6e-8 (ct6e-standard-8t). O framework é o MaxText (maxtext[tpu-post-train]) com tecnologia JAX.

Objetivos

• Configurar uma instância de VM do Cloud TPU.
• Instalar o MaxText e as dependências dele.
• 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 Hugging Face para disponibilização.

Custos

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

• Cloud TPU

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.

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

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:

• Administrador de TPU (roles/tpu.admin)
• Usuário da conta de serviço (roles/iam.serviceAccountUser)
• Editor de computação (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 com 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: ID do seu Google Cloud projeto
• 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 de TPU do Cloud TPU

Autentique-se com Google Cloud executando o seguinte comando:

gcloud auth login

Criar a 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 instância de VM da 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 o MaxText exige, 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 um ambiente virtual chamado maxtext_venv e ative-o.

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, é necessário 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 você 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 seguinte script. 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 de SFT.

1. Configure os parâmetros de treinamento da carga de trabalho de 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 Hugging Face

Depois que a carga de trabalho de treinamento for concluída, converta o modelo de volta para o formato 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 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 o armazenamento permanente ou faça o upload dele para o Hugging Face Hub.

Limpar

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

Excluir a instância de VM da TPU

Saia da instância de VM de TPU do 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 de arquitetura e configuração da v6e-8 TPU, consulte TPU v6e.