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 treinamento de aprendizado por reforço em uma VM de TPU usando o MaxText

Este tutorial oferece um guia detalhado para executar o treinamento de aprendizado por reforço (RL, na sigla em inglês) 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 alta performance baseada em JAX para modelos de linguagem grandes (LLMs).

Resumo da carga de trabalho

A execução dessa carga de trabalho de pós-treinamento de RL no Cloud TPU envolve os seguintes componentes:

Modelo:Llama 3.1 8B Instruct (llama3.1-8b-Instruct), um LLM de pesos abertos otimizado para tarefas de diálogo e de acompanhamento de instruções. 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 distribuída em uma TPU v6e-8 usando JAX.
Algoritmo de aprendizado por reforço (RL):otimização de política relativa de grupo (GRPO, na sigla em inglês), em que a carga de trabalho executa um pipeline de pós-treinamento escalonável usando o treinador de RL integrado do MaxText. O GRPO elimina a necessidade de um modelo de valor ou crítico separado, avaliando um grupo de saídas geradas para o mesmo prompt, usando recompensas relativas para otimizar a política do ator de maneira eficiente.
Conjunto de dados e modelo:Grade School Math 8K (GSM8K), um conjunto de dados de alta qualidade que consiste em problemas de palavras matemáticas do ensino fundamental linguisticamente diversos. Essa carga de trabalho é usada para gerar implementações, calcular recompensas e melhorar os recursos de raciocínio em várias etapas durante a fase de alinhamento de RL.
Infraestrutura e framework:um único nó v6e-8 do Cloud TPU (ct6e-standard-8t) usando o MaxText (maxtext[tpu-post-train]), uma pilha de treinamento de alta performance 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 otimização de política relativa de grupo (GRPO) de RL na TPU.
Converter o modelo treinado de volta para o formato Hugging Face para exibiçã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.

Antes de começar

Você precisa de um token de acesso do Hugging Face para usar este tutorial. Inscreva-se 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 llama3.1-8b-Instruct.

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.

Configurar 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 TPU_NAME="YOUR_TPU_NAME"

Substitua:

YOUR_PROJECT_ID: ID do 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 "${TPU_NAME}" \
    --zone="${ZONE}" \
    --project="${PROJECT}" \
    --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 "${TPU_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 de pós-treinamento.

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

Execute o seguinte comando para instalar as dependências restantes:

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.

Forneça os valores a seguir:

Seu token de acesso do Hugging Face
O nome do modelo que você quer usar
O diretório em que você quer salvar o modelo no formato MaxText
Opções de carregamento e armazenamento

export HF_TOKEN="YOUR_HF_TOKEN"
export MODEL_NAME='llama3.1-8b-Instruct'
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 RL.

Configure os parâmetros de treinamento da carga de trabalho de RL.

# -- 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 CHIPS_PER_VM=8
export NUM_BATCHES=50
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

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.rl.train_rl \
    model_name="${MODEL_NAME?}" \
    load_parameters_path="${MAXTEXT_CKPT_PATH?}" \
    run_name="${RUN_NAME?}" \
    base_output_directory="${BASE_OUTPUT_DIRECTORY?}" \
    chips_per_vm="${CHIPS_PER_VM?}" \
    num_batches="${NUM_BATCHES?}" \
    num_test_batches=10 \
    rollout_data_parallelism=1 \
    rollout_tensor_parallelism=-1

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.

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

export HF_EXPORT=/dev/shm/$MODEL_NAME/hf-trained/
export HF_MODEL_NAME=llama3.1-8b
export POST_TRAIN_PATH=$BASE_OUTPUT_DIRECTORY/$RUN_NAME/checkpoints/actor/$NUM_BATCHES/model_params

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

python3 -m maxtext.checkpoint_conversion.to_huggingface \
    model_name="${HF_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/$MODEL_NAME/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

Exclua a instância de VM de TPU do Cloud TPU.

gcloud compute instances delete "${TPU_NAME}" --zone="${ZONE}" --project="${PROJECT}" --quiet

A seguir

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