Ejecutar código PyTorch en slices de TPU

Antes de ejecutar los comandos de este documento, asegúrate de haber seguido las instrucciones que se indican en Configurar una cuenta y un proyecto de Cloud TPU.

Una vez que hayas ejecutado el código de PyTorch en una sola máquina virtual de TPU, puedes ampliarlo ejecutándolo en un sector de TPU. Las porciones de TPU son varias placas de TPU conectadas entre sí a través de conexiones de red específicas de alta velocidad. Este documento es una introducción a la ejecución de código de PyTorch en sectores de TPU.

Crear un slice de TPU de Cloud

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=europe-west4-b
   export ACCELERATOR_TYPE=v5p-32
   export RUNTIME_VERSION=v2-alpha-tpuv5

   Descripciones de las 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} \
       --zone=${ZONE} \
       --project=${PROJECT_ID} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --version=${RUNTIME_VERSION}

Instalar PyTorch/XLA en tu slice

Después de crear el segmento de TPU, debes instalar PyTorch en todos los hosts del segmento de TPU. Para ello, puedes usar el comando gcloud compute tpus tpu-vm ssh con los parámetros --worker=all y --commamnd.

Si los siguientes comandos fallan debido a un error de conexión SSH, puede deberse a que las VMs de TPU no tienen direcciones IP externas. 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.

1. Instala PyTorch/XLA en todos los trabajadores de la VM de TPU:

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
       --zone=${ZONE} \
       --project=${PROJECT_ID} \
       --worker=all \
       --command="pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html"

2. Clona XLA en todos los trabajadores de la VM de TPU:

   gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
       --zone=${ZONE} \
       --project=${PROJECT_ID} \
       --worker=all \
       --command="git clone https://github.com/pytorch/xla.git"

Ejecutar una secuencia de comandos de entrenamiento en tu segmento de TPU

Ejecuta la secuencia de comandos de entrenamiento en todos los trabajadores. La secuencia de comandos de entrenamiento usa una estrategia de partición de datos única para varios programas (SPMD). Para obtener más información sobre SPMD, consulta la guía de usuario de SPMD de PyTorch/XLA.

gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
   --zone=${ZONE} \
   --project=${PROJECT_ID} \
   --worker=all \
   --command="PJRT_DEVICE=TPU python3 ~/xla/test/spmd/test_train_spmd_imagenet.py  \
   --fake_data \
   --model=resnet50  \
   --num_epochs=1 2>&1 | tee ~/logs.txt"

La formación dura unos 15 minutos. Cuando se complete, debería aparecer un mensaje similar al siguiente:

Epoch 1 test end 23:49:15, Accuracy=100.00
     10.164.0.11 [0] Max Accuracy: 100.00%

Limpieza

Cuando hayas terminado de usar tu VM de TPU, sigue estos pasos para limpiar tus recursos.

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 tus recursos de TPU de Cloud.

   $ gcloud compute tpus tpu-vm delete  \
       --zone=${ZONE}

3. Para comprobar que los recursos se han eliminado, ejecuta gcloud compute tpus tpu-vm list. El proceso de eliminación puede tardar varios minutos. El resultado del siguiente comando no debe incluir ninguno de los recursos creados en este tutorial:

   $ gcloud compute tpus tpu-vm list --zone=${ZONE}