Executar um cálculo em uma VM da Cloud TPU usando o PyTorch

Criar um Cloud TPU usando gcloud

1. Defina algumas variáveis de ambiente para facilitar o uso dos comandos.

   export PROJECT_ID=your-project-id
   export TPU_NAME=your-tpu-name
   export ZONE=us-east5-a
   export ACCELERATOR_TYPE=v5litepod-8
   export RUNTIME_VERSION=v2-alpha-tpuv5-lite

   Descrições das variáveis de ambiente

   Variável	Descrição
   PROJECT_ID	O ID do seu Google Cloud projeto. Use um projeto existente ou crie um novo.
   TPU_NAME	O nome da TPU.
   ZONE	A zona em que a VM TPU será criada. Para mais informações sobre as zonas com suporte, consulte Regiões e zonas de TPU.
   ACCELERATOR_TYPE	O tipo de acelerador especifica a versão e o tamanho da Cloud TPU que você quer criar. Para mais informações sobre os tipos de aceleradores compatíveis com cada versão de TPU, consulte Versões de TPU.
   RUNTIME_VERSION	A versão do software do Cloud TPU.

2. Crie a VM da TPU executando o seguinte comando:

   $ gcloud compute tpus tpu-vm create $TPU_NAME \
       --project=$PROJECT_ID \
       --zone=$ZONE \
       --accelerator-type=$ACCELERATOR_TYPE \
       --version=$RUNTIME_VERSION

Conecte-se à VM do Cloud TPU

Conecte-se à VM da TPU por SSH usando o seguinte comando:

$ gcloud compute tpus tpu-vm ssh $TPU_NAME \
    --project=$PROJECT_ID \
    --zone=$ZONE

Se você não conseguir se conectar a uma VM da TPU usando SSH, talvez seja porque a VM da TPU não tem um endereço IP externo. Para acessar uma VM da TPU sem um endereço IP externo, siga as instruções em Conectar-se a uma VM da TPU sem um endereço IP público.

Instalar o PyTorch/XLA na VM da TPU

$ (vm) sudo apt-get update
$ (vm) sudo apt-get install libopenblas-dev -y
$ (vm) pip install numpy
$ (vm) pip install torch torch_xla[tpu] -f https://storage.googleapis.com/libtpu-releases/index.html

Verificar se o PyTorch pode acessar TPUs

Use o comando a seguir para verificar se o PyTorch pode acessar suas TPUs:

$ (vm) PJRT_DEVICE=TPU python3 -c "import torch_xla.core.xla_model as xm; print(xm.get_xla_supported_devices(\"TPU\"))"

A saída do comando vai ficar assim:

['xla:0', 'xla:1', 'xla:2', 'xla:3', 'xla:4', 'xla:5', 'xla:6', 'xla:7']

Fazer um cálculo básico

1. Crie um arquivo chamado tpu-test.py no diretório atual e copie e cole o seguinte script nele:

   import torch
   import torch_xla.core.xla_model as xm
   
   dev = xm.xla_device()
   t1 = torch.randn(3,3,device=dev)
   t2 = torch.randn(3,3,device=dev)
   print(t1 + t2)
   

2. Execute o script:

   (vm)$ PJRT_DEVICE=TPU python3 tpu-test.py

   A saída do script mostra o resultado do cálculo:

   tensor([[-0.2121,  1.5589, -0.6951],
           [-0.7886, -0.2022,  0.9242],
           [ 0.8555, -1.8698,  1.4333]], device='xla:1')