Como pré-treinar o FairSeq RoBERTa no Cloud TPU usando PyTorch

---

Neste tutorial, mostramos como pré-treinar o RoBERTa da FairSeq em um Cloud TPU. Especificamente, ele segue o tutorial do FairSeq, pré-treinando o modelo no conjunto de dados público do wikitext-103.

Atenção: neste modelo, é usado um conjunto de dados de terceiros. O Google não oferece declarações, proteções ou outras garantias sobre a validade ou quaisquer outros aspectos desse conjunto de dados.

Objetivos

• Crie e configure o ambiente PyTorch
• Preparar o conjunto de dados
• Executar o job de treinamento
• Verificar se é possível visualizar os resultados de saída

Custos

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

• Compute Engine
• 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 uma avaliação gratuita.

Antes de começar

Antes de começar o tutorial, verifique se o projeto do Google Cloud foi configurado corretamente.

1. Faça login na sua conta do Google Cloud. Se você começou a usar o Google Cloud agora, crie uma conta para avaliar o desempenho de nossos produtos em situações reais. Clientes novos também recebem US$ 300 em créditos para executar, testar e implantar cargas de trabalho.

2. No console do Google Cloud, na página do seletor de projetos, selecione ou crie um projeto do Google Cloud.

   Observação: se você não pretende manter os recursos criados neste procedimento, crie um projeto novo em vez de selecionar um que já existe. Depois de concluir essas etapas, é possível excluir o projeto. Para fazer isso, basta remover todos os recursos associados a ele.

   Acessar o seletor de projetos

3. Verifique se a cobrança está ativada para o seu projeto do Google Cloud.

4. No console do Google Cloud, na página do seletor de projetos, selecione ou crie um projeto do Google Cloud.

   Observação: se você não pretende manter os recursos criados neste procedimento, crie um projeto novo em vez de selecionar um que já existe. Depois de concluir essas etapas, é possível excluir o projeto. Para fazer isso, basta remover todos os recursos associados a ele.

   Acessar o seletor de projetos

5. Verifique se a cobrança está ativada para o seu projeto do Google Cloud.

Este tutorial usa componentes faturáveis do Google Cloud. Consulte a página de preços da Cloud TPU para fazer uma estimativa dos custos. Para evitar cobranças desnecessárias, não se esqueça de apagar os recursos criados ao terminar de usá-los.

Configurar uma instância do Compute Engine

1. Abra uma janela do Cloud Shell.

   Abrir V

2. Crie uma variável para o ID do seu projeto.

   export PROJECT_ID=project-id
   

3. Configure o Google Cloud CLI para usar o projeto em que você quer criar a Cloud TPU.

   gcloud config set project ${PROJECT_ID}
   

   Na primeira vez que você executar esse comando em uma nova VM do Cloud Shell, será exibida uma página Authorize Cloud Shell. Clique em Authorize na parte inferior da página para permitir que gcloud faça chamadas da API Google Cloud com suas credenciais.

4. No Cloud Shell, inicie o recurso do Compute Engine necessário para este tutorial.

   gcloud compute instances create roberta-tutorial \
   --zone=us-central1-a \
   --machine-type=n1-standard-16  \
   --image-family=torch-xla \
   --image-project=ml-images  \
   --boot-disk-size=200GB \
   --scopes=https://www.googleapis.com/auth/cloud-platform
   

   Observação :ao executar esse comando, você vai ver um aviso:

   Disk size: '200 GB' is larger than image size
   

   Você pode ignorar a mensagem de aviso com segurança.

5. Conecte-se à nova instância do Compute Engine.

   gcloud compute ssh roberta-tutorial --zone=us-central1-a
   

   Importante: a partir desse ponto, o prefixo (vm) $ significa que é necessário executar o comando na instância de VM do Compute Engine.

Iniciar um recurso da Cloud TPU

1. Inicie um recurso do Cloud TPU na máquina virtual do Compute Engine, usando o seguinte comando:

   (vm) $ gcloud compute tpus create roberta-tutorial \
   --zone=us-central1-a \
   --network=default \
   --version=pytorch-2.0  \
   --accelerator-type=v3-8
   

2. Identifique o endereço IP do recurso do Cloud TPU.

   (vm) $ gcloud compute tpus describe --zone=us-central1-a roberta-tutorial
   

   Importante: o endereço IP está localizado na coluna NETWORK_ENDPOINTS. Você precisará desse endereço IP, sem o número da porta, ao criar e configurar o ambiente PyTorch.

Crie e configure o ambiente PyTorch

1. Inicie um ambiente conda.

   (vm) $ conda activate torch-xla-2.0
   

2. Configure variáveis de ambiente para o recurso Cloud TPU.

   Observação: a variável TPU_IP_ADDRESS precisa ser igual ao endereço IP do Cloud TPU que você identificou quando lançou o recurso. É o endereço IP interno exibido no Console do Google Cloud, em Compute Engine > TPUs.

   (vm) $ export TPU_IP_ADDRESS=ip-address
   

   (vm) $ export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470"
   

Configurar os dados

1. Instale o FairSeq executando:

   (vm) $ pip install --editable /usr/share/torch-xla-2.0/tpu-examples/deps/fairseq
   

2. Crie um diretório, pytorch-tutorial-data para armazenar os dados do modelo.

   (vm) $ mkdir $HOME/pytorch-tutorial-data
   (vm) $ cd $HOME/pytorch-tutorial-data
   

3. Siga as instruções na seção "Pré-processar os dados" do FairSeq RoBERTa. Leva aproximadamente 10 minutos para preparar o conjunto de dados.

Treine o modelo

Para treinar o modelo, primeiro configure algumas variáveis de ambiente:

(vm) $ export TOTAL_UPDATES=125000    # Total number of training steps
(vm) $ export WARMUP_UPDATES=10000    # Warmup the learning rate over this many updates
(vm) $ export PEAK_LR=0.0005          # Peak learning rate, adjust as needed
(vm) $ export TOKENS_PER_SAMPLE=512   # Max sequence length
(vm) $ export UPDATE_FREQ=16          # Increase the batch size 16x
(vm) $ export DATA_DIR=${HOME}/pytorch-tutorial-data/data-bin/wikitext-103

Depois, execute este script:

(vm) $ python3 \
      /usr/share/torch-xla-pytorch-2.0/tpu-examples/deps/fairseq/train.py $DATA_DIR \
      --task=masked_lm --criterion=masked_lm \
      --arch=roberta_base --sample-break-mode=complete \
      --tokens-per-sample=512 \
      --optimizer=adam \
      --adam-betas='(0.9,0.98)' \
      --adam-eps=1e-6 \
      --clip-norm=0.0 \
      --lr-scheduler=polynomial_decay \
      --lr=0.0005 \
      --warmup-updates=10000 \
      --dropout=0.1 \
      --attention-dropout=0.1 \
      --weight-decay=0.01 \
      --update-freq=16 \
      --train-subset=train \
      --valid-subset=valid \
      --num_cores=8 \
      --metrics_debug \
      --save-dir=checkpoints \
      --log_steps=30 \
      --log-format=simple \
      --skip-invalid-size-inputs-valid-test \
      --suppress_loss_report \
      --input_shapes 16x512 18x480 21x384 \
      --max-epoch=1

O script de treinamento é executado por aproximadamente 15 minutos e, quando termina, gera uma mensagem semelhante a:

saved checkpoint /home/user/checkpoints/checkpoint1.pt
(epoch 1 @ 119 updates) (writing took 25.19265842437744 seconds)
| done training in 923.8 seconds

Verificar resultados de saída

Depois que o job de treinamento for concluído, será possível encontrar os pontos de verificação de modelo no seguinte diretório:

$HOME/checkpoints

Limpar

Execute uma limpeza para evitar cobranças desnecessárias na sua conta depois de usar os recursos criados:

1. Desconecte-se da instância do Compute Engine, caso ainda não tenha feito isso:

   (vm) $ exit
   

   Agora, o prompt será user@projectname, mostrando que você está no Cloud Shell.

2. No Cloud Shell, use a Google Cloud CLI para excluir a instância do Compute Engine.

   $ gcloud compute instances delete roberta-tutorial --zone=us-central1-a
   

3. Use o Google Cloud CLI para excluir o recurso do Cloud TPU.

   $ gcloud compute tpus delete roberta-tutorial --zone=us-central1-a
   

A seguir

Teste as colabs do PyTorch:

• Primeiros passos com o PyTorch em Cloud TPUs
• Como treinar MNIST em TPUs (em inglês)
• Como treinar o ResNet18 em TPUs com o conjunto de dados Cifar10 (em inglês)
• Inferência com modelo ResNet50 pré-treinado (em inglês)
• Transferência rápida de estilo neural
• Treinamento de vários núcleos da AlexNet no Fashion MNIST
• Treinamento de núcleo único da AlexNet no Fashion MNIST