Ejecuta un cálculo en una VM de Cloud TPU con PyTorch
En esta guía de inicio rápido, se muestra cómo crear una Cloud TPU, instalar PyTorch y ejecutar un cálculo simple en una Cloud TPU. Para ver un instructivo más detallado que muestre cómo entrenar un modelo en Cloud TPU ver una de las Instructivos de PyTorch de Cloud TPU.
Antes de comenzar
Antes de seguir esta guía de inicio rápido, debes crear una Google Cloud Platform
, instala Google Cloud CLI. y configura el comando gcloud
.
Para obtener más información, consulta
Configura una cuenta y un proyecto de Cloud TPU.
Crea una Cloud TPU con gcloud
Para crear una VM de TPU en el proyecto de usuario predeterminado, la red y la zona ejecuta lo siguiente:
$ gcloud compute tpus tpu-vm create tpu-name \
--zone=us-central2-b \
--accelerator-type=v4-8 \
--version=tpu-ubuntu2204-base
Descripciones de las marcas de comandos
zone
- La zona en la que planeas crear tu Cloud TPU.
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.
version
- El software de Cloud TPU version.
Mientras creas la TPU, puedes pasar los --network
adicionales y
--subnetwork
marca si
debes especificar la red y la subred predeterminadas.
Si no deseas usar la red predeterminada, debes pasar el
--network
. La marca --subnetwork
es opcional y se puede usar para
especificar una subred predeterminada para cualquier red que uses (predeterminada o
especificada por el usuario). Consulta la gcloud
Página de referencia de la API para obtener detalles sobre estas marcas.
Conéctate a tu VM de Cloud TPU
$ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central2-b
Instala PyTorch/XLA en tu VM de TPU
(vm)$ pip install torch~=2.4.0 torch_xla[tpu]~=2.4.0 torchvision -f https://storage.googleapis.com/libtpu-releases/index.html
Establece la configuración del entorno de ejecución de TPU
Asegúrate de que el entorno de ejecución de PyTorch/XLA use la TPU.
(vm) $ export PJRT_DEVICE=TPU
Realiza un cálculo simple:
Crea un archivo llamado
tpu-test.py
en el directorio actual y cópialo pega la siguiente secuencia de comandos.import torch import torch_xla.core.xla_model as xm dev = xm.xla_device() t1 = torch.randn(3,3,device=dev) t2 = torch.randn(3,3,device=dev) print(t1 + t2)
Ejecuta la secuencia de comandos:
(vm)$ python3 tpu-test.py
El resultado de la secuencia de comandos muestra el resultado del procesamiento:
tensor([[-0.2121, 1.5589, -0.6951], [-0.7886, -0.2022, 0.9242], [ 0.8555, -1.8698, 1.4333]], device='xla:1')
Limpia
Sigue estos pasos para evitar que se apliquen cargos a tu cuenta de Google Cloud por los recursos que usaste en esta página.
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 tu Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name \ --zone=us-central2-b
El resultado de este comando debería confirmar que se borró la TPU.
¿Qué sigue?
Obtén más información acerca de las VMs de Cloud TPU: