Crea una instancia de TPU

En este documento, se muestra cómo crear una instancia de máquina virtual (VM) de TPU v6e (Trillium), conectarte a ella y ejecutar un cálculo en ella. La instancia de TPU que creas en esta guía de inicio rápido usa la opción de consumo predeterminada, que es a pedido.

Antes de comenzar

  1. Accede a tu Google Cloud cuenta de. Si eres nuevo en Google Cloud, crea una cuenta para evaluar el rendimiento de nuestros productos en situaciones reales. Los clientes nuevos también obtienen $300 en créditos gratuitos para ejecutar, probar y, además, implementar cargas de trabajo.

  2. Instala Google Cloud CLI.

  3. Si usas un proveedor de identidad externo (IdP), primero debes acceder a la gcloud CLI con tu identidad federada.

  4. Para inicializar gcloud CLI, ejecuta el siguiente comando: 

    gcloud init

  5. Crea o selecciona un Google Cloud proyecto.

    Roles obligatorios para seleccionar o crear un proyecto

    • Seleccionar un proyecto: Para seleccionar un proyecto, no se requiere un rol de IAM específico. Puedes seleccionar cualquier proyecto en el que se te haya otorgado un rol.
    • Crear un proyecto: Para crear un proyecto, necesitas el rol de creador de proyectos (roles/resourcemanager.projectCreator), que contiene el resourcemanager.projects.create permiso. Obtén información para otorgar roles.

    • Crea un proyecto de: Google Cloud 

      gcloud projects create PROJECT_ID

      Reemplaza PROJECT_ID por un nombre para el Google Cloud proyecto de que estás creando.

    • Selecciona el Google Cloud proyecto de que creaste:

      gcloud config set project PROJECT_ID

      Reemplaza PROJECT_ID por el nombre de tu Google Cloud proyecto de.

  6. Si usas un proyecto existente en esta guía, verifica que tengas los permisos necesarios para completarla. Si creaste un proyecto nuevo, ya tienes los permisos necesarios.

  7. Verifica que la facturación esté habilitada para tu Google Cloud proyecto.

  8. Habilita la API de Compute Engine con este comando:

    Roles necesarios para habilitar las APIs

    Para habilitar las APIs, necesitas el permiso serviceusage.services.enable. Si creaste el proyecto, es probable que ya tengas este permiso a través del rol de propietario (roles/owner). De lo contrario, puedes obtener este permiso a través del rol de administrador de Service Usage (roles/serviceusage.serviceUsageAdmin). Obtén información para otorgar roles. 

    gcloud services enable compute.googleapis.com

  9. Instala Google Cloud CLI.

  10. Si usas un proveedor de identidad externo (IdP), primero debes acceder a la gcloud CLI con tu identidad federada.

  11. Para inicializar gcloud CLI, ejecuta el siguiente comando: 

    gcloud init

  12. Crea o selecciona un Google Cloud proyecto.

    Roles obligatorios para seleccionar o crear un proyecto

    • Seleccionar un proyecto: Para seleccionar un proyecto, no se requiere un rol de IAM específico. Puedes seleccionar cualquier proyecto en el que se te haya otorgado un rol.
    • Crear un proyecto: Para crear un proyecto, necesitas el rol de creador de proyectos (roles/resourcemanager.projectCreator), que contiene el resourcemanager.projects.create permiso. Obtén información para otorgar roles.

    • Crea un proyecto de: Google Cloud 

      gcloud projects create PROJECT_ID

      Reemplaza PROJECT_ID por un nombre para el Google Cloud proyecto de que estás creando.

    • Selecciona el Google Cloud proyecto de que creaste:

      gcloud config set project PROJECT_ID

      Reemplaza PROJECT_ID por el nombre de tu Google Cloud proyecto de.

  13. Si usas un proyecto existente en esta guía, verifica que tengas los permisos necesarios para completarla. Si creaste un proyecto nuevo, ya tienes los permisos necesarios.

  14. Verifica que la facturación esté habilitada para tu Google Cloud proyecto.

  15. Habilita la API de Compute Engine con este comando:

    Roles necesarios para habilitar las APIs

    Para habilitar las APIs, necesitas el permiso serviceusage.services.enable. Si creaste el proyecto, es probable que ya tengas este permiso a través del rol de propietario (roles/owner). De lo contrario, puedes obtener este permiso a través del rol de administrador de Service Usage (roles/serviceusage.serviceUsageAdmin). Obtén información para otorgar roles. 

    gcloud services enable compute.googleapis.com
    16.

Roles obligatorios

Si deseas obtener los permisos que necesitas para crear una instancia de TPU, conectarte a ella con SSH y ejecutar comandos, pídele a tu administrador que te otorgue los siguientes roles de IAM en tu proyecto:

Para obtener más información sobre cómo otorgar roles, consulta Administra el acceso a proyectos, carpetas y organizaciones.

También puedes obtener los permisos necesarios mediante roles personalizados o cualquier otro rol predefinido.

Crea una instancia de TPU v6e (Trillium)

Ejecuta los siguientes comandos en tu terminal local o con Cloud Shell.

  1. Para crear una instancia de TPU v6e, ejecuta el siguiente gcloud compute instances create comando:

    gcloud compute instances create quickstart-tpu \
    --machine-type=ct6e-standard-4t \
    --zone=europe-west4-a \
    --image-project=ubuntu-os-accelerator-images \
    --image-family=ubuntu-accel-2204-amd64-tpu-v5e-v5p-v6e \
    --maintenance-policy=TERMINATE

  2. De manera opcional, verifica que la instancia se esté ejecutando con el gcloud compute instances describe comando:

    gcloud compute instances describe quickstart-tpu --zone=europe-west4-a

    En el resultado, verifica el campo status para comprobar si la instancia se está ejecutando o aún está pendiente.

Conéctate a la instancia de TPU

Puedes conectarte a una instancia de TPU con la Google Cloud consola o la gcloud CLI.

Para conectarte a la VM con SSH, ejecuta el siguiente gcloud compute ssh comando:

gcloud compute ssh quickstart-tpu --zone=europe-west4-a

Instala JAX

Para instalar el framework de JAX, ejecuta el siguiente comando en tu instancia de TPU:

pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

Ejecuta una verificación del sistema

Para verificar que JAX pueda acceder a la TPU y ejecutar operaciones básicas, ejecuta los siguientes comandos en tu instancia de TPU:

  1. Inicia el intérprete de Python 3:

    python3

  2. Importa el módulo de JAX:

    >>> import jax

  3. Muestra la cantidad de núcleos de TPU disponibles:

    >>> jax.device_count()

    El resultado muestra la cantidad de núcleos de TPU. La cantidad de núcleos depende de la versión de TPU que uses. Para obtener más información, consulta Versiones de TPU.

Cómo hacer un cálculo

  1. Ejecuta el siguiente comando para sumar dos números:

    >>> jax.numpy.add(1, 1)

    El resultado debería ser similar al siguiente ejemplo:

    >>> Array(2, dtype=int32, weak_type=True)

  2. Sal del intérprete de Python:

    >>> 
exit()

Limpia

Para evitar que se apliquen cargos a tu Google Cloud cuenta de por los recursos que usaste en esta página, borra el Google Cloud proyecto de que tiene los recursos.

Como alternativa, si deseas conservar el proyecto, puedes borrar solo la VM con el comando gcloud compute instances delete :

gcloud compute instances delete quickstart-tpu --zone=europe-west4-a

¿Qué sigue?