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

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

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

En este documento se ofrece 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 TPUs de un solo host solo tienen una VM de TPU. Para ejecutar código en TPUs con más de una VM de TPU (por ejemplo, v5litepod-32 o una versión posterior), consulta Ejecutar código PyTorch en sectores de TPU de Cloud.

Antes de empezar

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

Roles obligatorios

Para obtener los permisos que necesitas para crear una TPU y conectarte a ella mediante SSH, pide a tu administrador que te conceda los siguientes roles de gestión de identidades y accesos en tu proyecto:

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

Para obtener más información sobre cómo conceder roles, consulta el artículo Gestionar acceso a proyectos, carpetas y organizaciones.

También puedes conseguir los permisos necesarios a través de roles personalizados u otros roles predefinidos.

Crear una TPU de Cloud 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 variables de entorno

   Variable	Descripción
   PROJECT_ID	El ID de tu proyecto Google Cloud . Usa un proyecto que ya tengas o crea uno.
   TPU_NAME	El nombre de la TPU.
   ZONE	La zona en la que se creará la VM de TPU. Para obtener más información sobre las zonas admitidas, consulta Regiones y zonas de TPU.
   ACCELERATOR_TYPE	El tipo de acelerador especifica la versión y el tamaño de la TPU de Cloud que quieres 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	La versión de software de la TPU de Cloud.

2. Crea tu VM de TPU ejecutando el siguiente comando:

   $ 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 TPU de Cloud

Conéctate a tu VM de TPU mediante 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 mediante SSH, puede deberse a que la VM de TPU no tiene 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 Conectarse a una VM de TPU sin una dirección IP pública.

Instalar PyTorch/XLA en tu máquina virtual 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] -f https://storage.googleapis.com/libtpu-releases/index.html

Verificar que PyTorch puede acceder a las TPUs

Usa el siguiente comando para verificar que PyTorch puede acceder a tus TPUs:

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

La salida del comando debería tener el siguiente aspecto:

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

Hacer un cálculo básico

1. Crea un archivo llamado tpu-test.py en el directorio actual y copia y pega el siguiente script 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 del script 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')
   

Limpieza

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

1. Desconéctate de la instancia de TPU de Cloud, si aún no lo has hecho:

   (vm)$ exit

   A continuación, se mostrará el mensaje username@projectname, que indica que estás en Cloud Shell.

2. Elimina tu TPU de Cloud.

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

3. Para comprobar que los recursos se han eliminado, ejecuta el siguiente comando. Comprueba que tu TPU ya no aparezca en la lista. El proceso de eliminación puede tardar varios minutos.

   $ gcloud compute tpus tpu-vm list \
       --zone=$ZONE

Siguientes pasos

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

• Ejecutar código de PyTorch en sectores de TPU
• Gestionar TPUs
• Arquitectura del sistema de TPU de Cloud
• Documentación de PyTorch/XLA