Guía del usuario de PyTorch/XLA 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
  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
  3. Establece la configuración del dispositivo de TPU PJRT o XRT.

    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)

    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

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.

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

(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:/#

Verificar libtpu

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