Guia do usuário do Cloud TPU PyTorch/XLA

Execute cargas de trabalho de ML com o PyTorch/XLA

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

O PyTorch 1.8.1 e o PyTorch/XLA 1.8.1 estã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"

Para modelos com alocações frequentes e grandes, o tcmalloc melhora o desempenho significativamente em comparação com a função de ambiente de execução C/C++ malloc. O malloc padrão usado na VM da TPU é tcmalloc. É possível forçar o ambiente de execução da VM da TPU a usar o malloc padrão desmarcando a variável de ambiente LD_PRELOAD:

   (vm)$ unset LD_PRELOAD

Como alterar a versão do PyTorch

Se você não quiser usar a versão PyTorch pré-instalada nas VMs da TPU, instale a versão que quer usar. Exemplo:

(vm)$ sudo bash /var/scripts/docker-login.sh
(vm)$ sudo docker rm libtpu || true
(vm)$ sudo docker create --name libtpu gcr.io/cloud-tpu-v2-images/libtpu:pytorch-1.9 "/bin/bash"
(vm)$ sudo docker cp libtpu:libtpu.so /lib
(vm)$ sudo pip3 uninstall --yes torch torch_xla torchvision
(vm)$ sudo pip3 install torch==1.9.0
(vm)$ sudo pip3 install torchvision==0.10.0
(vm)$ sudo pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-1.9-cp38-cp38-linux_x86_64.whl

Fazer um cálculo simples

  1. Inicie o interpretador Python na VM da TPU:

    (vm)$ python3
    
  2. Importe os pacotes PyTorch a seguir:

    import torch
    import torch_xla.core.xla_model as xm
    
  3. Digite o seguinte script:

    dev = xm.xla_device()
    t1 = torch.randn(3,3,device=dev)
    t2 = torch.randn(3,3,device=dev)
    print(t1 + t2)
    

    Você verá a seguinte resposta:

    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

Se necessário, execute qualquer código PyTorch / XLA. Por exemplo, é possível executar um modelo ResNet com dados falsos:

(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 uma época e leva cerca de sete minutos. Ele retorna uma saída semelhante a esta:

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

Após o término do treinamento do ResNet, exclua a VM da TPU.

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

A exclusão pode levar vários minutos. Execute gcloud alpha compute tpus list --zone=${ZONE} para verificar se os recursos foram excluídos.

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 que o interpretador do Python. 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á nas execuções de treinamento. Ao executar o servidor XLA em um processo separado, as informações de geração de registros do lado do servidor são gravadas em /tmp/xrt_server_log.

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

Caracterização de perfil de desempenho da VM da TPU

Para mais informações sobre como criar perfis dos modelos na VM da TPU, consulte criação de perfil de desempenho do PyTorch XLA.

Pods

O PyTorch/XLA exige que todas as VMs de TPU tenham acesso ao código e aos dados do modelo. Uma maneira de fazer isso é usar o script de inicialização a seguir ao criar o pod de VM da TPU. Executa o download de 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:

    $ 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'
    

    Conforme você continuar as instruções, execute cada comando que começa com (vm)$ no shell da VM da TPU.

  4. ssh para qualquer worker da TPU, por exemplo, worker 0. Verifique se o download dos dados/modelos foi concluído. O download geralmente é muito rápido, mas o processamento de arquivos pode levar vários minutos. Para verificar se o processamento do arquivo foi concluído, execute ls -al no diretório de download. O download dos arquivos é feito na ordem especificada no script de inicialização. Exemplo:

    download file1
    download file2
    download file3
    

    Se file3 foi o último arquivo no script de inicialização e é exibido pelo comando ls, o processamento do arquivo foi concluído.

    Quando o processamento de arquivos for concluído, gere ssh-keys para ssh entre os workers da VM no pod. Em seguida, inicie o treinamento.

    (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
    

    Ele leva cerca de três minutos. Quando ela for concluída, 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%
    

    É possível encontrar a geração de registros do servidor em /tmp/xrt_server_log em cada worker.

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

    Se você quiser reiniciar a XRT_SERVER (caso o servidor esteja em um estado não íntegro), poderá transmitir --restart-tpuvm-pod-server ao executar xla_dist. As novas configurações do servidor XRT, como variáveis de ambiente, como LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4, só entrarão em vigor quando você reiniciar o servidor.

  5. Após o término do treinamento do ResNet, exclua a VM da TPU.

    (vm)$ exit
    
    $ gcloud alpha compute tpus tpu-vm delete tpu-name \
     --zone=zone
    
  6. Execute gcloud alpha compute tpus list --zone=${ZONE} para verificar se os recursos foram excluídos. A exclusão pode levar vários minutos.

Pods com coordenador remoto

Recomendamos o uso de um coordenador remoto para permitir que o pod recupere automaticamente quando ocorrer um evento de manutenção de uma TPU. Um coordenador remoto é uma VM padrão do Compute Engine, não uma VM de TPU. Ele emite comandos para as VMs de TPU no pod.

  1. Exporte as seguintes variáveis de ambiente:

    $ export TPU_NAME=tpu-name
    $ export ZONE=zone
    $ export PROJECT_ID=project-id
    
  2. Crie uma fração de pod de VM de TPU:

    $ 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 seu treinamento nele.

    (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
    

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

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

  7. Após o término do treinamento do ResNet, saia da VM da TPU e exclua a VM do coordenador remoto e a VM da TPU.

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