Esta página foi traduzida pela API Cloud Translation.

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 da TPU e ssh nela, conforme descrito no Guia do usuário da VM do Cloud TPU.

Consulte Versões com suporte do PyTorch para uma lista das versões de software da TPU disponíveis para Pytorch/XLA.

Configuração básica

Defina a configuração do dispositivo XRT TPU:

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

   (vm)$ export TPU_NUM_DEVICES=4

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

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.13:

(tpuvm):$ cd /usr/share/
(tpuvm):$ sudo git clone -b release/1.13 --recursive https://github.com/pytorch/pytorch
(tpuvm):$ cd pytorch/
(tpuvm):$ sudo git clone -b r1.13 --recursive https://github.com/pytorch/xla.git
(tpuvm):$ cd xla/
(tpuvm):$ yes | sudo pip3 uninstall torch_xla
(tpuvm):$ yes | sudo pip3 uninstall torch
(tpuvm):$ yes | sudo pip3 uninstall torch_vision
(tpuvm):$ sudo pip3 install torch==1.13.0
(tpuvm):$ sudo pip3 install torchvision==0.14.0
(tpuvm):$ sudo pip3 install https://storage.googleapis.com/tpu-pytorch/wheels/tpuvm/torch_xla-1.13-cp38-cp38-linux_x86_64.whl
(tpuvm):$ sudo rm -rf /usr/local/lib/python3.8/dist-packages/libtpu*
(tpuvm):$ sudo pip3 install torch_xla[tpuvm]

Fazer um cálculo simples

Inicie o interpretador Python na VM de TPU:
```
(vm)$ python3
```

Importe os pacotes PyTorch a seguir:

import torch
import torch_xla.core.xla_model as xm

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

Configuração avançada

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

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.

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

(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

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