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

Importante: é possível usar TPUs usando duas arquiteturas diferentes: nós e VMs de TPU. Neste tutorial, presumimos que você esteja usando VMs de TPU. As duas arquiteturas de VM são descritas em Arquitetura do sistema. A menos que você saiba que precisa de nós de TPU, recomendamos o uso de VMs de TPU. Para mais informações sobre como usar o PyTorch em pods de TPU usando nós de TPU, consulte Como treinar modelos PyTorch em pods da Cloud TPU.

Executar cargas de trabalho de ML com o PyTorch/XLA

Veja neste guia como realizar um cálculo simples em uma TPU v4 usando o PyTorch.

Configuração básica

1. Crie uma VM de TPU com uma TPU v4 que execute o ambiente de execução da VM de TPU para Pytorch 2.0:

     gcloud compute tpus tpu-vm create your-tpu-name \
     --zone=us-central2-b \
     --accelerator-type=v4-8 \
     --version=tpu-vm-v4-pt-2.0

   Observação :se você quiser instalar uma versão diferente do PyTorch, substitua tpu-vm-v4-pt-2.0 pela versão do ambiente de execução da TPU que quer usar. Para ver uma lista de ambientes de execução da TPU com o PyTorch, consulte Versões compatíveis com PyTorch.

2. Conecte-se à VM da TPU usando SSH:

     gcloud compute tpus tpu-vm ssh your-tpu-name \
     --zone=us-central2-b \
     --accelerator-type=v4-8

   Importante: neste guia, o prefixo (vm) $ significa que você precisa executar o comando na instância da VM da TPU.

3. Defina a configuração do dispositivo TPU PJRT ou XRT.

   Observação :há duas opções de ambiente de execução do PyTorch/XLA: PJRT e XRT. Recomendamos que você use o PJRT, a menos que tenha um motivo para usar o XRT. Para saber mais sobre as diferentes configurações de ambiente de execução e ver por que usar o XRT, consulte a documentação do ambiente de execução do PJRT.

   JRT

       (vm)$ export PJRT_DEVICE=TPU
    

   XRT

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

4. Se você estiver treinando com a Cloud TPU v4, também defina a seguinte variável de ambiente:

     (vm)$ export TPU_NUM_DEVICES=4

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)

   Observação: a mensagem de erro OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op é conhecida e será corrigida em uma versão futura.

   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 compute tpus tpu-vm delete tpu-name \
--zone=zone

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

Tópicos avançados

Para modelos que têm alocações frequentes e consideráveis, o tcmalloc melhora o desempenho em comparação com a função de ambiente de execução C/C++ malloc. O malloc padrão usado na VM de TPU é tcmalloc. É possível forçar o software de VM da TPU a usar malloc padrão, basta desmarcar a variável de ambiente LD_PRELOAD:

   (vm)$ unset LD_PRELOAD

Nos exemplos anteriores (o cálculo simples e o ResNet50), o programa PyTorch/XLA inicia 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

Observação :a porta especificada anteriormente é a mesma que você especificou no comando export XRT_TPU_CONFIG="localservice;0;localhost:51011" anteriormente.

A vantagem dessa abordagem é que o cache de compilação persiste entre as execuções de 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.

Exemplos de pods de TPU do PyTorch/XLA

Consulte Pod de VMs da TPU do PyTorch para informações de configuração e exemplos para executar o PyTorch/XLA em um pod de VM de TPU.

Docker na VM da TPU

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

Observação: no momento, a execução do Docker na VM da TPU é compatível apenas com TPUs de dispositivo único (v2-8, v3-8).

Imagens do Docker disponíveis

Consulte o README (link em inglês) do GitHub para encontrar todas as imagens do Docker de VM de TPU disponíveis.

Executar imagens do Docker na VM da TPU

Observação: a saída do último comando, root:/#, é a solicitação que você vê depois da execução da imagem do Docker.

(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:/#

Observação: é preciso executar o contêiner usando a sinalização --privileged para garantir que os dispositivos TPU estejam visíveis dentro do contêiner. Há um ambiente conda pytorch pré-instalado neste contêiner. Todos os comandos que começam com (pytorch) precisam ser executados no contêiner.

Verificar libtpu

Observação:libtpu precisa ser pré-instalado em imagens mais recentes do Docker da TPUVM, mas é sempre melhor verificá-lo manualmente.

Para verificar se o libtpu está instalado, execute:

(pytorch) root:/# ls /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/ | grep libtpu

A saída vai ser semelhante a esta:

libtpu
libtpu_nightly-0.1.dev20220518.dist-info

Se nenhum resultado for exibido, instale 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 precisa ser pré-instalada em imagens mais recentes do Docker da VM da TPU, mas é sempre melhor verificá-la manualmente. Execute o comando a seguir para verificar se a biblioteca está instalada.

(pytorch) root:/# echo $LD_PRELOAD

O resultado será semelhante a este:

/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 de TPU está disponível executando:

(pytorch) root:/# ls /dev | grep accel

Isso gera os resultados abaixo.

accel0
accel1
accel2
accel3

Se nenhum resultado for exibido, é provável que você não tenha iniciado o contêiner com a sinalização --privileged.

Executar um modelo

Verifique se o dispositivo de VM de 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