Ejecuta el código de PyTorch en porciones de pod de TPU
PyTorch/XLA requiere que todas las VM de la TPU puedan acceder al código y a los datos del modelo. Puedes usar una secuencia de comandos de inicio para descargar el software necesario para distribuir los datos del modelo a todas las VMs de TPU.
Si conectas tus VMs de TPU a una nube privada virtual (VPC), debes agregar una regla de firewall en tu proyecto para permitir la entrada de los puertos del 8470 al 8479. Para obtener más información sobre cómo agregar reglas de firewall, consulta Usa reglas de firewall.
Configure su entorno
En Cloud Shell, ejecuta el siguiente comando para asegurarte de que que ejecuta la versión actual de
gcloud
:$ gcloud components update
Si necesitas instalar
gcloud
, usa el siguiente comando:$ sudo apt install -y google-cloud-sdk
Crea algunas variables de entorno:
$ export PROJECT_ID=project-id $ export TPU_NAME=tpu-name $ export ZONE=us-central2-b $ export RUNTIME_VERSION=tpu-ubuntu2204-base $ export ACCELERATOR_TYPE=v4-32
Crea la VM de TPU
$ gcloud compute tpus tpu-vm create ${TPU_NAME} \ --zone=${ZONE} \ --project=${PROJECT_ID} \ --accelerator-type=${ACCELERATOR_TYPE} \ --version ${RUNTIME_VERSION}
Configura y ejecuta la secuencia de comandos de entrenamiento
Agrega el certificado SSH a tu proyecto:
ssh-add ~/.ssh/google_compute_engine
Instala PyTorch/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=" pip install torch~=2.4.0 torch_xla[tpu]~=2.4.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html"
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 -b r2.4 https://github.com/pytorch/xla.git"
Ejecuta la secuencia de comandos de entrenamiento en todos los trabajadores
gcloud compute tpus tpu-vm ssh ${TPU_NAME} \ --zone=${ZONE} \ --project=${PROJECT_ID} \ --worker=all \ --command="PJRT_DEVICE=TPU python3 ~/xla/test/test_train_mp_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.
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.Borra tus recursos de Cloud TPU y Compute Engine.
$ gcloud compute tpus tpu-vm delete \ --zone=${ZONE}
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}