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.

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

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

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

    Acessar a página do seletor de projetos

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

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

    Abra o Cloud Shell

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

    export PROJECT_ID=project-id

  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_ID}

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

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

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

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 dlrm-tutorial \
--zone=us-central1-a \
--network=default \
--version=pytorch-1.6  \
--accelerator-type=v3-8

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

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

Crie e configure o ambiente PyTorch

  1. Inicie um ambiente conda.

    (vm) $ conda activate torch-xla-1.6

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

    (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

  1. Receba DLRM compatível com TPU executando:

    (vm) $ git clone --recursive https://github.com/pytorch-tpu/examples.git

  2. Instale as dependências.

    (vm) $ pip install onnx

  3. Executar o modelo em dados aleatórios.

    (vm) $ python 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-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

  1. Faça o download do conjunto de dados.

    Faça o download do conjunto de dados do Criteo Kaggle seguindo estas instruções. Depois de fazer o download do arquivo dac.tar.gz, coloque-o no diretório ./criteo-kaggle/.

     (vm) $ mkdir criteo-kaggle
 (vm) $ cd criteo-kaggle
 (vm) $ # Download dataset from above link here.
 (vm) $ tar -xzvf dac.tar.gz
 (vm) $ cd ..

  2. Pré-processar o conjunto de dados.

    Pré-processe o conjunto de dados Criteo com o comando a seguir. Ele é praticamente o mesmo comando que o comando de treinamento. No entanto, ele é executado apenas em um núcleo, porque queremos encerrar a execução do treinamento depois que ele for feito com pré-processamento e produção do kaggleAdDisplayChallenge_processed.npz. O pré-processamento demora um pouco (hora+).

    (vm) $ python 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 \
    --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 \
    --test-freq=101376 \
    --use-tpu \
    --num-indices-per-lookup=1 \
    --num-indices-per-lookup-fixed \
    --tpu-model-parallel-group-len 1 \
    --tpu-metrics-debug \
    --tpu-cores=1

    Quando você vir registros como o seguinte, o conjunto de dados foi pré-processado e é possível encerrar (Ctrl^C) a execução e executar pkill python apenas para garantir que não haja processos vazados.

    Load 6548659/6548660  Split: 0  Label True: 0  Stored: 0
Saved /home/jysohn/test/train_day_0.npz!
Load 6548659/6548660  Split: 1  Label True: 1  Stored: 1
Saved /home/jysohn/test/train_day_1.npz!
Load 3742283/6548660  Split: 2  Label True: 0  Stored: 0
...
time/loss/accuracy (if enabled): 2020-08-26 13:17:18.139566
Finished training it 1024/306968 of epoch 0, -1.00 ms/it, loss 0.518172, accuracy 75.720 %, 131072 samples, @ 2020-08-26 13:18:04.515570

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

    Você verá o kaggleAdDisplayChallenge_processed.npz produzido dentro do diretório criteo-kaggle.

  4. Execute treinamento no conjunto de dados pré-processado do Criteo Kaggle.

    (vm) $ python 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 \
    --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 precisa ser concluído em mais de três horas, com o acúmulo de mais de 78,75%.

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 dlrm-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 dlrm-tutorial --zone=us-central1-a

A seguir

Teste as colabs do PyTorch: