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 de pila completa que permite que un trabajo de entrenamiento use varios segmentos de TPU dentro de un solo segmento o en segmentos de varios Pods con paralelismo de datos estándar. Con los chips de 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 divisió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 configuraciones de Multislice, los chips de TPU en cada porción se comunican a través de la interconexión entre chips (ICI). Los chips de 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 de Cloud TPU con Multislice.

Los desarrolladores no tienen que escribir código para implementar la comunicación entre segmentos de la DCN. El compilador de 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 segmento 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 segmentación experimenta un evento de mantenimiento, una interrupción o una falla de hardware, Cloud TPU creará una nueva segmentación. Si no hay suficientes recursos para crear una nueva segmentación, la creación no se completará hasta que el hardware esté disponible. Después de crear el nuevo segmento, se reiniciarán todos los demás segmentos del entorno de Multislice para que pueda continuar el entrenamiento. Con un script de inicio configurado correctamente, el script de entrenamiento puede reiniciarse automáticamente 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 una configuración Multislice.

Programación en grupo

Cuando todas las segmentaciones de TPU se aprovisionan juntas, al mismo tiempo, lo que garantiza que todas o ninguna de las segmentaciones se aprovisionen correctamente.

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.

Node

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 de TPU.

Tensor

Es una 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 crees tus recursos.

• Spot: Se orienta a la cuota interrumpible con VMs Spot. Es posible que tus recursos se interrumpan para dejar lugar a solicitudes de un trabajo de mayor prioridad. Usa la marca --spot cuando crees 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 cola 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.

Comenzar

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 de TPU deben administrarse como recursos en cola.

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 se adapta desde un solo segmento hasta entornos de varios segmentos.

Paralelismo de ICI

ICI hace referencia a la interconexión de alta velocidad que conecta las TPUs en una sola porción. El sharding de ICI corresponde al sharding 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 determinarán 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 división.

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

	ici_data_parallelism	ici_fsdp_parallelism	ici_tensor_parallelism
FSDP de 4 vías	1	4	1
Paralelismo de tensor de 4 vías	1	1	4
FSDP bidireccional + paralelismo de tensor 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 lo suficientemente rápida como para preferir casi siempre el FSDP al paralelismo de datos.

En este ejemplo, se supone que sabes cómo ejecutar 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 segmentació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 cola.
   TPU_NAME	Es el nombre que el usuario asignó a la TPU.
   PROYECTO	Nombre del proyectoGoogle Cloud
   ZONA	Especifica la zona en la que se crearán los recursos.
   NETWORK_NAME	Nombre de las redes de VPC.
   SUBNETWORK_NAME	Nombre de la subred en las redes de VPC
   RUNTIME_VERSION	La versión de software de Cloud TPU.
   ACCELERATOR_TYPE	v4-16
   EXAMPLE_TAG_1, EXAMPLE_TAG_2 …	Son las etiquetas que se usan para identificar fuentes o destinos válidos para los firewalls de red.
   SLICE_COUNT	Cantidad de segmentos. Se limita a un máximo de 256 segmentos.
   STARTUP_SCRIPT	Si especificas una secuencia de comandos de inicio, esta se ejecutará cuando se aprovisione o reinicie el segmento 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 CLI de Google Cloud 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 QR.

   Console

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

      Ir a TPUs

   2. Haz clic en Crear TPU.

   3. En el campo Nombre, ingresa un nombre para 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 del software de TPU especifica la versión del tiempo 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 de activación Habilitar la cola.

   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 cola 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 cola, puede tardar de uno a cinco minutos en completarse, según el tamaño del recurso en cola. Puedes verificar el estado de una solicitud de recursos en cola con gcloud CLI o la consola de Google Cloud :

   gcloud

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

   Console

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

      Ir a TPUs

   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 segmento 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 segmento 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 aproximadamente el 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 cola y la división de TPU. Debes ejecutar estos pasos de limpieza desde el entorno en el que configuraste la segmentació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 gcloud CLI, puedes ejecutar este comando en segundo plano con la marca opcional --async.

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

    gcloud

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

    Console

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

       Ir a TPUs

    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 múltiples cortes 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 restringirse de modo que dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism sea igual a la cantidad de segmentos.

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

	dcn_data_parallelism	dcn_fsdp_parallelism	dcn_tensor_parallelism	Cantidad de cortes
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 segmentos, pero esto depende de la aplicación.

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

Ejecuta 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 Getting Started: Quick Experiments on Multiple slices 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 segmentos. Puedes ejecutar multihost_runner.py en una máquina local o en cualquier VM de Compute Engine en el mismo proyecto que tus segmentos. No se admite la ejecución de multihost_runner.py en un trabajador.

multihost_runner.py se conecta automáticamente 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 de 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 no conservar tu versión existente.

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

4. Agrega una variable de entorno para establecer el recuento de segmentos 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 a pedido. Para obtener información sobre los tipos de cuotas, consulta la Política de cuotas.

   En el siguiente comando, se muestra cómo crear una TPU de Multislice v5e. Para usar una versión de TPU diferente, 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]

   Console

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

      Ir a TPUs

   2. Haz clic en Crear TPU.

   3. En el campo Nombre, ingresa un nombre para 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 tiempo 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 de activación Habilitar la cola.

   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 de porciones múltiples.

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

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

6. Cuando comienza el aprovisionamiento de recursos en cola, puede tardar hasta cinco minutos en completarse, según el tamaño del recurso en cola. Espera hasta que el recurso en cola esté en el estado ACTIVE. Puedes verificar el estado de una solicitud de recursos en cola con 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
   ...

   Console

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

      Ir a TPUs

   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 Google Cloud representante de cuenta 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 aproximadamente 230 TFLOP por segundo de rendimiento en los archivos de registro.

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

9. Cuando termines, limpia las TPU y el recurso en cola. La eliminación tardará entre dos y cinco minutos en completarse. Si usas 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 paralelismo 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 afirmar que len(jax.devices()) es igual a la cantidad de chips en tu entorno de Multislice. Por ejemplo, si usas cuatro segmentos de v4-16, tienes ocho chips por segmento * 4 segmentos, por lo que len(jax.devices()) debería devolver 32.

Cómo elegir tamaños de segmentos 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 división v4-512, Multislice logrará aproximadamente el doble de rendimiento agregando una segunda división v4-512 y duplicando el tamaño del lote global. Para obtener más información, consulta Sharding With Multislice for Maximum Performance (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. Usa la secuencia de comandos del ejecutor de producción, multihost_job.py
3. Usa un enfoque manual

Secuencia de comandos del ejecutor de la 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 el 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 relativamente 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 directamente la API de Create Queued Resource. Usa multihost_job.py como ejemplo práctico que activa la llamada a la API de Created Queued Resource con la secuencia de comandos de inicio adecuada para ejecutar tu entrenamiento y reanudarlo en caso de interrupción. El script 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 directamente con comandos de QR, consulta Administra un entorno de Multislice.

Administra un entorno de Multislice

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

Crea recursos en cola

gcloud

1. Crea una solicitud de recursos en cola 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 a pedido. Para obtener información sobre los tipos de cuotas, consulta la Política de cuotas.

curl

1. Crea un archivo llamado queued-resource-req.json y copia el siguiente código 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: Tu Google Cloud número de proyecto
   • 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 un solo segmento. Multislice solo se admite en 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: 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: 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: Aunque 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 a pedido.

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: 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: 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 y obtener información sobre la solicitud de recursos en cola.

Console

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

   Ir a TPUs

2. Haz clic en Crear TPU.

3. En el campo Nombre, ingresa un nombre para 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 del software de TPU especifica la versión del tiempo 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 de activación Habilitar la cola.

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 de porciones múltiples.

10. En el campo Cantidad de segmentos, ingresa la cantidad de segmentos 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"
}

Console

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

   Ir a TPUs

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 de 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 algunos segmentos 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 Google Cloud equipo de asistencia.

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 la referencia de la API de recursos en cola. Si la cantidad de segmentos 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 un script de entrenamiento en más de 64 segmentos, por ejemplo, 100, el comando SSH dividirá los segmentos en lotes. Ejecutará la secuencia de comandos de entrenamiento en el primer lote de 64 segmentos y esperará a que se completen las secuencias de comandos antes de ejecutar la secuencia de comandos en el lote restante de 36 segmentos. Sin embargo, el primer lote de 64 segmentos no puede completarse hasta que los 36 segmentos restantes comiencen a ejecutar la secuencia de comandos, lo que provoca un bloqueo.

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 segmentos, 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 segmentos específicos 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. Afirma que len(jax.devices()) es igual a la cantidad de chips en tu entorno de Multislice. Para ello, ejecuta el siguiente comando en cada segmento:

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

Si ejecutas este código en cuatro segmentos de v4-16, habrá ocho chips por segmento y cuatro segmentos, 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
...

Console

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

   Ir a TPUs

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 divisió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"

Cómo restablecer las respuestas rápidas

La API de ResetQueuedResource se puede usar para restablecer todas las VMs en un QR de 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}

Se están borrando 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 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]

Console

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

   Ir a TPUs

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 división de reemplazo, se detiene el entrenamiento.

Para reanudar el entrenamiento automáticamente 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 automáticamente cada vez que se reasigna una porción o se restablece una VM. Especificas un script 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 Create QRs) te permite recuperarte automáticamente 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 repo de MaxText antes de probar esto.

Generación de perfiles y depuración

El perfilado es el mismo en los entornos de un solo segmento y de varios segmentos. Para obtener más información, consulta Cómo generar perfiles de programas de JAX.

Optimiza las capacitaciones

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 el sharding 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 fragmentos. Para todas estas estrategias, puedes mantener la misma estrategia de fragmentación dentro de una divisió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 particionado 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 correctamente, pero te gustaría mejorar su rendimiento escalando en varias particiones.

Para lograr un escalamiento fuerte en varios segmentos, el tiempo necesario para realizar la operación de reducción total en la DCN debe ser menor que el tiempo necesario para realizar un pase hacia atrás. La DCN se usa para la comunicación entre porciones y es un factor limitante en el rendimiento de la carga de trabajo.

Cada chip de 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á de 32 a 64 bits de ancho de banda de la DCN para cada parámetro por paso. Si usas dos segmentos, tu modelo usará 32 bits de ancho de banda de la DCN. Si usas más de dos segmentos, el compilador realizará una operación de reducción total de mezcla 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 pase hacia adelante y un pase hacia atrás 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 pase hacia atrás 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 pase hacia atrás 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, haz lo siguiente: Operational intensity = 4 * B / 32

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

Esto sugiere un tamaño de lote mínimo de entre 176 000 y 352 000 para los modelos de lenguaje basados en Transformer. Dado que la red de DCN puede descartar paquetes brevemente, 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 350 000 (dos Pods) a 700 000 (muchos Pods).

Para otras arquitecturas de modelos, deberás estimar el tiempo de ejecución de tu pase hacia atrás por segmento (ya sea cronometrándolo con un generador de perfiles o contando las FLOPS). Luego, puedes comparar ese valor con el tiempo de ejecución esperado para 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 pases hacia adelante y hacia atrás, los pesos se recopilan de forma completa para que cada segmento tenga los pesos que necesita. En lugar de sincronizar los gradientes con all-reduce, los gradientes se reducen y dispersan a medida que se producen. De esta manera, cada segmento solo obtiene los gradientes de los pesos de los que es responsable.

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 segmentos. FSDP disminuirá la presión de memoria a medida que aumentes la cantidad de segmentos. Esto se debe a que la cantidad de pesos y el estado del optimizador por segmento disminuyen, pero a costa de un mayor tráfico de red y una mayor posibilidad de bloqueo debido a un colectivo retrasado.

En la práctica, el FSDP en todas las porciones es mejor si aumentas el lote por porción, almacenas más activaciones para minimizar la nueva materialización durante el pase hacia atrás o aumentas la cantidad de parámetros en tu red neuronal.

Las operaciones de all-gather y all-reduce en FSDP funcionan de manera similar a las de DP, por 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 los segmentos 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 efectivo 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 de 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 lo suficientemente 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 el rendimiento. En concreto, si tienes N segmentos, 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 lo suficientemente grande como para que las operaciones colectivas de la DCN se puedan ocultar detrás de la aritmética en el pase hacia atrás. Esto minimizará la ralentización que introduce el paralelismo de canalización y te permitirá superar el límite del tamaño del lote global.

Prácticas recomendadas para el corte múltiple

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.

Controles

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.

Configuraciones admitidas

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

Formas

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

Organización

La organización es compatible con GKE. Para obtener más información, consulta TPUs 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 tuGoogle Cloud representante de cuenta.

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.