Executar treinamento de aprendizado por reforço em uma VM de TPU usando o MaxText

Este tutorial fornece 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) usando o MaxText, uma pilha de treinamento de alto desempenho baseada em JAX para modelos de linguagem grandes (LLMs). Google Cloud

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 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 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:

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

export PROJECT="YOUR_PROJECT_ID"
export ZONE="ZONE_NAME"
export RESERVATION="RESERVATION_NAME"
export TPU_NAME="TPU_MACHINE_NAME"

Substitua:

  • YOUR_PROJECT_ID: ID do seu Google Cloud projeto
  • ZONE_NAME: a zona que você quer usar
  • RESERVATION_NAME: sua reserva de capacidade
  • TPU_MACHINE_NAME: o nome da instância de VM 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 do Cloud TPU com oito chips de TPU v6e, vinculados à sua reserva de capacidade.

gcloud alpha compute tpus tpu-vm create $TPU_NAME \
    --zone=$ZONE \
    --project=$PROJECT \
    --accelerator-type=v6e-8 \
    --version=v2-alpha-tpuv6e \
    --provisioning-model=reservation-bound \
    --reservation=$RESERVATION

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

gcloud compute tpus tpu-vm 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 python3.12 python3.12-venv

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 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 pip install maxtext[tpu-post-train]==0.2.2 --resolution=lowest

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

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 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 de RL.

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

    # -- MaxText configuration --
export BASE_OUTPUT_DIRECTORY=/dev/shm/$MODEL_NAME/post-train/
export RUN_NAME=$(date +%Y-%m-%d-%H-%M-%S)
export CHIPS_PER_VM=8
export NUM_BATCHES=50
export MAXTEXT_CKPT_PATH=$MODEL_CHECKPOINT_DIRECTORY/0/items

  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.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.

  1. 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

  2. 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 do Cloud TPU.

gcloud alpha compute tpus tpu-vm delete $TPU_NAME --zone=$ZONE --project=$PROJECT --quiet

A seguir