Ejecuta un cálculo en una VM de Cloud TPU con PyTorch

En este documento, se proporciona una breve introducción sobre cómo trabajar con PyTorch y Cloud TPU.

Antes de comenzar

Antes de ejecutar los comandos de este documento, debes crear una cuenta de Google Cloud , instalar Google Cloud CLI y configurar el comando gcloud. Para obtener más información, consulta Configura el entorno de Cloud TPU.

Crea una Cloud TPU con `gcloud`

Define algunas variables de entorno para que los comandos sean más fáciles de usar.

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

Descripciones de las variables de entorno

Variable	Descripción
`PROJECT_ID`	Es el ID de tu proyecto de Google Cloud . Usa un proyecto existente o crea uno nuevo.
`TPU_NAME`	Es el nombre de la TPU.
`ZONE`	Es la zona en la que se creará la VM de TPU. Para obtener más información sobre las zonas compatibles, consulta Regiones y zonas de TPU.
`ACCELERATOR_TYPE`	El tipo de acelerador especifica la versión y el tamaño de Cloud TPU que deseas crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.
`RUNTIME_VERSION`	Es la versión de software de Cloud TPU.

Ejecuta el siguiente comando para ejecutar tu VM de TPU:

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

Conéctate a la VM de tu Cloud TPU

Conéctate a tu VM de TPU a través de SSH con el siguiente comando:

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

Si no puedes conectarte a una VM de TPU con SSH, es posible que la VM de TPU no tenga una dirección IP externa. Para acceder a una VM de TPU sin una dirección IP externa, sigue las instrucciones que se indican en Conéctate a una VM de TPU sin una dirección IP pública.

Instala PyTorch/XLA en la VM de tu 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

Verifica que PyTorch pueda acceder a las TPU

Usa el siguiente comando para verificar que PyTorch pueda acceder a tus TPU:

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

El resultado del comando debería verse de la siguiente manera:

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

Realiza un cálculo básico

Crea un archivo llamado tpu-test.py en el directorio actual y copia y pega la siguiente secuencia de comandos en él:

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)

Ejecuta la secuencia de comandos:

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

El resultado de la secuencia de comandos muestra el resultado del cálculo:

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

Realiza una limpieza

Sigue estos pasos para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos que usaste en esta página.

Desconéctate de la instancia de Cloud TPU, si aún no lo hiciste:
```
(vm)$ exit
```
El mensaje ahora debería mostrar username@projectname, que indica que estás en Cloud Shell.

Borra tu Cloud TPU.

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

Ejecuta el siguiente comando para verificar que los recursos se hayan borrado. Asegúrate de que tu TPU ya no aparezca en la lista. Este proceso puede tardar varios minutos.
```
$ gcloud compute tpus tpu-vm list \
    --zone=$ZONE
```

¿Qué sigue?

Obtén más información sobre las VMs de Cloud TPU: