Descripción general de Cloud TPU Multislice

Cloud TPU Multislice es una tecnología de escalamiento del rendimiento full stack que permite que un trabajo de entrenamiento use varias porciones de TPU dentro de una sola porción o porciones de varios Pods con paralelismo de datos estándar. Con los chips TPU v4, esto significa que los trabajos de entrenamiento pueden usar más de 4,096 chips en una sola ejecución. Para los trabajos de entrenamiento que requieren menos de 4,096 chips, una sola porción puede ofrecer el mejor rendimiento. Sin embargo, varias porciones más pequeñas están disponibles con mayor rapidez, lo que permite un tiempo de inicio más rápido cuando se usa Multislice con porciones más pequeñas.

Cuando se implementan en parámetros de configuración de Multislice, los chips TPU en cada porción se comunican a través de interconexiones entre chips (ICI). Los chips TPU en diferentes segmentos se comunican transfiriendo datos a las CPU (hosts), que, a su vez, transmiten los datos a través de la red del centro de datos (DCN). Para obtener más información sobre el escalamiento con Multislice, consulta Cómo escalar el entrenamiento de IA hasta decenas de miles de chips Cloud TPU con Multislice.

Los desarrolladores no tienen que escribir un código para implementar la comunicación entre porciones de la DCN. El compilador XLA genera ese código por ti y superpone la comunicación con el cálculo para lograr el máximo rendimiento.

Conceptos

Tipo de acelerador

Es la forma de cada porción de TPU que compone una Multislice. Cada porción de una solicitud de varios segmentos es del mismo tipo de acelerador. Un tipo de acelerador consta de un tipo de TPU (v4 o posterior) seguido de la cantidad de TensorCores. Por ejemplo, v5litepod-128 especifica una TPU v5e con 128 TensorCores.

Reparación automática

Cuando una porción experimenta un evento de mantenimiento, una interrupción o una falla de hardware, Cloud TPU creará una porción nueva. Si no hay suficientes recursos para crear una nueva porción, la creación no se completará hasta que el hardware esté disponible. Después de crear la nueva porción, se reiniciarán todos los demás segmentos del entorno de Multislice para que pueda continuar el entrenamiento. Con una secuencia de comandos de inicio configurada correctamente, el entrenamiento puede reiniciarse de forma automática sin intervención del usuario, cargarse y reanudarse desde el último punto de control.

Redes de centros de datos (DCN)

Una red de menor capacidad de procesamiento y mayor latencia (en comparación con la ICI) que conecta porciones de TPU en un parámetro de configuración Multislice.

Programación en grupo

Cuando todas las porciones de TPU se aprovisionan juntas, al mismo tiempo, lo que garantiza que todas o ninguna de las porciones se aprovisionen de forma correcta.

Interconexión entre chips (ICI)

Vínculos internos de alta velocidad y baja latencia que conectan las TPU dentro de un pod de TPU.

Multislice

Dos o más porciones de chips TPU que pueden comunicarse a través de la DCN.

Nodo

En el contexto de Multislice, el término nodo hace referencia a una sola porción de TPU. A cada porción de TPU en Multislice se le asigna un ID de nodo.

Secuencia de comandos de inicio

Una secuencia de comandos de inicio de Compute Engine estándar que se ejecuta cada vez que se inicia o reinicia una VM. En el caso de Multislice, se especifica en la solicitud de creación de QR. Para obtener más información sobre las secuencias de comandos de inicio de Cloud TPU, consulta Administra recursos TPU.

Tensor

Estructura de datos que se usa para representar datos multidimensionales en un modelo de aprendizaje automático.

Tipos de capacidad de Cloud TPU

Las TPU se pueden crear a partir de diferentes tipos de capacidad (consulta las opciones de uso en Cómo funcionan los precios de las TPU):

• Reserva: Para utilizar una reserva, debes tener un acuerdo de reserva con Google. Usa la marca --reserved cuando creas tus recursos.

• Spot: Se orienta a la cuota interrumpible con VMs Spot. Es posible que tus recursos se interrumpan para dar lugar a solicitudes de un trabajo de mayor prioridad. Usa la marca --spot cuando creas tus recursos.

• Según demanda: Segmenta la cuota según demanda, que no necesita una reserva y no se interrumpirá. La solicitud de TPU se pondrá en la fila en una cola de cuota según demanda que ofrece Cloud TPU, pero no se garantiza la disponibilidad de los recursos. Se selecciona de forma predeterminada y no se necesitan marcas.

Comienza

1. Configura tu entorno de Cloud TPU.

2. In the Google Cloud console, activate Cloud Shell.

   Activate Cloud Shell

   At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

Para usar Multislice, tus recursos TPU deben administrarse como recursos en fila.

Ejemplo introductorio

En este instructivo, se usa código del repositorio de GitHub de MaxText. MaxText es un LLM básico de alto rendimiento, escalable de forma arbitraria, de código abierto y bien probado, escrito en Python y Jax. MaxText se diseñó para entrenar de manera eficiente en Cloud TPU.

El código en shardings.py está diseñado para ayudarte a comenzar a experimentar con diferentes opciones de paralelismo. Por ejemplo, el paralelismo de datos, el paralelismo de datos completamente fragmentados (FSDP) y el paralelismo de tensores. El código escala desde una sola porción hasta entornos de Multislice.

Paralelismo de ICI

ICI hace referencia a la interconexión de alta velocidad que conecta las TPU en una sola porción. La fragmentación de ICI corresponde a fragmentar dentro de una porción. shardings.py proporciona tres parámetros de paralelismo de ICI:

• ici_data_parallelism
• ici_fsdp_parallelism
• ici_tensor_parallelism

Los valores que especifiques para estos parámetros determinan la cantidad de fragmentos para cada método de paralelismo.

Estas entradas deben restringirse de modo que ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism sea igual a la cantidad de chips en la porción.

En la siguiente tabla, se muestran ejemplos de entradas del usuario para el paralelismo de ICI en los cuatro chips disponibles en v4-8:

	ici_data_parallelism	ici_fsdp_parallelism	ici_tensor_parallelism
FSDP de 4 vías	1	4	1
Paralelismo de tensor de 4 direcciones	1	1	4
FSDP bidireccional + paralelismo de tensores bidireccional	1	2	2

Ten en cuenta que, en la mayoría de los casos, ici_data_parallelism debe ser 1, ya que la red de ICI es tan rápida como para preferir casi siempre el FSDP al paralelismo de datos.

En este ejemplo, se supone que sabes cómo ejecutar un código en una sola porción de TPU, como se explica en Ejecuta un cálculo en una VM de Cloud TPU con JAX. En este ejemplo, se muestra cómo ejecutar shardings.py en una sola porción.

1. Configura el entorno:

   $ gcloud auth login
   $ export QR_ID=your-queued-resource-id
   $ export TPU_NAME=your-tpu-name
   $ export PROJECT=your-project-name
   $ export ZONE=us-central1-a
   $ export NETWORK_NAME=your-network-name
   $ export SUBNETWORK_NAME=your-subnetwork-name
   $ export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   $ export ACCELERATOR_TYPE=v5litepod-16
   $ export EXAMPLE_TAG_1=your-tag-1
   $ export EXAMPLE_TAG_2=your-tag-2
   $ export SLICE_COUNT=4
   $ export STARTUP_SCRIPT='#!/bin/bash\n'

   Descripciones de las variables

   Entrada	Descripción
   QR_ID	Es el ID asignado por el usuario del recurso en línea.
   TPU_NAME	Es el nombre que el usuario asignó a la TPU.
   PROJECT	Es el nombre del proyecto deGoogle Cloud .
   ZONE	Especifica la zona en la que se crearán los recursos.
   NETWORK_NAME	Nombre de las redes de VPC.
   SUBNETWORK_NAME	Es el nombre de la subred en las redes de VPC.
   RUNTIME_VERSION	Es la versión de software de Cloud TPU.
   ACCELERATOR_TYPE	v4-16
   EXAMPLE_TAG_1, EXAMPLE_TAG_2…	Las etiquetas se usan para identificar objetivos o fuentes válidos para firewalls de red.
   SLICE_COUNT	Es la cantidad de porciones. Se limita a un máximo de 256 porciones.
   STARTUP_SCRIPT	Si especificas una secuencia de comandos de inicio, esta se ejecutará cuando se aprovisione o reinicie la porción de TPU.

2. Crea claves SSH para gcloud. Te recomendamos que dejes la contraseña en blanco (presiona Intro dos veces después de ejecutar el siguiente comando). Si se te indica que el archivo google_compute_engine ya existe, reemplaza la versión existente.

   $ ssh-keygen -f ~/.ssh/google_compute_engine

3. Aprovisiona tus TPU:

   gcloud

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --node-id=${TPU_NAME} \
       --zone=${ZONE} \
       [--reserved |--spot]

   La Google Cloud CLI no admite todas las opciones de creación de códigos QR, como las etiquetas. Para obtener más información, consulta Crea códigos QRs.

   Consola

   1. En la consola de Google Cloud , ve a la página TPU:

      Ir a TPU

   2. Haz clic en Crear una TPU.

   3. En el campo Nombre, ingresa el nombre de tu TPU.

   4. En el cuadro Zona, selecciona la zona en la que deseas crear la TPU.

   5. En el cuadro Tipo de TPU, selecciona un tipo de acelerador. 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.

   6. En el cuadro Versión de software de TPU, selecciona una versión de software. Cuando creas una VM de Cloud TPU, la versión de software de TPU especifica la versión del entorno de ejecución de TPU que se instalará. Para obtener más información, consulta Versiones de software de TPU.

   7. Haz clic en el botón para activar o desactivar la opción Habilitar las colas.

   8. En el campo Nombre del recurso en cola, ingresa un nombre para tu solicitud de recurso en cola.

   9. Haz clic en Crear para crear tu solicitud de recursos en cola.

4. Espera hasta que el recurso en fila esté en el estado ACTIVE, lo que significa que los nodos trabajadores están en el estado READY. Una vez que comience el aprovisionamiento de recursos en fila, puede tardar de uno a cinco minutos en completarse, según el tamaño del recurso en fila. Puedes verificar el estado de una solicitud de recursos en fila con la gcloud CLI o la consola de Google Cloud :

   gcloud

   $ gcloud compute tpus queued-resources \
       list --filter=${QR_ID} --zone=${ZONE}

   Consola

   1. En la consola de Google Cloud , ve a la página TPU:

      Ir a TPU

   2. Haz clic en la pestaña Recursos en cola.

   3. Haz clic en el nombre de la solicitud de recurso en cola.

5. Conéctate a la VM de TPU con SSH:

   $ gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE}

6. Clona MaxText (que incluye shardings.py) en tu VM de TPU:

   $ git clone https://github.com/AI-Hypercomputer/maxtext && cd maxtext

7. Instala Python 3.10:

   $ sudo apt-get update
   $ sudo apt install python3.10
   $ sudo apt install python3.10-venv

8. Crea y activa un entorno virtual:

   $ python3 -m venv your-venv-name
   $ source your-venv-name/bin/activate

9. Dentro del directorio del repositorio de MaxText, ejecuta la secuencia de comandos de configuración para instalar JAX y otras dependencias en tu porción de TPU. La secuencia de comandos de configuración tarda unos minutos en ejecutarse.

   $ bash setup.sh

10. Ejecuta el siguiente comando para ejecutar shardings.py en tu porción de TPU.

    $ python3 -m pedagogical_examples.shardings \
      --ici_fsdp_parallelism 4 \
      --batch_size 131072 \
      --embedding_dimension 2048

    Puedes ver los resultados en los registros. Tus TPU deberían alcanzar alrededor de 260 TFLOP por segundo o una impresionante utilización de FLOPS del 90% o más. En este caso, seleccionamos alrededor del lote máximo que cabe en la memoria de ancho de banda alto (HBM) de la TPU.

11. Puedes explorar otras estrategias de fragmentación en ICI. Por ejemplo, puedes probar la siguiente combinación:

    $ python3 -m pedagogical_examples.shardings \
      --ici_tensor_parallelism 4 \
      --batch_size 131072 \
      --embedding_dimension 2048

12. Cuando termines, borra el recurso en fila y la porción de TPU. Debes ejecutar estos pasos de limpieza desde el entorno en el que configuraste la porción (primero, ejecuta exit para salir de la sesión de SSH). La eliminación tardará entre dos y cinco minutos en completarse. Si usas la gcloud CLI, puedes ejecutar este comando en segundo plano con la marca opcional --async.

    gcloud

    $ gcloud compute tpus queued-resources \
        delete ${QR_ID} --force (--async)

    Consola

    1. En la consola de Google Cloud , ve a la página TPU:

       Ir a TPU

    2. Haz clic en la pestaña Recursos en cola.

    3. Selecciona la casilla de verificación junto a la solicitud de recurso en cola.

    4. Haz clic en delete Borrar.

Fragmentación de Multislice con paralelismo de DCN

La secuencia de comandos shardings.py toma tres parámetros que especifican el paralelismo de la DCN, que corresponden a la cantidad de fragmentos de cada tipo de paralelismo de datos:

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

Los valores de estos parámetros deben estar restringidos de modo que dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism sea igual a la cantidad de porciones.

Como ejemplo para dos porciones, usa --dcn_data_parallelism = 2.

	dcn_data_parallelism	dcn_fsdp_parallelism	dcn_tensor_parallelism	Cantidad de porciones
Paralelismo de datos bidireccional	2	1	1	2

dcn_tensor_parallelism siempre debe establecerse en 1 porque la DCN no se ajusta bien a ese tipo de fragmentación. Para las cargas de trabajo típicas de LLM en chips v4, dcn_fsdp_parallelism también debe establecerse en 1 y, por lo tanto, dcn_data_parallelism debe establecerse en la cantidad de porciones, pero esto depende de la aplicación.

A medida que aumentas la cantidad de porciones (suponiendo que mantienes constantes el tamaño del segmento y el lote por porción), aumentas la cantidad de paralelismo de datos.

Ejecución de shardings.py en un entorno de Multislice

Puedes ejecutar shardings.py en un entorno de Multislice con multihost_runner.py o ejecutando shardings.py en cada VM de TPU. Aquí usamos multihost_runner.py. Los siguientes pasos son muy similares a los de Cómo comenzar: Experimentos rápidos en porciones múltiples del repositorio de MaxText, excepto que aquí ejecutamos shardings.py en lugar del LLM más complejo en train.py.

La herramienta multihost_runner.py está optimizada para experimentos rápidos, ya que reutiliza las mismas TPU de forma repetida. Dado que la secuencia de comandos multihost_runner.py depende de conexiones SSH de larga duración, no la recomendamos para ningún trabajo de larga duración. Si deseas ejecutar un trabajo más largo (por ejemplo, de horas o días), te recomendamos que uses multihost_job.py.

En este instructivo, usamos el término ejecutor para indicar la máquina en la que ejecutas la secuencia de comandos multihost_runner.py. Usamos el término trabajadores para indicar las VMs de TPU que componen tus porciones. Puedes ejecutar multihost_runner.py en una máquina local o en cualquier VM de Compute Engine en el mismo proyecto que tus porciones. No se admite la ejecución de multihost_runner.py en un trabajador.

multihost_runner.py se conecta de forma automática a los trabajadores de TPU a través de SSH.

En este ejemplo, ejecutas shardings.py en dos porciones v5e-16, un total de cuatro VMs y 16 chips TPU. Puedes modificar el ejemplo para que se ejecute en más TPU.

Configura tu entorno

1. Clona MaxText en tu máquina ejecutora:

   $ git clone https://github.com/AI-Hypercomputer/maxtext

2. Ve al directorio del repositorio.

   $ cd maxtext

3. Crea claves SSH para gcloud. Te recomendamos que dejes la contraseña en blanco (presiona Intro dos veces después de ejecutar el siguiente comando). Si se te indica que el archivo google_compute_engine ya existe, selecciona la opción para no conservar tu versión existente.

     $ ssh-keygen -f ~/.ssh/google_compute_engine
     

4. Agrega una variable de entorno para establecer el recuento de porciones de TPU en 2.

     $ export SLICE_COUNT=2
     

5. Crea un entorno de Multislice con el comando queued-resources create o la consola de Google Cloud .

   gcloud

   En el siguiente comando, se muestra cómo crear una TPU v5e de Multislice. Para usar una versión de TPU distinta, especifica un accelerator-type y un runtime-version diferentes.

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --node-count=${SLICE_COUNT} \
       --node-prefix=${TPU_NAME} \
       --zone=${ZONE} \
       [--reserved|--spot]

   Consola

   1. En la consola de Google Cloud , ve a la página TPU:

      Ir a TPU

   2. Haz clic en Crear una TPU.

   3. En el campo Nombre, ingresa el nombre de tu TPU.

   4. En el cuadro Zona, selecciona la zona en la que deseas crear la TPU.

   5. En el cuadro Tipo de TPU, selecciona un tipo de acelerador. El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Multislice solo es compatible con Cloud TPU v4 y versiones posteriores de TPU. Para obtener más información sobre las versiones de TPU, consulta Versiones de TPU.

   6. En el cuadro Versión de software de TPU, selecciona una versión de software. Cuando creas una VM de Cloud TPU, la versión del software de TPU especifica la versión del entorno de ejecución de TPU que se instalará en las VMs de TPU. Para obtener más información, consulta Versiones de software de TPU.

   7. Haz clic en el botón para activar o desactivar la opción Habilitar las colas.

   8. En el campo Nombre del recurso en cola, ingresa un nombre para tu solicitud de recurso en cola.

   9. Haz clic en la casilla de verificación Convertirlo en una TPU Multislice.

   10. En el campo Cantidad de porciones, ingresa la cantidad de porciones que deseas crear.

   11. Haz clic en Crear para crear tu solicitud de recursos en cola.

6. Cuando comienza el aprovisionamiento de recursos en fila, puede tardar hasta cinco minutos en completarse, según el tamaño del recurso en fila. Espera hasta que el recurso en fila esté en el estado ACTIVE. Puedes verificar el estado de una solicitud de recursos en fila con la gcloud CLI o la consola de Google Cloud :

   gcloud

   $ gcloud compute tpus queued-resources list \
       --filter=${QR_ID} --zone=${ZONE} --project=${PROJECT}

   Esto debería generar un resultado similar al siguiente:

   NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
   ...
   que-res-id  us-central2-b  4           v5litepod-16             ACTIVE
   ...

   Consola

   1. En la consola de Google Cloud , ve a la página TPU:

      Ir a TPU

   2. Haz clic en la pestaña Recursos en cola.

   3. Haz clic en el nombre de la solicitud de recurso en cola.

   Comunícate con tu representante de cuenta de Google Cloud si el estado del código QR es WAITING_FOR_RESOURCES o PROVISIONING durante más de 15 minutos.

7. Instala las dependencias.

   $ python3 multihost_runner.py \
       --TPU_PREFIX=${TPU_NAME} \
       --ZONE=${ZONE} \
       --COMMAND="bash setup.sh"

8. Ejecuta shardings.py en cada trabajador con multihost_runner.py.

   $ python3 multihost_runner.py \
       --TPU_PREFIX=${TPU_NAME} \
       --ZONE=${ZONE} \
       --COMMAND="python3 -m pedagogical_examples.shardings \
       --dcn_data_parallelism ${SLICE_COUNT} \
       --ici_fsdp_parallelism 16 \
       --batch_size 131072 \
       --embedding_dimension 2048"

   Verás alrededor de 230 TFLOPS por segundo de rendimiento en los archivos de registro.

   Para obtener más información sobre cómo configurar el paralelismo, consulta Fragmentación Multislice con paralelismo de DCN y shardings.py.

9. Cuando termines, limpia las TPU y el recurso en fila. La eliminación tardará entre dos y cinco minutos en completarse. Si usas la gcloud CLI, puedes ejecutar este comando en segundo plano con la marca opcional --async.

Cómo escalar una carga de trabajo a Multislice

Antes de ejecutar tu modelo en un entorno de Multislice, realiza los siguientes cambios en el código:

• Usa jax.experimental.mesh_utils.create_hybrid_device_mesh en lugar de jax.experimental.mesh_utils.create_device_mesh cuando crees tu malla.

Estos deberían ser los únicos cambios de código necesarios cuando se migra a Multislice. Para lograr un alto rendimiento, la DCN debe asignarse a ejes de paralelismo de datos, paralelismo de datos completamente fragmentados o de canalización. Las consideraciones sobre el rendimiento y las estrategias de fragmentación se analizan con más detalle en Fragmentación con Multislice para obtener el máximo rendimiento.

Para validar que tu código puede acceder a todos los dispositivos, puedes declarar que len(jax.devices()) es igual a la cantidad de chips en tu entorno de Multislice. Por ejemplo, si usas cuatro porciones de v4-16, tienes ocho chips por segmento * 4 porciones, por lo que len(jax.devices()) debería devolver 32.

Cómo elegir tamaños de porciones para entornos de Multislice

Para obtener una aceleración lineal, agrega nuevas porciones del mismo tamaño que la porción existente. Por ejemplo, si usas una porción de v4-512, Multislice logrará alrededor del doble de rendimiento agregando una segunda porción de v4-512 y duplicando el tamaño del lote global. Para obtener más información, consulta Fragmentación con Multislice para obtener el máximo rendimiento.

Cómo ejecutar tu trabajo en varias porciones

Existen tres enfoques diferentes para ejecutar tu carga de trabajo personalizada en un entorno de Multislice:

1. Con la secuencia de comandos del ejecutor de experimentos, multihost_runner.py
2. Con la secuencia de comandos del ejecutor de producción, multihost_job.py
3. Con un enfoque manual

Secuencia de comandos del ejecutor de experimentación

La secuencia de comandos multihost_runner.py distribuye el código a un entorno de Multislice existente, ejecuta tu comando en cada host, copia tus registros y hace un seguimiento del estado de error de cada comando. La secuencia de comandos multihost_runner.py se documenta en README de MaxText.

Dado que multihost_runner.py mantiene conexiones SSH persistentes, solo es adecuado para experimentos de tamaño modesto y de ejecución, en cierto modo, más corta. Puedes adaptar los pasos del instructivo de multihost_runner.py a tu carga de trabajo y configuración de hardware.

Secuencia de comandos del ejecutor de producción

Para los trabajos de producción que necesitan resiliencia ante fallas de hardware y otras interrupciones, lo mejor es integrar de forma directa la API para crear recurso en fila. Usa multihost_job.py como ejemplo práctico que activa la llamada a la API para crear recurso en fila con la secuencia de comandos de inicio adecuada para ejecutar tu entrenamiento y reanudarlo en caso de interrupción. La secuencia de comandos multihost_job.py se documenta en el README de MaxText.

Como multihost_job.py debe aprovisionar recursos para cada ejecución, no proporciona un ciclo de iteración tan rápido como multihost_runner.py.

Enfoque manual

Te recomendamos que uses o adaptes multihost_runner.py o multihost_job.py para ejecutar tu carga de trabajo personalizada en tu configuración de Multislice. Sin embargo, si prefieres aprovisionar y administrar tu entorno de forma directa con comandos de QR, consulta Administra un entorno de Multislice.

Administra un entorno de Multislice

Para aprovisionar y administrar códigos QRs de forma manual sin usar las herramientas que se proporcionan en el repositorio de MaxText, lee las siguientes secciones.

Crea recursos en fila

gcloud

1. Crea una solicitud de recurso en fila con el siguiente comando:

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --project=${PROJECT} \
       --zone=${ZONE} \
       --node-count=${SLICE_COUNT} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --network=${NETWORK_NAME} \
       --subnetwork=${SUBNETWORK_NAME} \
       --tags=${EXAMPLE_TAG_1},${EXAMPLE_TAG_2} \
       --metadata=startup-script="${STARTUP_SCRIPT}" \
       [--reserved|--spot]

Asegúrate de tener la cuota correspondiente antes de seleccionar --reserved, --spot o la cuota predeterminada según demanda. Para obtener información sobre los tipos de cuotas, consulta Política de cuotas.

curl

1. Crea un archivo llamado queued-resource-req.json y copia el siguiente archivo JSON en él.

   {
   "guaranteed": { "reserved": true },
   "tpu": {
       "node_spec": [
       {
       "parent": "projects/your-project-number/locations/your-zone",
           "node": {
           "accelerator_type": "accelerator-type",
           "runtime_version": "tpu-vm-runtime-version",
           "network_config": {
               "network": "your-network-name",
               "subnetwork": "your-subnetwork-name",
               "enable_external_ips": true
           },
           "tags" : ["example-tag-1"]
           "metadata": {
               "startup-script": "your-startup-script"
           }
       },
       "multi_node_params": {
           "node_count": slice-count,
           "node_id_prefix": "your-queued-resource-id"
       }
       }
       ]
   }
   }

   Reemplaza los siguientes valores:

   • your-project-number: Es el número de tu proyecto de Google Cloud .
   • your-zone: Es la zona en la que deseas crear el recurso en cola.
   • accelerator-type: Es la versión y el tamaño de una sola porción. Multislice solo es compatible con Cloud TPU v4 y versiones posteriores de TPU.
   • tpu-vm-runtime-version: Es la versión del entorno de ejecución de la VM de TPU que deseas usar.
   • your-network-name: Es opcional y representa una red a la que se adjuntará el recurso en cola.
   • your-subnetwork-name: Es opcional y representa una subred a la que se adjuntará el recurso en cola.
   • example-tag-1: Es opcional, una cadena de etiqueta arbitraria.
   • your-startup-script: Es una secuencia de comandos de inicio que se ejecutará cuando se asigne el recurso en cola.
   • slice-count: Es la cantidad de porciones de TPU en tu entorno de Multislice.
   • your-queued-resource-id: Es el ID proporcionado por el usuario para el recurso en cola.

   Para obtener más información, consulta la documentación de la API de REST de recursos en cola para conocer todas las opciones disponibles.

   Para usar la capacidad de Spot, reemplaza lo siguiente:

   "guaranteed": { "reserved": true } con "spot": {}

   Quita la línea para usar la capacidad predeterminada según demanda.

2. Envía la solicitud de creación de recursos en cola con la carga útil JSON:

   $ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" \
   -H "Content-Type: application/json" \
   -d @queuedresourcereq.json \
   https://tpu.googleapis.com/v2alpha1/projects/your-project-id/locations/your-zone/queuedResources\?queued_resource_id\=your-queued-resource-id

   Reemplaza los siguientes valores:

   • your-project-id: Es el ID de tu proyecto de Google Cloud .
   • your-zone: Es la zona en la que deseas crear el recurso en cola.
   • your-queued-resource-id: Es el ID proporcionado por el usuario para el recurso en cola.

La respuesta debería ser como la siguiente:

{
"name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qr-guid>",
"metadata": {
    "@type": "type.googleapis.com/google.cloud.common.OperationMetadata",
    "createTime": "2023-11-01T00:17:05.742546311Z",
    "target": "projects/<your-project-id>/locations/<your-zone>/queuedResources/<your-qa-id>",
    "verb": "create",
    "cancelRequested": false,
    "apiVersion": "v2alpha1"
},
"done": false
}

Usa el valor del GUID al final del valor de la cadena para el atributo name para obtener información sobre la solicitud de recursos en cola.

Consola

1. En la consola de Google Cloud , ve a la página TPU:

   Ir a TPU

2. Haz clic en Crear una TPU.

3. En el campo Nombre, ingresa el nombre de tu TPU.

4. En el cuadro Zona, selecciona la zona en la que deseas crear la TPU.

5. En el cuadro Tipo de TPU, selecciona un tipo de acelerador. El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Multislice solo es compatible con Cloud TPU v4 y versiones posteriores de TPU. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.

6. En el cuadro Versión de software de TPU, selecciona una versión de software. Cuando creas una VM de Cloud TPU, la versión de software de TPU especifica la versión del entorno de ejecución de TPU que se instalará. Para obtener más información, consulta Versiones de software de TPU.

7. Haz clic en el botón para activar o desactivar la opción Habilitar las colas.

8. En el campo Nombre del recurso en cola, ingresa un nombre para tu solicitud de recurso en cola.

9. Haz clic en la casilla de verificación Convertirlo en una TPU Multislice.

10. En el campo Cantidad de porciones, ingresa la cantidad de porciones que deseas crear.

11. Haz clic en Crear para crear tu solicitud de recursos en cola.

Recupera el estado de un recurso en cola

gcloud

$ gcloud compute tpus queued-resources describe ${QR_ID} --zone=${ZONE}

Para un recurso en cola que se encuentra en el estado ACTIVE, el resultado es similar al siguiente:

...
state:
    state: ACTIVE
...

curl

$ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2/projects/your-project-id/locations/your-zone/queuedResources/${YOUR_QR_ID}

Para un recurso en cola que se encuentra en el estado ACTIVE, el resultado es similar al siguiente:

{
"name": your-queued-res,
"tpu": {
    "nodeSpec": [
    {
        ... // node 1
    },
    {
        ... // node 2
    },
    ...
    ]
},
...
"state": "ACTIVE"
}

Consola

1. En la consola de Google Cloud , ve a la página TPU:

   Ir a TPU

2. Haz clic en la pestaña Recursos en cola.

3. Haz clic en el nombre de la solicitud de recurso en cola.

Después de que se aprovisione tu TPU, también puedes ver los detalles de tu solicitud de recursos en cola. Para ello, ve a la página TPU, busca tu TPU y haz clic en el nombre de la solicitud de recursos en cola correspondiente.

En un caso excepcional, es posible que encuentres tu recurso en cola en el estado FAILED mientras que algunas porciones están en ACTIVE. Si esto sucede, borra los recursos que se crearon y vuelve a intentarlo en unos minutos o comunícate con el equipo de asistencia deGoogle Cloud .

SSH y dependencias de instalación

En Ejecuta el código JAX en porciones de TPU, se describe cómo conectarte a tus VMs de TPU con SSH en una sola porción. Para conectarte a todas las VMs de TPU en tu entorno de Multislice a través de SSH y, luego, instalar las dependencias, usa el siguiente comando gcloud:

  $ gcloud compute tpus queued-resources ssh ${QR_ID} \
        --zone=${ZONE} \
        --node=all \
        --worker=all \
        --command="command-to-run" \
        --batch-size=4

Este comando gcloud envía el comando especificado a todos los nodos y trabajadores de QR a través de SSH. El comando se agrupa en lotes de cuatro y se envía de forma simultánea. El siguiente lote de comandos se envía cuando finaliza la ejecución del lote actual. Si falla uno de los comandos, se detiene el procesamiento y no se envían más lotes. Para obtener más información, consulta Referencia de la API de recursos en cola. Si la cantidad de porciones que usas supera el límite de subprocesos de tu computadora local (también llamado límite de procesamiento por lotes), se producirá un interbloqueo. Por ejemplo, supongamos que el límite de procesamiento por lotes en tu máquina local es 64. Si intentas ejecutar una secuencia de comandos de entrenamiento en más de 64 segmentos, por ejemplo, 100, el comando SSH dividirá las porciones en lotes. Ejecutará la secuencia de comandos de entrenamiento en el primer lote de 64 porciones y esperará a que se completen las secuencias de comandos antes de ejecutarlas en el lote restante de 36 segmentos. Sin embargo, el primer lote de 64 segmentos no puede completarse hasta que los 36 porciones restantes comiencen a ejecutar la secuencia de comandos, lo que provoca un interbloqueo.

Para evitar esta situación, puedes ejecutar la secuencia de comandos de entrenamiento en segundo plano en cada VM agregando un signo et (&) al comando de la secuencia de comandos que especificas con la marca --command. Cuando lo hagas, después de iniciar la secuencia de comandos de entrenamiento en el primer lote de porciones, el control volverá de inmediato al comando SSH. Luego, el comando SSH puede comenzar a ejecutar la secuencia de comandos de entrenamiento en el lote restante de 36 segmentos. Deberás canalizar tus transmisiones de stdout y stderr de forma adecuada cuando ejecutes los comandos en segundo plano. Para aumentar el paralelismo dentro del mismo QR, puedes seleccionar porciones específicas con el parámetro --node.

Configuración de red

Sigue estos pasos para asegurarte de que las porciones de TPU puedan comunicarse entre sí. Instala JAX en cada una de las porciones. Para obtener más información, consulta Ejecuta el código JAX en porciones de TPU. Declara que len(jax.devices()) es igual a la cantidad de chips en tu entorno de Multislice. Para ello, ejecuta el siguiente comando en cada porción:

  $ python3 -c 'import jax; print(jax.devices())'

Si ejecutas este código en cuatro porciones de v4-16, habrá ocho chips por porción y cuatro porciones, por lo que jax.devices() debería devolver un total de 32 chips (dispositivos).

Enumera los recursos en cola

gcloud

Puedes ver el estado de los recursos en cola con el comando queued-resources list:

$ gcloud compute tpus queued-resources list --zone=${ZONE}

El resultado es similar al siguiente:

NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
...
que-res-id  us-central1-a  4           v5litepod-16             ACTIVE
...

Consola