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

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

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

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.

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

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