Como treinar o FairSeq Transformer no Cloud TPU usando PyTorch

Alfa

Este produto ou recurso está em um estado de pré-lançamento e pode ser alterado ou ter suporte limitado. Para mais informações, consulte Estágios de lançamento do produto.

Para mais informações sobre o acesso a esta versão, consulte a página de solicitação de acesso.

Este tutorial se concentra especificamente na versão FairSeq do Transformer e na tarefa de tradução do WMT 18, traduzindo do inglês para o alemão.

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

• Preparar o conjunto de dados.
• Executar o job de treinamento.
• Verificar os resultados da saída.

Custos

Neste tutorial, há componentes faturáveis do Google Cloud Platform, entre eles:

• Compute Engine
• Cloud TPU

Use a calculadora de preços para gerar uma estimativa de custos baseada na projeção de uso. Os novos usuários do GCP 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 Platform foi configurado corretamente.

1. Faça login na sua Conta do Google.

   Se você ainda não tiver uma, inscreva-se.

2. Selecione ou crie um projeto do Google Cloud Platform.

   Observação: se você não pretende manter os recursos criados neste tutorial, crie um novo projeto em vez de selecionar um existente. Depois de terminar, é possível excluir o projeto, removendo todos os recursos associados ao projeto e ao tutorial.

   Acessar a página Gerenciar recursos

3. Verifique se o faturamento foi ativado no projeto do Google Cloud Platform.

   Saiba como ativar o faturamento

Este tutorial inclui componentes faturáveis do Google Cloud Platform. Consulte a página de preços do Cloud TPU para fazer uma estimativa dos custos. Para evitar cobranças desnecessárias, certifique-se de limpar os recursos criados quando terminar de usá-los.

Configurar uma instância do Compute Engine

1. Abra uma janela do Cloud Shell.

   Abrir o Cloud Shell

2. Crie uma variável para o nome do projeto.

   export PROJECT_NAME=project_name
   

3. Configure a ferramenta de linha de comando gcloud para usar o projeto em que a Cloud TPU será criada.

   gcloud config set project ${PROJECT_NAME}
   

4. Inicie o recurso do Compute Engine necessário para este tutorial.

   gcloud compute --project=${PROJECT_NAME} instances create transformer-pytorch-tutorial \
   --zone=us-central1-a  \
   --machine-type=n1-standard-8  \
   --image-family=torch-xla \
   --image-project=ml-images  \
   --boot-disk-size=500GB
   

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

   gcloud compute ssh transformer-pytorch-tutorial --zone=us-central1-a
   

Daqui em diante, o prefixo (vm)$ significa que você precisa 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 transformer-pytorch-tutorial \
   --zone=us-central1-a \
   --network=default \
   --range=10.2.3.0 \
   --version=pytorch-0.5 \
   --accelerator-type=v3-8
   

   Para mais informações sobre intervalos de endereços IP e o Cloud TPU, consulte Como configurar endereços IP internos TPU.

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

   (vm) $ gcloud compute tpus list --zone=us-central1-a
   

   O endereço IP está localizado abaixo da coluna NETWORK_ENDPOINTS. Você precisará desse endereço IP ao criar e configurar o ambiente PyTorch.

Fazer o download dos dados

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

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

2. Navegue até o diretório pytorch-tutorial-data.

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

3. Faça o download dos dados do modelo.

   (vm) $ wget https://dl.fbaipublicfiles.com/fairseq/data/wmt18_en_de_bpej32k.zip
   

4. Extraia os dados.

   (vm) $ sudo apt-get install unzip && \
   unzip wmt18_en_de_bpej32k.zip
   

Crie e configure o ambiente PyTorch

1. Inicie um ambiente conda.

   (vm) $ conda activate torch-xla-0.5
   

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

   (vm) $ export TPU_IP_ADDRESS=ip-address; \
   export XRT_TPU_CONFIG="tpu_worker;0;$TPU_IP_ADDRESS:8470";
   

   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.

Treine o modelo

Para treinar o modelo, execute o seguinte script:

Observação: devido ao tempo de CPU necessário para preparar o conjunto de dados e as compilações iniciais de gráficos de TPU, pode levar aproximadamente 20 minutos para treinar 100 etapas e gerar a primeira saída de registro.

(vm) $ python /usr/share/torch-xla-0.5/tpu-examples/deps/fairseq/train.py \
  $HOME/pytorch-tutorial-data/wmt18_en_de_bpej32k \
  --save-interval=1 \
  --arch=transformer_vaswani_wmt_en_de_big \
  --max-target-positions=64 \
  --attention-dropout=0.1 \
  --no-progress-bar \
  --criterion=label_smoothed_cross_entropy \
  --source-lang=en \
  --lr-scheduler=inverse_sqrt \
  --min-lr 1e-09 \
  --skip-invalid-size-inputs-valid-test \
  --target-lang=de \
  --label-smoothing=0.1 \
  --update-freq=1 \
  --optimizer adam \
  --adam-betas '(0.9, 0.98)' \
  --warmup-init-lr 1e-07 \
  --lr 0.0005 \
  --warmup-updates 4000 \
  --share-all-embeddings \
  --dropout 0.3 \
  --weight-decay 0.0 \
  --valid-subset=valid \
  --max-epoch=25 \
  --input_shapes 128x64 \
  --num_cores=8 \
  --metrics_debug \
  --log_steps=100

Observação: alterar o valor do hiperparâmetro ‘input_shapes’ pode melhorar o desempenho. Por exemplo:

--input_shapes 640x16 512x32 256x64

Essas alterações podem resultar em compilações iniciais significativamente mais lentas, mas tempos de época mais rápidos após a estabilização.

Ative o uso do bfloat16 para usar esses formatos de entrada. Para isso, execute o seguinte comando:

(vm) $ export XLA_USE_BF16=1

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:

(vm) $ $HOME/checkpoints

Como fazer a limpeza

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 seu Cloud Shell, use a ferramenta de linha de comando gcloud para excluir a instância do Compute Engine.

   $ gcloud compute instances delete transformer-pytorch-tutorial --zone="us-central1-a"
   

3. Use a ferramenta de linha de comando gcloud para excluir o recurso do Cloud TPU.

   $ gcloud compute tpus delete transformer-pytorch-tutorial --zone="us-central1-a"
   

A seguir

Teste as colabs do PyTorch:

• 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 ResNet pré-treinado (em inglês)