Guía del usuario de PyTorch/XLA de Cloud TPU

Importante :Puedes usar TPU con dos arquitecturas diferentes: TPU Node y TPU VM. En este instructivo, se supone que usas VM de TPU. Las dos arquitecturas de VM se describen en Arquitectura del sistema. A menos que sepas que necesitas usar nodos de TPU, te recomendamos usar las VM de TPU. Para obtener más información sobre el uso de PyTorch en pods de TPU con nodos TPU, consulta Entrena modelos de PyTorch en pods de Cloud TPU.

Ejecuta cargas de trabajo del AA con PyTorch/XLA

En esta guía, se explica cómo realizar un cálculo simple en una TPU v4 con PyTorch.

Configuración básica

1. Crea una VM de TPU con una TPU v4 que ejecute el entorno de ejecución de VM de TPU para Pytorch 2.0:

     gcloud compute tpus tpu-vm create your-tpu-name \
     --zone=us-central2-b \
     --accelerator-type=v4-8 \
     --version=tpu-vm-v4-pt-2.0

   Nota :Si deseas instalar una versión diferente de PyTorch, reemplaza tpu-vm-v4-pt-2.0 por la versión del entorno de ejecución de TPU que deseas usar. Para obtener una lista de los entornos de ejecución de TPU con PyTorch, consulta Versiones compatibles con PyTorch.

2. Conéctate a la VM de TPU con SSH:

     gcloud compute tpus tpu-vm ssh your-tpu-name \
     --zone=us-central2-b \
     --accelerator-type=v4-8

   Punto clave: En esta guía, el prefijo (vm) $ significa que debes ejecutar el comando en la instancia de VM de TPU.

3. Establece la configuración del dispositivo de TPU PJRT o XRT.

   Nota: Hay dos opciones de tiempo de ejecución de PyTorch/XLA: PJRT y XRT. Te recomendamos que uses PJRT, a menos que tengas un motivo para usar XRT. Para obtener más información sobre las diferentes configuraciones del entorno de ejecución y ver por qué tienes un motivo para usar XRT, consulta la documentación del entorno de ejecución de PJRT.

   PJRT

       (vm)$ export PJRT_DEVICE=TPU
    

   XRT

       (vm)$ export XRT_TPU_CONFIG="localservice;0;localhost:51011"
    

4. Si entrenas con Cloud TPU v4, configura también la siguiente variable de entorno:

     (vm)$ export TPU_NUM_DEVICES=4

Realiza un cálculo simple

1. Inicia el intérprete de Python en la VM de TPU:

   (vm)$ python3

2. Importa los siguientes paquetes de PyTorch:

   import torch
   import torch_xla.core.xla_model as xm

3. Ingresa la siguiente secuencia de comandos:

   dev = xm.xla_device()
   t1 = torch.randn(3,3,device=dev)
   t2 = torch.randn(3,3,device=dev)
   print(t1 + t2)

   Nota: Se espera el mensaje de error OpKernel ('op: "TPURoundRobin" device_type: "CPU"') for unknown op y se corregirá en una próxima versión.

   Se muestra el siguiente resultado:

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

Ejecuta Resnet en una TPU de un solo dispositivo

En este punto, puedes ejecutar cualquier código PyTorch / XLA que prefieras. Por ejemplo, puedes ejecutar un modelo ResNet con datos falsos:

(vm)$ git clone --recursive https://github.com/pytorch/xla.git
(vm)$ python3 xla/test/test_train_mp_imagenet.py --fake_data --model=resnet50 --num_epochs=1

El ejemplo de ResNet se entrena durante 1 ciclo de entrenamiento y tarda alrededor de 7 minutos. Muestra un resultado similar al siguiente:

Epoch 1 test end 20:57:52, Accuracy=100.00 Max Accuracy: 100.00%

Una vez que finalice el entrenamiento de ResNet, borra la VM de TPU.

(vm)$ exit

$ gcloud compute tpus tpu-vm delete tpu-name \
--zone=zone

La eliminación puede tardar varios minutos. Ejecuta gcloud compute tpus list --zone=${ZONE} para verificar que los recursos se hayan borrado.

Temas avanzados

Para los modelos que tienen asignaciones frecuentes y de gran tamaño, tcmalloc mejora el rendimiento en comparación con la función malloc del entorno de ejecución de C/C++. El malloc predeterminado que se usa en la VM de TPU es tcmalloc. Puedes forzar que el software de VM de TPU use la malloc estándar si anulas la configuración de la variable de entorno LD_PRELOAD:

   (vm)$ unset LD_PRELOAD

En los ejemplos anteriores (el cálculo simple y ResNet50), el programa PyTorch/XLA inicia el servidor XRT local en el mismo proceso que el intérprete de Python. También puedes optar por iniciar el servicio local de XRT en un proceso independiente:

(vm)$ python3 -m torch_xla.core.xrt_run_server --port 51011 --restart

Nota: El puerto que especificaste antes es el mismo que especificaste en el comando de export XRT_TPU_CONFIG="localservice;0;localhost:51011" antes.

La ventaja de este enfoque es que la caché de compilación persiste en todas las ejecuciones de entrenamiento. Cuando ejecutas el servidor de XLA en un proceso independiente, la información de registro del servidor se escribe en /tmp/xrt_server_log.

(vm)$ ls /tmp/xrt_server_log/

server_20210401-031010.log

Creación de perfiles de rendimiento de VM de TPU

Si deseas obtener más información para generar perfiles de tus modelos en la VM de TPU, consulta la siguiente información: Creación de perfiles de rendimiento de PyTorch XLA.

Ejemplos de pods de TPU de PyTorch/XLA

Consulta Pod de VM de PyTorch TPU a fin de obtener información de configuración y ejemplos para ejecutar PyTorch/XLA en un pod de VM de TPU.

Docker en la VM de TPU

En esta sección, se muestra cómo ejecutar Docker en una VM de TPU con PyTorch/XLA preinstalado.

Nota :En este momento, la ejecución de Docker en una VM de TPU solo es compatible con TPU de dispositivo único (v2-8, v3-8).

Imágenes de Docker disponibles

Puedes consultar el archivo README de GitHub para encontrar todas las imágenes de Docker de VM de TPU disponibles.

Ejecuta imágenes de Docker en la VM de TPU

Nota: El último resultado del comando, root:/#, es el mensaje que ves después de ejecutar la imagen de Docker.

(tpuvm): sudo docker pull gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm
(tpuvm): sudo docker run --privileged  --shm-size 16G --name tpuvm_docker -it -d  gcr.io/tpu-pytorch/xla:nightly_3.8_tpuvm
(tpuvm): sudo docker exec --privileged -it tpuvm_docker /bin/bash
(pytorch) root:/#

Nota: Debes ejecutar el contenedor mediante la marca --privileged para asegurarte de que los dispositivos de TPU sean visibles dentro del contenedor. Hay un entorno conda pytorch preinstalado en este contenedor. Cada comando que comienza con (pytorch) debe ejecutarse dentro del contenedor.

Verificar libtpu

Nota:libtpu debe estar preinstalado en imágenes de Docker de TPUVM más nuevas, pero siempre es mejor verificarlo de forma manual.

Para verificar que libtpu esté instalada, ejecuta lo siguiente:

(pytorch) root:/# ls /root/anaconda3/envs/pytorch/lib/python3.8/site-packages/ | grep libtpu

Esto debería generar un resultado similar al siguiente:

libtpu
libtpu_nightly-0.1.dev20220518.dist-info

Si no se muestran resultados, puedes instalar la libtpu correspondiente de forma manual con el siguiente comando:

(pytorch) root:/# pip install torch_xla[tpuvm]

Verificar tcmalloc

tcmalloc es el malloc predeterminado que usamos en la VM de TPU. Para obtener más información, consulta esta sección. Esta biblioteca debe instalarse previamente en imágenes de Docker de VM de TPU más nuevas, pero es mejor verificarla de forma manual. Puedes ejecutar el siguiente comando para verificar que la biblioteca esté instalada.

(pytorch) root:/# echo $LD_PRELOAD

Esto debería generar un resultado similar al siguiente:

/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4

Si no se configura LD_PRELOAD, puedes ejecutarlo de forma manual:

(pytorch) root:/# sudo apt-get install -y google-perftools
(pytorch) root:/# export LD_PRELOAD="/usr/lib/x86_64-linux-gnu/libtcmalloc.so.4"

Verificar dispositivo

Para verificar que el dispositivo de VM de TPU esté disponible, ejecuta lo siguiente:

(pytorch) root:/# ls /dev | grep accel

Esto debería generar los siguientes resultados.

accel0
accel1
accel2
accel3

Si no se muestran resultados, es probable que no hayas iniciado el contenedor con la marca --privileged.

Ejecuta un modelo

Para verificar si el dispositivo de VM de TPU está disponible, ejecuta lo siguiente:

(pytorch) root:/# export XRT_TPU_CONFIG="localservice;0;localhost:51011"
(pytorch) root:/# python3 pytorch/xla/test/test_train_mp_imagenet.py --fake_data --num_epochs 1