Guía de 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 de Cloud TPU y las soluciones a los problemas comunes que puedes encontrar cuando comienzas a usar las VM de Cloud TPU.

Las VM de Cloud TPU se ejecutan en la máquina anfitrión de TPU (la máquina conectada al dispositivo de Cloud TPU) y ofrecen un rendimiento y una usabilidad significativamente mejores cuando se trabaja con TPU.

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

La vista previa de la VM de Cloud TPU presenta una nueva arquitectura de Cloud TPU. La arquitectura actual se llama nodos TPU. Para obtener más información sobre las arquitecturas del nodo de Cloud TPU y de la 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 mediante 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.

Prepara un proyecto de GCP

Accede a tu Cuenta de Google. Si aún no tienes una, regístrate para obtener una cuenta 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 mediante 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 siguientes comandos a fin de configurar gcloud con el fin de usar tu proyecto de GCP e 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 las VM de Cloud TPU mediante gcloud o curl. Para obtener más información, consulta Administra las 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
Es la zona en la que deseas crear la 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 crearás la Cloud TPU.
node_name
El nombre de la VM de TPU que crearás.

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 se editó la configuración de red predeterminada, es posible que debas habilitar de forma explícita el acceso SSH mediante la adición de una regla de firewall:

$ 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 estás conectando.
zone
La zona en la que crearás la Cloud TPU. Por el momento, 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 que se usa para la autenticación cuando te conectas a la VM de Cloud TPU mediante SSH, con el prefijo $USER@ para el nombre de TPU, por ejemplo: my-email-account@tpu-node-1.
worker
Para los pods de Cloud TPU, puedes elegir la VM de trabajador a la que deseas acceder con 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 claves SSH. El valor predeterminado es ~/.ssh/google_compute_engine.
internal-ip
Conéctate a las VM de TPU con una dirección IP interna. Para que esta conexión funcione, debes configurar tus redes y 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 que se ejecuta en la VM de TPU. El comando se ejecuta en la VM de TPU de destino y, luego, sale.
tunnel-through-iap
Crea un túnel a la conexión SSH a través de Cloud Identity-Aware Proxy para la redirección de TCP. Para obtener más información, consulta [Descripción general del reenvío 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 tiene el índice basado en 0.

Enumera los 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
Es la zona en la que deseas crear la Cloud TPU.

Con este comando, se enumeran los recursos de Cloud TPU en la zona especificada. Si no hay recursos configurados, el resultado solo mostrará guiones para la VM y la TPU.

Recupera información sobre Cloud TPU

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

$ 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 sola Cloud TPU con el siguiente comando. No puedes detener un pod de TPU.

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

Campos obligatorios

tpu-name
Es el nombre de la Cloud TPU que se detendrá.
zone
La zona en la que creaste tu Cloud TPU.

Inicia tus recursos de Cloud TPU

Si se detuvo tu Cloud TPU, puedes reiniciarla mediante 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 la VM y los recursos de Cloud TPU

Puedes borrar tu Cloud TPU cuando termines de usarlos.

$ 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 con una secuencia de comandos de línea de comandos o TensorBoard. Para obtener instrucciones sobre cómo instalar TensorBoard, consulta Configuración de TensorBoard.

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

A fin de capturar de forma manual los datos del perfil de los 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)"

Si deseas capturar datos de perfil para modelos de PyTorch mediante la línea de comandos, usa el siguiente comando en tu 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 Crea perfiles de programas de JAX.

Visualiza datos de 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, selecciona Cambiar puerto y escribe 6006.
  5. Haz clic en el perfil para ver una página de resumen.
  6. Cómo navegar para hacer un seguimiento de usuarios con las herramientas

Solicita una cuota de TPU mayor

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

Genera y visualiza la información del 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 la VM de TPU, deberás usar una instalación de gcloud que no esté administrada 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

Si ejecutas cualquier comando que comience con gcloud alpha compute tpus tpu-vm, se mostrará 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ó correctamente. Para verificar esto, ejecuta gcloud --version. La primera línea del resultado debe 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 todavía ves 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 la Llave SSH. Intenta mover las claves generadas de manera automática a una ubicación de copia de seguridad para forzar a gcloud a que vuelvan a crearlas:

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 VM de Cloud TPU?
Se permite obtener o enumerar mutaciones, pero las mutaciones solo están disponibles en la versión V2Alpha1 de la API.
¿Cómo sé si las TPU usan VM de Cloud TPU?
Realiza una llamada GetNode a la TPU, establece el campo APIVersion en V2_ALPHA1.