Ejecuta el código de PyTorch en porciones de pod de TPU

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

Una vez que tu código de PyTorch se ejecute en una sola VM de TPU, puedes escalarlo de forma vertical en una porción de pod de TPU. Las porciones de pod de TPU son varios paneles de TPU conectados entre sí en conexiones de red dedicadas de alta velocidad. Este documento es una introducción a la ejecución de código PyTorch en porciones de pod de TPU.

Crea una porción de pod de Cloud TPU

  1. Define algunas variables de entorno para facilitar el uso de los comandos.

    export PROJECT_ID=your-project
export ACCELERATOR_TYPE=v5p-32
export ZONE=europe-west4-b
export RUNTIME_VERSION=v2-alpha-tpuv5
export TPU_NAME=your-tpu-name

    Descripciones de las variables de entorno

    PROJECT_ID
    El ID Google Cloud de tu proyecto.
    ACCELERATOR_TYPE
    El tipo de acelerador especifica la versión y el tamaño de la 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.
    ZONE
    Es la zona en la que deseas crear la Cloud TPU.
    RUNTIME_VERSION
    La versión del software de Cloud TPU.
    TPU_NAME
    El nombre asignado por el usuario a tu Cloud TPU.

  2. Ejecuta el siguiente comando para crear tu VM de TPU:

    $ gcloud compute tpus tpu-vm create $TPU_NAME \
--zone=$ZONE \
--project=$PROJECT_ID \
--accelerator-type=$ACCELERATOR_TYPE \
--version=$RUNTIME_VERSION

Instala PyTorch/XLA en tu porción de pod

Después de crear la porción de pod de TPU, debes instalar PyTorch en todos los hosts en la porción de pod de TPU. Puedes hacerlo con el comando gcloud compute tpus tpu-vm ssh usando los parámetros --worker=all y --commamnd.

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

Ejecuta una secuencia de comandos de entrenamiento en tu porción de pod de TPU

Ejecuta la secuencia de comandos de entrenamiento en todos los trabajadores. La secuencia de comandos de entrenamiento usa una estrategia de fragmentación de datos múltiples de un solo programa (SPMD). Para obtener más información sobre SPMD, consulta la Guía del 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"

El entrenamiento dura aproximadamente 5 minutos. Cuando se complete, deberías ver un mensaje similar al siguiente:

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

Limpia

Cuando termines de usar la VM de TPU, sigue estos pasos para limpiar los recursos.

  1. Desconéctate de la instancia de Compute Engine, 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 tus recursos de Cloud TPU y Compute Engine.

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

  3. Ejecuta gcloud compute tpus execution-groups list para verificar que los recursos se hayan borrado. La eliminación puede tardar varios minutos. El resultado del siguiente comando no debe incluir ninguno de los recursos creados en este instructivo:

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