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

En este documento, se proporciona una breve introducción al trabajo 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

  1. Define algunas variables de entorno para facilitar el uso de los comandos.

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

    Descripciones de las variables de entorno

    PROJECT_ID
    El ID Google Cloud de tu proyecto.
    ACCELERATOR_TYPE
    El tipo de acelerador especifica la versión y el tamaño de la 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.
    ZONE
    Es la zona en la que deseas crear la Cloud TPU.
    RUNTIME_VERSION
    La versión del entorno de ejecución de Cloud TPU.
    TPU_NAME
    El nombre asignado por el usuario a tu Cloud TPU.

  2. Ejecuta el siguiente comando para crear 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 tu VM de 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

Instala PyTorch/XLA en tu VM de 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]~=2.5.0 -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']

Realiza un cálculo básico

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

  2. Ejecuta la secuencia de comandos:

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

    El resultado de la secuencia de comandos muestra el resultado del procesamiento:

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

Limpia

Para evitar que se apliquen cargos a tu Google Cloud cuenta por los recursos que usaste en esta página, sigue estos pasos.

  1. Desconéctate de la instancia de Compute Engine, si aún no lo hiciste:

    (vm)$ exit

    El mensaje ahora debería mostrar username@projectname, que indica que estás en Cloud Shell.

  2. Borra tu Cloud TPU.

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

El resultado de este comando debería confirmar que se borró tu TPU.

¿Qué sigue?

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