Como treinar HuggingFace GPT2 no Cloud TPU (TF 2.x)


Se você não conhece o Cloud TPU, recomendamos que consulte o quickstart para aprender a criar uma VM de TPU.

Neste tutorial, mostramos como treinar o modelo HuggingFace GPT2 no Cloud TPU.

Objetivos

  • Criar uma Cloud TPU
  • Instalar dependências
  • Executar o job de treinamento

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.

    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.

    Acessar o seletor de projetos

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

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

Treinar o HuggingFace GPT2 com Cloud TPUs

  1. Abra uma janela do Cloud Shell.

    Abrir o Cloud Shell

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

    export PROJECT_ID=your-project-id
    
  3. Configure a CLI do Google Cloud para usar o projeto do Google Cloud em que você quer criar um 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 de API do Google Cloud com suas credenciais.

  4. Crie uma conta de serviço para o projeto da Cloud TPU.

    As contas de serviço permitem que o serviço do Cloud TPU acesse outros serviços do Google Cloud.

    $ gcloud beta services identity create --service tpu.googleapis.com --project $PROJECT_ID
    

    O comando retorna uma conta de serviço do Cloud TPU com o formato a seguir:

    service-PROJECT_NUMBER@cloud-tpu.iam.gserviceaccount.com
    

Criar uma Cloud TPU

  1. Crie uma VM do Cloud TPU usando o comando gcloud. O comando a seguir cria uma TPU v4-8. Também é possível criar um podslice de TPU definindo a sinalização --accelerator-type como um tipo de podslice, por exemplo, v4-32.

    $ gcloud compute tpus tpu-vm create hf-gpt2 \
      --zone=us-central2-b \
      --accelerator-type=v4-8 \
      --version=tpu-vm-tf-2.16.1-pjrt
    

    Descrições de sinalizações de comando

    zone
    A zona em que você planeja criar o Cloud TPU.
    accelerator-type
    O tipo de acelerador especifica a versão e o tamanho do Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, acesse Versões de TPU.
    version
    A versão do software do Cloud TPU.
  2. Conecte-se à VM do Cloud TPU executando o comando ssh a seguir.

    gcloud compute tpus tpu-vm ssh hf-gpt2 --zone=us-central2-b
    

Instalar dependências

  1. Clone o repositório HuggingFace Transformers:

    (vm)$ cd /tmp
    (vm)$ git clone https://github.com/huggingface/transformers.git
    (vm)$ cd transformers
    
  2. Instale as dependências:

    (vm)$ pip install .
    (vm)$ pip install -r examples/tensorflow/_tests_requirements.txt
    (vm)$ cd /tmp/transformers/examples/tensorflow/language-modeling
    (vm)$ pip install -r requirements.txt
    
  3. Crie um diretório temporário:

    (vm)$ mkdir /tmp/gpt2-wikitext
    
  4. Ao criar sua TPU, se você definir o parâmetro --version como uma versão que termina com -pjrt, defina as seguintes variáveis de ambiente para ativar o ambiente de execução PJRT:

      (vm)$ export NEXT_PLUGGABLE_DEVICE_USE_C_API=true
      (vm)$ export TF_PLUGGABLE_DEVICE_LIBRARY_PATH=/lib/libtpu.so
    

Executar script de treinamento

(vm)$ python3 run_clm.py \
  --model_name_or_path distilgpt2 \
  --max_train_samples 1000 \
  --max_eval_samples 100 \
  --num_train_epochs 1 \
  --output_dir /tmp/gpt2-wikitext \
  --dataset_name wikitext \
  --dataset_config_name wikitext-103-raw-v1

Descrições de sinalizações de comando

model_name_or_path
O nome do modelo a ser treinado.
max_train_samples
O número máximo de amostras a serem usadas no treinamento.
max_eval_samples
O número máximo de amostras a serem usadas na avaliação.
num_train_epochs
O número de períodos para treinar o modelo.
output_dir
O diretório de saída do script de treinamento.
dataset_name
O nome do conjunto de dados a ser usado.
dataset_config_name
O nome de configuração do conjunto de dados

Quando o treinamento terminar, uma mensagem semelhante a esta será exibida:

  125/125 [============================] - ETA: 0s - loss: 3.61762023-07-07 21:38:17.902850: W tensorflow/core/framework/dataset.cc:956] Input of GeneratorDatasetOp::Dataset will not be optimized because the dataset does not implement the AsGraphDefInternal() method needed to apply optimizations.
  125/125 [============================] - 763s 6s/step - loss: 3.6176 - val_loss: 3.4233
  Configuration saved in /tmp/gpt2-wikitext/config.json
  Configuration saved in /tmp/gpt2-wikitext/generation_config.json
  Model weights saved in /tmp/gpt2-wikitext/tf_model.h5
  D0707 21:38:45.640973681   12027 init.cc:191]                          grpc_shutdown starts clean-up now
  

Limpar

  1. Desconecte-se da instância de VM da TPU:

    (vm)$ exit
    

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

  2. Exclua o recurso da TPU.

    $ gcloud compute tpus tpu-vm delete hf-gpt2 \
    --zone=us-central2-b
    

A seguir

Teste um dos outros modelos de referência compatíveis.