Guía del usuario de VM de Cloud TPU

En esta guía, se describe cómo configurar un proyecto de Google Cloud Platform para usar las VM de Cloud TPU. Se describen los comandos para usar las VM y soluciones de Cloud TPU a problemas comunes que pueden surgir al comenzar a usar las VM de Cloud TPU.

Las VM de Cloud TPU se ejecutan en la máquina host de TPU (la máquina conectada al dispositivo de Cloud TPU) y ofrecen un rendimiento y una usabilidad mucho mejor cuando funcionan con TPU.

Si eres nuevo en Cloud TPU, consulta la guía para principiantes de TPU.

La vista previa de VM de Cloud TPU introduce una nueva arquitectura de Cloud TPU. La arquitectura actual se llama nodos de TPU. Para obtener más información sobre el tipo de nodo de Cloud TPU y de las VM de Cloud TPU, consulta Arquitectura del sistema.

Configura un proyecto de GCP

Después de instalar el SDK de Google Cloud, instala los componentes de alpha para gcloud con el siguiente comando:

gcloud components install alpha

Para obtener más información sobre los componentes de gcloud, consulta Cómo administrar componentes del SDK.

Preparar un proyecto de GCP

Accede a tu Cuenta de Google. Si aún no tienes una cuenta, regístrate para obtener una nueva. En Google Cloud Console, selecciona o crea un proyecto de Cloud desde la página de selección de proyecto. Asegúrate de tener habilitada la facturación para tu proyecto. Configura el ID del proyecto con gcloud en Cloud Shell. El ID del proyecto es el nombre de tu proyecto que se muestra en Cloud Console.

$ gcloud config set project project-id

Habilita la API de Cloud TPU

Habilita la API de Cloud TPU con el siguiente comando de gcloud en Cloud Shell. (También puedes habilitarlo desde Google Cloud Console).

$ gcloud services enable tpu.googleapis.com

Configura el comando de gcloud

Ejecuta los comandos siguientes a fin de configurar gcloud a fin de usar tu proyecto de GCP y, luego, instalar los componentes necesarios para la vista previa de la VM de TPU.

$ gcloud config set account your-email-account
$ gcloud config set project your-project

Administra TPU

Puedes administrar la VM de Cloud TPU con gcloud o curl. Para obtener más información, consulta Administra Cloud TPU.

Crea una VM de Cloud TPU con gcloud

$ gcloud alpha compute tpus tpu-vm create tpu-name \
  --zone=zone \
  --accelerator-type=v3-8 \
  --version=v2-alpha

Campos obligatorios

zone
La zona en la que planeas crear tu Cloud TPU.
accelerator-type
El tipo de Cloud TPU que se creará.
version
La versión del entorno de ejecución de Cloud TPU.

Crea una VM de Cloud TPU con curl

$ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" \
  -H "Content-Type: application/json" -d "{accelerator_type: 'v2-8', \
  runtime_version:'[runtime_version]', \
  network_config: {enable_external_ips: true}}" \
  https://tpu.googleapis.com/v2alpha1/projects/[project]/locations/[zone]/nodes?node_id=[node_name]

Campos obligatorios

runtime_version
La versión del entorno de ejecución que deseas usar.
project
El nombre de tu proyecto de GCP inscrito.
zone
La zona en la que creas el Cloud TPU.
node_name
El nombre de la VM de TPU que estás creando.

Conéctate a una VM de Cloud TPU

  1. (opcional). Configura un firewall para SSH

La red predeterminada viene preconfigurada para permitir el acceso SSH a todas las VM. Si no usas la red predeterminada o la configuración de red predeterminada se editó, es posible que debas agregar una regla de firewall para habilitar explícitamente el acceso SSH:

$ gcloud compute firewall-rules create --network=network allow-ssh --allow=tcp:22
  1. Establece una conexión SSH a las VM de TPU
$ gcloud alpha compute tpus tpu-vm ssh tpu-name --zone zone --project project-id

Campos obligatorios

tpu_name
El nombre de la VM de TPU a la que te conectas.
zone
La zona en la que creas el Cloud TPU. Actualmente, se admiten europe-west4-a y us-central1-a.
project-id
El ID de tu proyecto de GCP.

Campos opcionales

user
Puedes elegir el nombre de usuario usado para autenticarte cuando te conectas a la VM de Cloud TPU a través de SSH, con el prefijo $USER@ en el nombre de la TPU, por ejemplo: my-email-account@tpu-node-1.
worker
Para los pods de Cloud TPU, puedes elegir en qué VM de trabajador establecer una conexión SSH. El valor predeterminado es el trabajador 0, la primera VM asociada con el pod de TPU.
ssh-key-file
La ruta de acceso al archivo de llaves SSH. El valor predeterminado es ~/.ssh/google_compute_engine.
internal-ip
Conéctate a las VM de TPU mediante una dirección IP interna. Para que esta conexión funcione, debes configurar tus redes y tu firewall a fin de permitir conexiones SSH a la dirección IP interna de la VM de TPU a la que deseas conectarte.
command
Es un comando para ejecutar en la VM de TPU. El comando se ejecuta en la VM de TPU de destino y, luego, se cierra.
tunnel-through-iap
Configura la conexión SSH mediante Cloud Identity-Aware Proxy para redireccionar TCP. Para obtener más información, consulta [Descripción general de la redirección de TCP](https://cloud.google.com/iap/docs/tcp-forwarding-overview).

Para establecer una conexión SSH a otras VM de TPU asociadas con el pod de TPU, agrega --worker ${WORKER_NUMBER} en el comando, en el que WORKER_NUMBER es un índice basado en 0.

Enumera tus recursos de Cloud TPU

Puedes enumerar todas tus Cloud TPU en una zona específica.

$ gcloud alpha compute tpus tpu-vm list --zone=zone

Campos obligatorios

zone
La zona en la que planeas crear tu Cloud TPU.

Este comando enumera los recursos de Cloud TPU de la zona especificada. Si no hay recursos configurados, el resultado solo mostrará guiones para la VM y la TPU.

Recupera información sobre tu Cloud TPU

Puedes recuperar información sobre una Cloud TPU específica con el comando siguiente.

$ gcloud alpha compute tpus tpu-vm describe tpu-name \
  --zone=zone

Campos obligatorios

tpu-name
El nombre de la Cloud TPU para crear.
zone
La zona en la que se creó tu Cloud TPU.

Detén tus recursos de Cloud TPU

Puedes detener una Cloud TPU con el comando siguiente. No puedes detener un pod de TPU.

$ gcloud alpha compute tpus tpu-vm stop tpu-name \
  --zone=zone

Campos obligatorios

tpu-name
El nombre de la Cloud TPU que deseas detener.
zone
La zona en la que creaste tu Cloud TPU.

Inicia tus recursos de Cloud TPU

Si tu Cloud TPU se detuvo, puedes reiniciarla con el siguiente comando.

$ gcloud alpha compute tpus tpu-vm start tpu-name --zone  zone

Descripciones de las marcas de comandos

tpu-name
El nombre de la Cloud TPU que se iniciará.
zone
La zona en la que se creó la Cloud TPU.

Borra tus VM y recursos de Cloud TPU

Puedes borrar tu Cloud TPU cuando hayas terminado de usarlas.

$ gcloud alpha compute tpus tpu-vm delete tpu-name \
  --zone=zone

Campos obligatorios

zone
La zona en la que se creó tu Cloud TPU.

Captura métricas de rendimiento

Puedes capturar un perfil de rendimiento mediante una secuencia de comandos de la línea de comandos o con Tensorboard. Si deseas obtener instrucciones para instalar Tensorboard, consulta Configuración de Tensorboard.

Para los modelos de Tensorflow, puedes capturar datos de perfil de forma automática mediante el método de devolución de llamada de generación de perfiles de Tensorflow estándar.

Para capturar manualmente los datos de perfil de modelos de TensorFlow, usa el siguiente comando en tu VM de TPU:

$ python3 -c "import tensorflow as tf; tf.profiler.experimental.client.trace('grpc://localhost:port', 'gs://model-dir', 1000)"

Para capturar datos de perfil de modelos de PyTorch con la línea de comandos, usa el siguiente comando en la VM de TPU:

$ python3 -c "import torch_xla.debug.profiler as xp; xp.trace('localhost:port', '/tmp/tb', 1000)"

Para obtener información sobre cómo capturar datos de perfil para modelos de JAX, consulta Cómo generar perfiles de programas JAX.

Cómo ver los datos del perfil

  1. Abre Cloud Shell
  2. Asegúrate de haber instalado Tensorboard.
  3. Ejecuta TensorBoard
  4. Desde Cloud Shell, haz clic en el botón Vista previa en la Web y selecciona Cambiar puerto y escribe 6006.
  5. Se muestra una página de resumen.
  6. Cómo navegar hasta el lector de seguimiento en Herramientas

Solicita una cuota de TPU mayor

Cloud TPU asigna cuota de TPU predeterminada para tu proyecto. Si necesitas más, consulta Solicita una cuota adicional.

Genera y visualiza la información de perfil

Puedes generar información de perfil y usar TensorBoard para visualizar las métricas de entrenamiento.

Solución de problemas de configuración de gcloud

Problema

gcloud components update muestra el siguiente mensaje de error:

ERROR: (gcloud.components.update)
You cannot perform this action because the Cloud SDK component manager
is disabled for this installation.
Solución

Para usar gcloud con una VM de TPU, deberás usar una instalación gcloud que no se administre a través de un administrador de paquetes. Sigue estos pasos para instalar gcloud desde el código fuente:

sudo apt-get remove google-cloud-sdk
curl -O https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
tar -xzf google-cloud-sdk-311.0.0-linux-x86_64.tar.gz
./google-cloud-sdk/install.sh
source ~/.bashrc
Problema

La ejecución de cualquier comando que comience con gcloud alpha compute tpus tpu-vm muestra la siguiente información:

ERROR: (gcloud.alpha.compute.tpus) Invalid choice: 'tpu-vm'.
Solución

Esto sucede cuando el repositorio de componentes no se actualizó de forma correcta. Para verificar esto, ejecuta gcloud --version. La primera línea del resultado debería ser "Google Cloud SDK HEAD"; Si el resultado es diferente, la actualización no se realizó. Si esto sucede, intenta actualizar los componentes de gcloud con el siguiente comando.

gcloud components update

Si aún recibes el mismo error, intenta reinstalar gcloud con el siguiente comando:

gcloud components reinstall
Problema

El comando gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} --zone ${ZONE} muestra el siguiente mensaje de error:

Waiting for SSH key to propagate.
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ssh: connect to host 34.91.136.59 port 22: Connection timed out
ERROR: (gcloud.alpha.compute.tpus.tpu-vm.ssh) Could not SSH into the instance.  It is possible that your SSH key has not propagated to the instance yet. Try running this command again.  If you still cannot connect, verify that the firewall and instance are set to accept ssh traffic.
Solución

Es posible que haya un problema con la propagación de llaves SSH. Intenta mover las claves generadas de forma automática a una ubicación de copia de seguridad para forzar gcloud a recrearlas:

mv ~/.ssh/google_compute_engine ~/.ssh/old-google_compute_engine
mv ~/.ssh/google_compute_engine.pub ~/.ssh/old-google_compute_engine.pub

Preguntas frecuentes

¿Puedo usar las API V1Alpha1 y V1 para administrar las VM de Cloud TPU?
Se permite Get/List, pero las mutaciones solo están disponibles en la versión de la API V2Alpha1.
¿Cómo sé si las TPU usan las VM de Cloud TPU?
Realiza una llamada GetNode en la TPU, configura el campo APIVersion como V2_ALPHA1.