Se usó la API de Cloud Translation para traducir esta página.

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 Cloud TPU. Para obtener un instructivo más detallado que te muestre cómo entrenar un modelo en Cloud TPU, consulta uno de los instructivos de PyTorch en Cloud TPU.

Antes de comenzar

Antes de seguir esta guía de inicio rápido, debes crear una cuenta de Google Cloud Platform, instalar Google Cloud CLI y configurar 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, la red y la zona del usuario predeterminados, ejecuta lo siguiente:

$ gcloud compute tpus tpu-vm create tpu-name \
   --zone=us-central1-a \
   --accelerator-type=v3-8 \
   --version=tpu-ubuntu2204-base

Descripciones de las marcas de comandos

zone: Es la zona en la que deseas crear la 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: La versión del software de Cloud TPU

Cuando crees tu TPU, puedes pasar las marcas --network y --subnetwork adicionales si deseas especificar la red y la subred predeterminadas. Si no quieres usar la red predeterminada, debes pasar la marca --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 página de referencia de la API de gcloud para obtener detalles sobre estas marcas.

Conéctate a tu VM de Cloud TPU

   $ gcloud compute tpus tpu-vm ssh tpu-name --zone=us-central1-a

Instala PyTorch/XLA en tu VM de TPU

   (vm)$ pip install torch~=2.5.0 torch_xla[tpu]~=2.5.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

Verifica que PyTorch pueda acceder a las TPU

Crea un archivo llamado tpu-count.py en el directorio actual y copia y pega la siguiente secuencia de comandos en él.

import torch
import torch_xla
import torch_xla.core.xla_model as xm
print(f'PyTorch can access {len(torch_xla.devices())} TPU cores')

Ejecuta la secuencia de comandos:
```
(vm)$ python3 tpu-count.py
```
El resultado de la secuencia de comandos muestra el resultado del procesamiento:
```
PyTorch can access 8 TPU cores
```

Realiza un cálculo básico

Crea un archivo llamado tpu-test.py en el directorio actual y copia y pega la siguiente secuencia de comandos en él.

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-central1-a

El resultado de este comando debería confirmar que se borró tu TPU.

¿Qué sigue?

Obtén más información sobre las VMs de Cloud TPU: