Descripción general de Cloud TPU Multislice

Importante: Multislice solo se admite en Cloud TPU v4 y versiones de TPU posteriores.

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:

   Nota: Crear un entorno virtual para cada proyecto es una práctica recomendada para administrar las dependencias de Python. Para obtener más información, consulta la documentación de venv.

   $ 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.

    Importante: Los recursos en fila se deben borrar después de usarlos. De lo contrario, seguirán consumiendo tu cuota de recursos en fila, que se hace un seguimiento por separado de la cuota de TPU.

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

   Nota: 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.

   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.

   Nota: Si bien node_spec es una lista, solo se admite una entrada de node_spec.

   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

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.

Inicia tu trabajo en un entorno aprovisionado

Puedes ejecutar cargas de trabajo de forma manual conectándote a todos los hosts de cada porción a través de SSH y ejecutando el siguiente comando en todos los hosts.

$ gcloud compute tpus tpu-vm ssh ${TPU_NAME} \
    --zone=${ZONE} \
    --worker=all \
    --command="command-to-run"

Restablece los QRs

La API de ResetQueuedResource se puede usar para restablecer todas las VMs en un QR ACTIVE. Cuando se restablecen las VMs, se borra por la fuerza la memoria de la máquina y se restablece la VM a su estado inicial. Los datos almacenados de forma local permanecerán intactos y se invocará la secuencia de comandos de inicio después del restablecimiento. La API de ResetQueuedResource puede ser útil cuando deseas reiniciar todas las TPU. Por ejemplo, cuando el entrenamiento se detiene y restablecer todas las VMs es más fácil que depurar.

El restablecimiento de todas las VMs se realiza en paralelo, y una operación de ResetQueuedResource tarda de uno a dos minutos en completarse. Para invocar la API, usa el siguiente comando:

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

Borra los recursos en cola

Para liberar recursos al final de la sesión de entrenamiento, borra el recurso en cola. 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} --zone=${ZONE} --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 Borrar.

Recuperación automática de fallas

En caso de interrupción, Multislice ofrece reparación sin intervención de la porción afectada y restablecimiento de todas las porciones después. La porción afectada se reemplaza por una nueva, y las demás porciones que no están afectadas se restablecen. Si no hay capacidad disponible para asignar una porción de reemplazo, se detiene el entrenamiento.

Para reanudar el entrenamiento de forma automática después de una interrupción, debes especificar una secuencia de comandos de inicio que verifique y cargue los últimos puntos de control guardados. La secuencia de comandos de inicio se ejecuta de forma automática cada vez que se reasigna una porción o se restablece una VM. Especificas una secuencia de comandos de inicio en la carga útil JSON que envías a la API de solicitud de creación de QR.

La siguiente secuencia de comandos de inicio (que se usa en Crear QRs) te permite recuperarte de forma automática de las fallas y reanudar el entrenamiento desde los puntos de control almacenados en un bucket de Cloud Storage durante el entrenamiento de MaxText:

{
 "tpu": {
   "node_spec": [
     {
      ...
         "metadata": {
               "startup-script": "#! /bin/bash \n pwd \n runuser -l user1 -c 'cd /home/user1/MaxText && python3 -m MaxText.train MaxText/configs/base.yml run_name=run_test_failure_recovery dcn_data_parallelism=4 ici_fsdp_parallelism=8 steps=10000 save_period=10 base_output_directory='gs://user1-us-central2'' EOF"
         }
     ...
     }
   ]
 }
}

Clona el repositorio de MaxText antes de probar esto.

Generación de perfiles y depuración

El perfilado es el mismo en los entornos de una sola porción y Multislice. Para obtener más información, consulta Crea perfiles de programas de JAX.

Optimiza el entrenamiento

En las siguientes secciones, se describe cómo puedes optimizar el entrenamiento de Multislice.

Fragmentación con Multislice para obtener el máximo rendimiento

Para lograr el máximo rendimiento en entornos de Multislice, es necesario tener en cuenta cómo realizar la fragmentación en las múltiples porciones. Por lo general, hay tres opciones (paralelismo de datos, paralelismo de datos completamente fragmentado y paralelismo de canalización). No recomendamos fragmentar las activaciones en las dimensiones del modelo (a veces, se denomina paralelismo de tensores) porque requiere demasiado ancho de banda entre porciones. Para todas estas estrategias, puedes mantener la misma estrategia de fragmentación dentro de una porción que te haya funcionado en el pasado.

Te recomendamos que comiences con el paralelismo de datos puro. El uso del paralelismo de datos completamente fragmentado es útil para liberar el uso de memoria. La desventaja es que la comunicación entre las porciones usa la red de DCN y ralentizará tu carga de trabajo. Usa el paralelismo de canalización solo cuando sea necesario según el tamaño del lote (como se analiza a continuación).

Cuándo usar el paralelismo de datos

El paralelismo de datos puro funcionará bien en los casos en los que tengas una carga de trabajo que se ejecute de forma correcta, pero puedes elegir mejorar su rendimiento escalando en varias porciones.

Para lograr un escalamiento fuerte en varias porciones, el tiempo necesario para realizar la operación all-reduce en la DCN debe ser menor que el tiempo necesario para realizar un retroceso. El DCN se usa para la comunicación entre porciones y es un factor limitante en la capacidad de procesamiento de la carga de trabajo.

Cada chip TPU v4 alcanza un máximo de 275 * 10¹² FLOPS por segundo.

Hay cuatro chips por host de TPU, y cada host tiene un ancho de banda de red máximo de 50 Gbps.

Esto significa que la intensidad aritmética es de 4 * 275 * 10¹² FLOPS / 50 Gbps = 22,000 FLOPS / bit.

Tu modelo usará entre 32 y 64 bits de ancho de banda de la DCN para cada parámetro por paso. Si usas dos porciones, tu modelo usará 32 bits de ancho de banda de la DCN. Si usas más de dos porciones, el compilador realizará una operación de all-reduce de Shuffle completa y usarás hasta 64 bits de ancho de banda de DCN para cada parámetro por paso. La cantidad de FLOPS necesarios para cada parámetro variará según tu modelo. Específicamente, para los modelos de lenguaje basados en Transformer, la cantidad de FLOPS requerida para un avance y un retroceso es de aproximadamente 6 * B * P, donde:

• B es el tamaño del lote en tokens.
• P es la cantidad de parámetros.

La cantidad de FLOPS por parámetro es 6 * B y la cantidad de FLOPS por parámetro durante el retroceso es 4 * B.

Para garantizar un escalamiento sólido en varias porciones, asegúrate de que la intensidad operativa supere la intensidad aritmética del hardware de la TPU. Para calcular la intensidad operativa, divide la cantidad de FLOPS por parámetro durante el retroceso por el ancho de banda de la red (en bits) por parámetro y por paso: Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

Por lo tanto, para un modelo de lenguaje basado en Transformer, si usas dos segmentos, sucede lo siguiente: Operational intensity = 4 * B / 32

Si usas más de dos segmentos, haz lo siguiente: Operational intensity = 4 * B/64

Esto supone un tamaño de lote mínimo de entre 176k y 352k para los modelos de lenguaje basados en Transformer. Dado que la red de DCN puede descartar paquetes con rapidez, es mejor mantener un margen de error significativo y, luego, implementar el paralelismo de datos solo si el tamaño del lote por Pod es de al menos 350k (dos Pods) a 700k (muchos Pods).

Para otras arquitecturas de modelos, deberás estimar el entorno de ejecución de tu retroceso por porción (ya sea marcando el tiempo con un generador de perfiles o contando las FLOPS). Luego, puedes comparar ese valor con el tiempo de ejecución esperado para la operación de all-reduce en la DCN y obtener una buena estimación de si el paralelismo de datos tendrá sentido para ti.

Cuándo usar el paralelismo de datos completamente fragmentado (FSDP)

El paralelismo de datos completamente fragmentado (FSDP) combina el paralelismo de datos (fragmentación de los datos en los nodos) con la fragmentación de los pesos en los nodos. Para cada operación en los avances y retrocesos, los pesos se recopilan de forma completa para que cada porción tenga los pesos que necesita. En lugar de sincronizar los gradientes con la operación all-reduce, se dispersan a medida que se producen. De esta manera, cada porción solo obtiene los gradientes de los pesos de los que se encarga.

Al igual que el paralelismo de datos, FSDP requerirá que el tamaño de lote global se escale de forma lineal con la cantidad de porciones. FSDP disminuirá la presión de memoria a medida que aumentes la cantidad de porciones. Esto se debe a que la cantidad de pesos y el estado del optimizador por porción disminuyen, pero a costa de un mayor tráfico de red y una mayor posibilidad de bloqueo debido a un conjunto retrasado.

En la práctica, el FSDP en todas las segmentaciones es mejor si aumentas el lote por porción, ya que almacenas más activaciones para minimizar la rematerialización durante el retroceso o aumentas la cantidad de parámetros en tu red neuronal.

Las operaciones all-gather y all-reduce en FSDP funcionan de manera similar a las de DP, lo que puedes determinar si tu carga de trabajo de FSDP está limitada por el rendimiento de la DCN de la misma manera que se describe en la sección anterior.

Cuándo usar el paralelismo de canalización

El paralelismo de canalización se vuelve relevante cuando se logra un alto rendimiento con otras estrategias de paralelismo que requieren un tamaño de lote global mayor que el tamaño de lote máximo preferido. El paralelismo de canalización permite que las porciones que componen una canalización "compartan" un lote. Sin embargo, el paralelismo de canalización tiene dos desventajas significativas:

1. Se produce la "burbuja de canalización", en la que los chips están inactivos porque esperan datos
2. Requiere microlotes, lo que disminuye el tamaño eficaz del lote, la intensidad aritmética y, en última instancia, la utilización de FLOP del modelo

El paralelismo de canalización solo se debe usar si las otras estrategias de paralelismo requieren un tamaño del lote global demasiado grande. Antes de probar el paralelismo de canalización, vale la pena experimentar para ver de forma empírica si la convergencia por muestra se ralentiza con el tamaño del lote necesario para lograr un FSDP de alto rendimiento. El FSDP tiende a lograr una mayor utilización de FLOP del modelo, pero si la convergencia por muestra se ralentiza a medida que aumenta el tamaño del lote, el paralelismo de canalización puede seguir siendo la mejor opción. La mayoría de las cargas de trabajo pueden tolerar tamaños de lote tan grandes como para no beneficiarse del paralelismo de canalización, pero tu carga de trabajo puede ser diferente.

Si el paralelismo de canalización es necesario, recomendamos combinarlo con el paralelismo de datos o el FSDP. Esto te permitirá minimizar la profundidad de la canalización y, al mismo tiempo, aumentar el tamaño del lote por canalización hasta que la latencia de la DCN deje de ser un factor importante en la capacidad de procesamiento. En concreto, si tienes porciones N, considera canalizaciones de profundidad 2 y N/2 réplicas de paralelismo de datos, luego, canalizaciones de profundidad 4 y N/4 réplicas de paralelismo de datos, y así sucesivamente, hasta que el lote por canalización sea tan grande como para que las operaciones colectivas de la DCN se puedan ocultar detrás de la aritmética en el retroceso. Esto minimizará la ralentización que ingresa el paralelismo de canalización y, al mismo tiempo, te permitirá superar el límite del tamaño del lote global.

Prácticas recomendadas para Multislice

En las siguientes secciones, se describen las prácticas recomendadas para el entrenamiento de Multislice.

Carga de datos

Durante el entrenamiento, cargamos lotes de un conjunto de datos de forma repetida para alimentar el modelo. Es importante tener un cargador de datos asíncrono y eficiente que divida el lote en fragmentos entre los hosts para evitar que las TPU se queden sin trabajo. El cargador de datos actual en MaxText hace que cada host cargue un subconjunto igual de los ejemplos. Esta solución es adecuada para el texto, pero requiere un nuevo fragmento dentro del modelo. Además, MaxText aún no ofrece la creación de instantáneas determinísticas, lo que permitiría que el iterador de datos cargue los mismos datos antes y después de la interrupción.

Punto de control

La biblioteca de puntos de control Orbax proporciona elementos primitivos para crear puntos de control de PyTrees de JAX en el almacenamiento local o en el almacenamiento de Google Cloud . Proporcionamos una integración de referencia con la creación de puntos de control síncronos en MaxText en checkpointing.py.

Parámetros de configuración admitidos

En las siguientes secciones, se describen las formas de porción, la organización, los frameworks y el paralelismo admitidos para Multislice.

Formas

Todas las porciones deben tener la misma forma (por ejemplo, el mismo AcceleratorType). No se admiten formas de porciones heterogéneas.

Organización

La organización es compatible con GKE. Para obtener más información, consulta TPU en GKE.

Frameworks

Multislice solo admite cargas de trabajo de JAX y PyTorch.

Paralelismo

Recomendamos a los usuarios que prueben Multislice con paralelismo de datos. Para obtener más información sobre la implementación del paralelismo de canalización con Multislice, comunícate con tu representante de cuenta deGoogle Cloud .

Asistencia y comentarios

Agradecemos todos los comentarios. Para compartir comentarios o solicitar asistencia, comunícate con nosotros a través del Formulario de comentarios o asistencia de Cloud TPU.