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

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

Nota: Esta página se aplica a la API de Cloud TPU. Para Ironwood (TPU7x), debes usar Google Kubernetes Engine (GKE). Para obtener más información, consulta Acerca de las TPU en GKE.

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

Nota: En este ejemplo, se muestra cómo ejecutar código en una TPU v5litepod-8 (v5e), que es una TPU de un solo host. Las TPU de host único solo tienen 1 VM de TPU. Para ejecutar código en TPU con más de una VM de TPU (por ejemplo, v5litepod-32 o más grande), consulta Ejecuta código de PyTorch en porciones de 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.

Roles requeridos

Si deseas obtener los permisos que necesitas para crear una TPU y conectarte a ella con SSH, pídele a tu administrador que te otorgue los siguientes roles de IAM en tu proyecto:

• Administrador de TPU (roles/tpu.admin)
• Usuario de la cuenta de servicio (roles/iam.serviceAccountUser)
• Visualizador de Compute (roles/compute.viewer)

Para obtener más información sobre cómo otorgar roles, consulta Administra el acceso a proyectos, carpetas y organizaciones.

También puedes obtener los permisos necesarios a través de roles personalizados o cualquier otro rol predefinido.

Crea una Cloud TPU con gcloud

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

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

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

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

2. Borra tu Cloud TPU.

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

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

• Ejecuta código de PyTorch en porciones de TPU
• Administra TPU
• Arquitectura de sistemas de Cloud TPU
• Documentación de PyTorch/XLA