Guia do usuário do Cloud TPU com PyTorch/XLA

Executar cargas de trabalho de ML com o PyTorch/XLA

Antes de iniciar os procedimentos neste guia, configure uma VM de TPU e ssh, conforme descrito em Preparar um projeto do Google Cloud.

Há duas versões de software de TPU disponíveis para Pytorch / XLA:

  1. A imagem mais recente disponível: tpu-vm-pt-1.11
  2. tpu-vm-pt-1.10
  3. v2-alpha

A imagem tpu-vm-pt-1.11 é pré-instalada com PyTorch 1.11.0 e Pytorch / XA 1.11.0. A imagem tpu-vm-pt-1.10 é pré-instalada com PyTorch 1.10.0 e Pytorch / XA 1.10.0. A imagem v2-alpha é pré-instalada com o PyTorch 1.8.1.

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 significativas e frequentes, o tcmalloc melhora o desempenho em comparação com a função de ambiente de execução malloc do C/C++. O malloc padrão usado na VM de TPU é tcmalloc. É possível forçar o software de VM da TPU a usar o malloc padrão desativando 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 do PyTorch pré-instalada nas VMs de TPU, instale a versão que você quer usar. Por exemplo, se você quiser usar 1.9:

(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 de 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 do PyTorch / XLA. Por exemplo, é possível executar um modelo do 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 um período e leva cerca de sete minutos. Ele retorna um resultado semelhante a este:

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 de 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 anteriores (cálculo simples e ResNet50), o programa PyTorch/XLA inicia o servidor XRT local no mesmo processo que o interpretador 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 persiste entre as execuções do treinamento. Ao executar o servidor XLA em um processo separado, as informações de registro 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 de TPU

Para mais informações sobre como criar perfis dos modelos na VM de 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. Use o script de inicialização a seguir ao criar o pod da VM da TPU. Ele faz o download dos dados para todas as VMs da TPU.

  1. No Cloud Shell, execute o seguinte comando para verificar se você está executando a versão mais recente de 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=tpu-vm-pt-1.11
    
  3. Criar a VM de 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 shell de VM de TPU.

  4. ssh para qualquer worker de TPU, por exemplo, worker 0. Verifique se o download dos dados/modelos foi concluído. No entanto, o download é 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 estiver na saída do comando ls, o processamento de arquivos estará concluído.

    Quando o processamento de arquivos for concluído, gere chaves ssh para o 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
    

    O treinamento leva cerca de três minutos. Quando ele 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%
    

    Há uma geração de registros no lado do servidor em /tmp/xrt_server_log em cada worker.

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

    Se você quiser reiniciar o XRT_SERVER (caso o servidor esteja em um estado não íntegro), transmita --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ó entram em vigor quando você reinicia o servidor.

  5. Após o término do treinamento do ResNet, exclua a VM de 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 usar um coordenador remoto para permitir que o Pod recupere automaticamente quando houver um evento de manutenção de 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 tpu-vm-pt-1.11 --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, que são 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.11 e execute o treinamento nele.

    (vm)$ gcloud compute config-ssh
    
    (vm)$ conda activate torch-xla-1.11
    (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 esta:

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

  7. Após terminar o treinamento do ResNet, saia da VM de TPU e exclua a VM do coordenador remoto e a VM de TPU.

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

Docker na VM da TPU

Nesta seção, mostramos como executar o Docker em VMs de TPU com o PyTorch/XLA pré-instalado.

Imagens do Docker disponíveis

Consulte o README do GitHub para encontrar todas as imagens do docker da VM de TPU disponíveis.

Executar imagens do Docker na VM da TPU

(tpuvm): sudo docker pull gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm
(tpuvm): sudo docker run --privileged  --shm-size 16G --name tpuvm_docker -it -d  gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm
(tpuvm): sudo docker exec --privileged -it tpuvm_docker /bin/bash
(pytorch) root:/#

Verificar libtpu

Para verificar se o libtpu está instalado, execute:

(pytorch) root:/# ls /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/ | grep libtpu
Isso gera uma saída semelhante a esta:
libtpu
libtpu_nightly-0.1.dev20220518.dist-info

Se nenhum resultado for exibido, instale manualmente o libtpu correspondente usando:

(pytorch) root:/# pip install torch_xla[tpuvm]

Verificar tcmalloc

tcmalloc é a malloc padrão que usamos na VM da TPU. Para mais informações, leia esta seção. Essa biblioteca deve ser pré-instalada em imagens mais recentes da VM do TPU, mas é sempre melhor verificá-la manualmente. Você pode executar o comando a seguir para verificar se a biblioteca está instalada.

(pytorch) root:/# echo $LD_PRELOAD
Isso gera uma saída semelhante a esta:
/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4

Se LD_PRELOAD não estiver definido, será possível executar manualmente:

(pytorch) root:/# sudo apt-get install -y google-perftools
(pytorch) root:/# export LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4"

Verificar dispositivo

Verifique se o dispositivo de VM da TPU está disponível executando:

(pytorch) root:/# ls /dev | grep accel
Isso deve gerar os seguintes resultados.
accel0
accel1
accel2
accel3

Se nenhum resultado for exibido, provavelmente você não iniciou o contêiner com a sinalização --privileged.

Executar um modelo

É possível verificar se o dispositivo de VM da TPU está disponível executando:

(pytorch) root:/# export XRT_TPU_CONFIG="localservice;0;localhost:51011"
(pytorch) root:/# python3 pytorch/xla/test/test_train_mp_imagenet.py --fake_data --num_epochs 1