Guia do usuário PyTorch/XLA da Cloud TPU

Executar cargas de trabalho de ML com PyTorch/XLA

Antes de iniciar os procedimentos deste guia, configure uma VM de TPU e ssh nela como descrito em Preparar um projeto do GCP. Isso configurará os recursos necessários para executar os comandos deste guia.

A PyTorch 1.8.1 e o PyTorch/XLA 1.8.1 são pré-instalados na VM da TPU.

Configuração básica

Defina a configuração do dispositivo XRT TPU:

   (vm)$ export XRT_TPU_CONFIG="localservice;0;localhost:51011"

No caso de modelos que têm alocações grandes e frequentes que usam, tcmalloc melhora o desempenho significativamente em comparação com o padrão malloc. Portanto, o malloc padrão usado na VM da TPU é tcmalloc. No entanto, dependendo da sua carga de trabalho (por exemplo, com DLRM que tem alocações muito grandes para as tabelas de embedding) tcmalloc pode causar uma lentidão. Nesse caso, você pode tentar cancelar a definição da variável a seguir usando o malloc padrão em vez disso:

   (vm)$ unset LD_PRELOAD

Realizar um cálculo simples

Inicie o interpretador do Python 3:

(vm)$ python3
   Python 3.6.9 (default, Jan 26 2021, 15:33:00)
   [GCC 8.4.0] on linux
   Type "help", "copyright", "credits" or "license" for more information.

No interpretador do Pytorch 3, importe os seguintes pacotes do PyTorch:

import torch
import torch_xla.core.xla_model as xm

Faça os seguintes cálculos:

dev = xm.xla_device()
A saída desse comando é semelhante a:
2021-04-01 23:20:23.268115: E   55362 tensorflow/core/framework/op_kernel.cc:1693] OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op: TPURoundRobin
2021-04-01 23:20:23.269345: E   55362 tensorflow/core/framework/op_kernel.cc:1693] OpKernel ('op: "TpuHandleToProtoKey" device_type: "CPU"') for unknown op: TpuHandleToProtoKey
t1 = torch.randn(3,3,device=dev)
t2 = torch.randn(3,3,device=dev)
print(t1 + t2)

Você verá a seguinte saída do comando:

tensor([[-0.2121,  1.5589, -0.6951],
        [-0.7886, -0.2022,  0.9242],
        [ 0.8555, -1.8698,  1.4333]], device='xla:1')

Como executar o Resnet em uma TPU de dispositivo único

Nesse ponto, você pode executar qualquer código do PyTorch / XLA. Por exemplo, é possível executar um ResNet com dados falsos usando:

   (vm)$ git clone --recursive https://github.com/pytorch/xla.git
   (vm)$ python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1

A amostra do ResNet treina por 1 época e leva cerca de 7 minutos. Ela retorna uma saída semelhante a esta:

Epoch 1 test end 20:57:52, Accuracy=100.00 Max Accuracy: 100.00%

Após o fim do treinamento do ResNet, exclua a VM do TPU.

   (vm)$ exit
   $ gcloud alpha compute tpus tpu-vm delete tpu-name \
   --zone=zone

Execute gcloud list para verificar se os recursos foram excluídos. A exclusão pode levar vários minutos. Uma resposta como esta indica que suas instâncias foram excluídas com sucesso.

   $ gcloud alpha compute tpus list --zone=zone
   

   Listed 0 items.
   

Configuração avançada

Nos exemplos acima (o cálculo simples e o ResNet50), seu programa PyTorch/XLA iniciará o servidor XRT local no mesmo processo. Também é possível iniciar o serviço local XRT em um processo separado:

(vm)$ python3 -m torch_xla.core.xrt_run_server --port 51011 --restart

A vantagem dessa abordagem é que o cache de compilação persistirá durante o treinamento. Neste caso, é possível encontrar o registro do servidor em /tmp/xrt_server_log.

(vm)$ ls /tmp/xrt_server_log/
server_20210401-031010.log

Criação de perfil de desempenho da VM do TPU

Para mais informações sobre como criar perfis de modelos na VM de TPU, consulte Criação de perfil de desempenho do PyTorch XLA (em inglês).

Pods

O PyTorch/XLA requer que todas as VMs de TPU acessem o código e os dados do modelo. Uma maneira fácil de fazer isso é usar o seguinte script de inicialização ao criar o pod de VM de TPU. Ele realiza o download dos dados em todas as VMs de TPU.

  1. No Cloud Shell, execute o seguinte comando para garantir que você esteja executando a versão mais atual do gcloud:

    $ sudo /opt/google-cloud-sdk/bin/gcloud components update
    
  2. Exporte as seguintes variáveis de ambiente necessárias para criar a VM do TPU do pod:

    $ export PROJECT_ID=project-id
    $ export TPU_NAME=tpu-name
    $ export ZONE=zone
    $ export RUNTIME_VERSION=v2-alpha
    
  3. Crie a VM da TPU

    $ gcloud alpha compute tpus tpu-vm create ${TPU_NAME} \
    --zone ${ZONE} --project ${PROJECT_ID} --accelerator-type v3-32 \
    --version ${RUNTIME_VERSION} --metadata startup-script='#! /bin/bash
    cd /usr/share/
    git clone --recursive https://github.com/pytorch/pytorch
    cd pytorch/
    git clone --recursive https://github.com/pytorch/xla.git
    EOF'
    

    Ao seguir essas instruções, execute cada comando iniciado por (vm)$ na janela de sessão da VM.

  4. ssh para qualquer worker de TPU, por exemplo, worker 0, verifique se o download de dados/modelos foi concluído (isso leva e inicia o treinamento depois de gerar as chaves SSH para ssh). entre os workers da VM em um pod.

    (vm)$ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone ${ZONE} --project ${PROJECT_ID}
    

    Atualize os metadados ssh do projeto:

    (vm)$ gcloud compute config-ssh
    
    (vm)$ export TPU_NAME=tpu-name
    
    (vm)$ python3 -m torch_xla.distributed.xla_dist \
       --tpu=${TPU_NAME} -- python3 /usr/share/pytorch/xla/test/test_train_mp_imagenet.py \
       --fake_data --model=resnet50 --num_epochs=1
    

    O treinamento leva cerca de três minutos. Após a conclusão, você verá uma mensagem semelhante a esta:

    Epoch 1 test end 23:49:15, Accuracy=100.00
    10.164.0.11 [0] Max Accuracy: 100.00%
    

    A geração de registros do servidor está em /tmp/xrt_server_log em cada worker.

       (vm)$ ls /tmp/xrt_server_log/
    
    server_20210401-031010.log
    

    Se você quiser reiniciar XRT_SERVER(caso o servidor esteja em um estado não íntegro), transmita --restart-tpuvm-pod-server ao executar xla_dist. Observe que as novas configurações do servidor XRT, como variáveis de ambiente, como LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4, não serão coletadas até que você reinicie o servidor.

  5. Após o fim do treinamento do ResNet, exclua a VM do TPU.

    (vm)$ exit
    
    $ gcloud alpha compute tpus tpu-vm delete tpu-name \
    --zone=zone
    
  6. Execute gcloud list para verificar se os recursos foram excluídos. A exclusão pode levar vários minutos. Uma resposta como esta indica que suas instâncias foram excluídas com sucesso.

    $ gcloud alpha compute tpus list --zone=zone
    
    Listed 0 items.
    

Pods com coordenador remoto

A prática recomendada é configurar um coordenador remoto a ser usado para iniciar o treinamento do pod. O coordenador remoto é uma VM padrão do Compute Engine, não uma VM de TPU. O coordenador remoto emite comandos para o pod da VM da TPU. A vantagem de usar um coordenador remoto é que, quando ocorrem eventos de manutenção de TPU, o treinamento pode parar e você pode perder o acesso ao pod da VM da TPU. No entanto, o coordenador remoto pode permitir que a VM de TPU se recupere automaticamente. para começar.

  1. Exporte as seguintes variáveis de ambiente necessárias para criar a VM do TPU:

    $ export TPU_NAME=tpu-name
    $ export ZONE=zone
    $ export PROJECT_ID=project-id
    
  2. Criar uma fração do pod da VM do TPU

    Crie uma fração do pod de TPU usando o comando gcloud. Por exemplo, para criar uma fração de pod v2-32, use o seguinte comando:

    $ gcloud alpha compute tpus tpu-vm create tpu-name \
    --zone europe-west4-a --project tpu-prod-env-one-vm --accelerator-type v3-32 \
    --version v2-alpha --metadata startup-script='#! /bin/bash
    cd /usr/share/
    git clone --recursive https://github.com/pytorch/pytorch
    cd pytorch/
    git clone --recursive https://github.com/pytorch/xla.git
    EOF'
    
  3. Exporte as seguintes variáveis de ambiente necessárias para criar a VM do coordenador remoto:

    $ export VM_NAME=vm-name
    $ export ZONE=zone
    
  4. Crie uma VM de coordenador remoto executando:

    (vm)$ gcloud compute --project=project-id instances create vm-name \
      --zone=zone  \
      --machine-type=n1-standard-1  \
      --image-family=torch-xla \
      --image-project=ml-images  \
      --boot-disk-size=200GB \
      --scopes=https://www.googleapis.com/auth/cloud-platform
    

    Quando o comando gcloud compute terminar a execução, verifique se o prompt do shell foi alterado de username@projectname para username@vm-name. Essa alteração mostra que você fez login na VM do coordenador remoto.

  5. ssh na instância do coordenador remoto:

    (vm)$ gcloud compute ssh vm-name --zone=zone
    
  6. Ative o ambiente torch-xla-1.8.1 e execute o treinamento a partir dele.

    (vm)$ gcloud compute config-ssh
    
    (vm)$ conda activate torch-xla-1.8.1
    (vm)$ python3 -m torch_xla.distributed.xla_dist --tpu=tpu-name --restart-tpuvm-pod --env LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4 -- python3 /usr/share/pytorch/xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1
    

    O treinamento leva cerca de três minutos para ser executado e gera uma mensagem semelhante a:

      Epoch 1 test end 23:19:53, Accuracy=100.00
      Max Accuracy: 100.00%
      

  7. Depois que o treinamento do ResNet terminar, saia da VM do TPU e exclua a VM do coordenação remoto e a VM do TPU.

    (vm)$ exit
    
    $ gcloud compute instances delete vm-name  \
      --zone=zone
    
    $ gcloud alpha compute tpus tpu-vm delete tpu-name \
      --zone zone