Como treinar o DLRM no Cloud TPU usando PyTorch

Neste tutorial, mostramos como treinar o Facebook Research DLRM em um Cloud TPU.

Objetivos

Crie e configure o ambiente PyTorch
Executar o job de treinamento com dados falsos
(Opcional) Treinar no conjunto de dados do Criteo Kaggle

Custos

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

Compute Engine
Cloud TPU

Use a calculadora de preços para gerar uma estimativa de custos baseada na projeção de uso. 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.

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.

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

Verifique se o faturamento está ativado para seu projeto na nuvem. Saiba como confirmar se o faturamento está ativado para o projeto.

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

Abra uma janela do Cloud Shell.

Abra o Cloud Shell
Crie uma variável para o ID do seu projeto.
```
export PROJECT_ID=project-id
```
Configure a ferramenta de linha de comando gcloud para usar o projeto em que a Cloud TPU será criada.
```
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 GCP com suas credenciais.
No Cloud Shell, inicie o recurso do Compute Engine necessário para este tutorial. Observação: use um n1-highmem-96 machine-type se estiver treinando no conjunto de dados do Criteo Kaggle.
```
gcloud compute instances create dlrm-tutorial \
--zone=us-central1-a \
--machine-type=n1-standard-64 \
--image-family=torch-xla \
--image-project=ml-images  \
--boot-disk-size=200GB \
--scopes=https://www.googleapis.com/auth/cloud-platform
```
Observação:
ao executar este comando, você verá um alerta:
```
Disk size: '200 GB' is larger than image size
```
É possível ignorar a mensagem de alerta com segurança.
Conecte-se à nova instância do Compute Engine.
```
gcloud compute ssh dlrm-tutorial --zone=us-central1-a
```
Ponto-chave: a partir desse ponto, um prefixo (vm) $ significa que é necessário executar o comando na instância de VM do Compute Engine.

Iniciar um recurso da Cloud TPU

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

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

Identifique o endereço IP do recurso do Cloud TPU.
```
(vm) $ gcloud compute tpus list --zone=us-central1-a
```
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

Inicie um ambiente conda.
```
(vm) $ conda activate torch-xla-1.9
```
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"
```

Executar o job de treinamento com dados falsos

Instale as dependências.
```
(vm) $ pip install onnx
```

Executar o modelo em dados aleatórios. Esse processo leva de 5 a 10 minutos.

(vm) $ python /usr/share/torch-xla-1.9/tpu-examples/deps/dlrm/dlrm_tpu_runner.py \
    --arch-embedding-size=1000000-1000000-1000000-1000000-1000000-1000000-1000000-1000000 \
    --arch-sparse-feature-size=64 \
    --arch-mlp-bot=512-512-64 \
    --arch-mlp-top=1024-1024-1024-1 \
    --arch-interaction-op=dot \
    --lr-num-warmup-steps=10 \
    --lr-decay-start-step=10 \
    --mini-batch-size=2048 \
    --num-batches=1000 \
    --data-generation='random' \
    --numpy-rand-seed=727 \
    --print-time \
    --print-freq=100 \
    --num-indices-per-lookup=100 \
    --use-tpu \
    --num-workers 7 \
    --num-indices-per-lookup-fixed \
    --tpu-model-parallel-group-len=8 \
    --tpu-metrics-debug \
    --tpu-cores=8

(Opcional) Treinar no conjunto de dados do Criteo Kaggle

Essas etapas são opcionais. Execute-as somente se quiser treinar no conjunto de dados do Criteo Kaggle.

Faça o download do conjunto de dados.

Faça o download do conjunto de dados do Criteo Kaggle seguindo estas instruções. Quando o download for concluído, copie o arquivo dac.tar.gz em um diretório chamado ./criteo-kaggle/. Use o comando tar -xzvf para extrair o conteúdo do arquivo tar.gz no diretório ./critero-kaggle.
```
 (vm) $ mkdir criteo-kaggle
 (vm) $ cd criteo-kaggle
 (vm) $ # Download dataset from above link here.
 (vm) $ tar -xzvf dac.tar.gz
 (vm) $ cd ..
```

Pré-processar o conjunto de dados.

Inicie esse script para pré-processar o conjunto de dados Criteo. Esse script produz um arquivo chamado kaggleAdDisplayChallenge_processed.npz e leva mais de três horas para pré-processar o conjunto de dados.

(vm) $ python /usr/share/torch-xla-1.9/tpu-examples/deps/dlrm/dlrm_data_pytorch.py \
    --data-generation=dataset \
    --data-set=kaggle \
    --raw-data-file=criteo-kaggle/train.txt \
    --mini-batch-size=128 \
    --memory-map \
    --test-mini-batch-size=16384 \
    --test-num-workers=4

Verifique se o pré-processamento foi bem-sucedido.

Você verá o arquivo kaggleAdDisplayChallenge_processed.npz no diretório criteo-kaggle.

Execute o script de treinamento no conjunto de dados pré-processado do Kaggle.

(vm) $ python /usr/share/torch-xla-1.9/tpu-examples/deps/dlrm/dlrm_tpu_runner.py \
    --arch-sparse-feature-size=16 \
    --arch-mlp-bot="13-512-256-64-16" \
    --arch-mlp-top="512-256-1" \
    --data-generation=dataset \
    --data-set=kaggle \
    --raw-data-file=criteo-kaggle/train.txt \
    --processed-data-file=criteo-kaggle/kaggleAdDisplayChallenge_processed.npz \
    --loss-function=bce \
    --round-targets=True \
    --learning-rate=0.1 \
    --mini-batch-size=128 \
    --print-freq=1024 \
    --print-time \
    --test-mini-batch-size=16384 \
    --test-num-workers=4 \
    --memory-map \
    --test-freq=101376 \
    --use-tpu \
    --num-indices-per-lookup=1 \
    --num-indices-per-lookup-fixed \
    --tpu-model-parallel-group-len 8 \
    --tpu-metrics-debug \
    --tpu-cores=8

O treinamento deve ser concluído em mais de duas horas e precisão de 78,75%.

Limpeza

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

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.
No Cloud Shell, use a ferramenta de linha de comando gcloud para excluir a instância do Compute Engine.
```
$ gcloud compute instances delete dlrm-tutorial --zone=us-central1-a
```
Use a ferramenta de linha de comando gcloud para excluir o recurso do Cloud TPU.
```
$ gcloud compute tpus delete dlrm-tutorial --zone=us-central1-a
```

A seguir

Teste as colabs do PyTorch: