Descripción general de Multislice de Cloud TPU

Importante: Multislice solo es compatible con Cloud TPU v4 y versiones posteriores de TPU.

Cloud TPU Multislice es una tecnología de escalamiento de rendimiento de pila completa que permite que una tarea de entrenamiento use varias porciones de TPU dentro de un solo pod o en porciones de varios pods con un paralelismo de datos simple. 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 las tareas de entrenamiento que requieren menos de 4096 chips, una sola porción puede ofrecer el mejor rendimiento. Sin embargo, varias porciones más pequeñas están más disponibles, 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 TPU de cada porción se comunican a través de interconexiones entre chips (ICI). Los chips de TPU en diferentes porciones 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).

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

Conceptos

Tipo de acelerador

La forma de cada porción de TPU que comprende una Multislice. Cada fragmento de una solicitud de varias porciones es del mismo tipo de acelerador. Un tipo de acelerador consiste en un tipo de TPU (v4 o v5e) seguido de la cantidad de TensorCores. Por ejemplo, v4-128 especifica una TPU v4 con 128 TensorCores.

Reparación automática

Cuando una porción encuentra un evento de mantenimiento, una interrupción o una falla de hardware, Cloud TPU creará una porción nueva. En el caso poco frecuente de que haya recursos insuficientes 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 todas las demás porciones del entorno de Multislice para que la capacitación pueda continuar.Con una secuencia de comandos de inicio configurada correctamente, la secuencia de comandos de entrenamiento se puede reiniciar automáticamente sin intervención del usuario, cargando y reanudando desde el punto de control más reciente.

Conjunto de datos

Son los datos que usa un modelo para el entrenamiento o la inferencia.

Redes de centros de datos (DCN)

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

Programación en grupo

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

Host

Un host es una computadora física que ejecuta VMs. Un host puede ejecutar, como máximo, cuatro VMs al mismo tiempo. Cada VM tiene una TPU dedicada.

Inferencia

Carga un modelo de aprendizaje automático previamente entrenado en un host y realiza predicciones sobre los datos.

Interconexión entre chips (ICI)

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

Multislice

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

Node

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

Pod

Es una colección de chips TPU conectados por interfaces de red ICI dedicadas. Un Pod te permite distribuir la carga de procesamiento en varias TPU.

Recurso en cola (QR)

Es una representación de los recursos de TPU que se usa para poner en cola y administrar una solicitud de un entorno de TPU de una o varias porciones.

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

Porción de TPU

Es una subsección lógica de un pod de TPU que consta de chips TPU. Todos los chips de una porción se comunican entre sí a través de la red ICI.

VM de TPU

Una máquina virtual que ejecuta Linux y tiene acceso a las TPU subyacentes. En el caso de las TPU v4, cada VM de TPU tiene acceso directo a cuatro chips. A veces, llamamos a una VM de TPU un trabajador.

Tensor

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

Unidad de procesamiento tensorial (TPU)

Chip de aceleración de AA desarrollado internamente por Google. Están diseñados para ofrecer un procesamiento rápido y eficiente en términos de energía para tareas clave de aprendizaje automático, como la multiplicación de matrices.

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 consumir 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 se usurpen tus recursos para dejar espacio para solicitudes de un trabajo de prioridad más alta. Usa la marca --spot cuando crees tus recursos.
• Según demanda: Se orienta a la cuota según demanda, que no necesita una reserva y no se anula. La solicitud de TPU se pondrá en cola en una cola de cuotas on demand que ofrece Cloud TPU. No se garantiza la disponibilidad de los recursos. Se selecciona de forma predeterminada, no se necesitan marcas.

Nota: Cloud TPU limita el tamaño de la porción para todos los tipos de aceleradores. Para conocer los límites de tamaño de las porciones de las solicitudes on demand, consulta Tamaños de porciones predeterminados para los códigos QR on demand.

Comenzar

Si nunca antes usaste TPU, primero instala Google Cloud CLI y configura tu entorno de Cloud TPU. Para usar Multislice, tus recursos de TPU deben administrarse como Recursos en cola.

Si eres un usuario existente de TPU v4 y tienes una reserva, es posible que debas migrar tu reserva a un sistema de reserva nuevo. Para obtener más información, comunícate con tu representante de cuenta de Google Cloud.

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 de 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 fragmentado (FSDP) y el paralelismo de tensores. El código se escala de una sola porción a entornos de varias porciones.

Paralelismo de ICI

La ICI hace referencia a la interconexión de alta velocidad que conecta las TPU en una sola porción. El fragmentación de ICI corresponde a la fragmentación 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 para 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 para los cuatro chips disponibles en las versiones 4 a 8:

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

Ten en cuenta que ici_data_parallelism debe dejarse como 1 en la mayoría de los casos, 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 da por sentado que conoces la ejecución de código en una sola porción de TPU, como 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
   $ gcloud config set project your-project-id
   $ gcloud config set compute/zone your-zone

2. Crea claves SSH para gcloud. Te recomendamos que dejes una 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 YOUR_QR_ID \
       --accelerator-type your-accelerator-type \
       --runtime-version tpu-ubuntu2204-base \
       --node-id qr-id \
       [--reserved |--spot]

   Descripciones de las marcas de comandos

   YOUR_QR_ID

   Es una cadena definida por el usuario que identifica la solicitud de código QR.

   accelerator-type

   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.

   runtime-version

   La versión de software de Cloud TPU.

   node-id

   El ID de los recursos de TPU que se crearán en respuesta a la solicitud de código QR.

   reserved

   Usa una reserva cuando crees las porciones.

   spot

   Usa VMs Spot cuando crees las porciones.

   Google Cloud CLI no admite todas las opciones para crear códigos QR, como las etiquetas. Para obtener más información, consulta Cómo crear códigos QR.

   Console

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

      Ve 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 se crea 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á. Para obtener más información, consulta Imágenes de VM de TPU.

   7. Haz clic en el botón de activación Habilitar 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 cola esté en el estado ACTIVE, lo que significa que los nodos de trabajo están en el estado READY. Una vez que comienza el aprovisionamiento de recursos en cola, puede tardar entre uno y 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=YOUR_QR_ID

   Console

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

      Ve a TPUs

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

   3. Haz clic en el nombre de tu solicitud de recursos en cola.

5. Una porción v4-8 tiene una sola VM de TPU. Conéctate a la VM de TPU con SSH:

   $ gcloud compute tpus tpu-vm ssh YOUR_QR_ID

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

7. En el 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

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

   $ python3 pedagogical_examples/shardings.py \
     --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 más del 90%de FLOP. En este caso, seleccionamos aproximadamente el lote máximo que se ajusta a la memoria de ancho de banda alto (HBM) de la TPU.

9. No dudes en explorar otras estrategias de fragmentación en ICI. Por ejemplo, puedes probar la siguiente combinación:

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

10. Borra el recurso en cola y la porción de TPU cuando termines. 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 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, estos seguirán consumiendo tu cuota de recursos en cola, que se registra separadamente de la cuota de TPU.

    gcloud

    $ gcloud compute tpus queued-resources
        delete YOUR_QR_ID --force (--async)

    Console

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

       Ve a TPUs

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

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

    4. Haz clic en Borrar.

Fragmentación de varias porciones con paralelismo de DCN

La secuencia de comandos shardings.py toma tres parámetros que especifican el paralelismo de DCN, que corresponde 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 rebanadas.

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

	dcn_data_parallelism	dcn_fsdp_parallelism	dcn_tensor_parallelism	Cantidad de rebanadas
Paralelismo de datos de 2 vías	2	1	1	2

dcn_tensor_parallelism siempre debe establecerse en 1 porque el DCN no es una buena opción para ese tipo de fragmentación. Para cargas de trabajo de LLM típicas 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 rebanadas, pero esto depende de la aplicación.

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

Cómo ejecutar 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 varias divisiones 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 y reutiliza las mismas TPU de forma reiterada. Debido a que la secuencia de comandos multihost_runner.py depende de conexiones SSH de larga duración, no la recomendamos para trabajos de larga duración. Si deseas ejecutar una tarea más larga (por ejemplo, horas o días), te recomendamos que use multihost_job.py.

En este instructivo, usamos el término ejecutor para indicar la máquina en la que se ejecuta la secuencia de comandos multihost_runner.py. Usamos el término trabajadores para indicar las VMs de TPU que conforman 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 divisiones. 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 mediante SSH.

En este ejemplo, ejecutamos shardings.py en dos fragmentos v4-16, un total de cuatro VM y 16 chips TPU. Puedes modificar el ejemplo para que se ejecute en más TPU.

Configure su entorno

1. Clona MaxText en tu máquina de ejecución.

2. Ve al directorio del repositorio.

3. Crea claves SSH para gcloud. Te recomendamos que dejes una 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 la 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 a pedido. 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 multislice v4. Para usar v5e, especifica un accelerator-type v5e (por ejemplo, v5litepod-16) y el runtime-version v5e (v2-alpha-tpuv5-lite).

   $ gcloud compute tpus queued-resources \
       create YOUR_QR_ID \
       --accelerator-type=your-accelerator-type \
       --runtime-version=tpu-vm-runtime-version \
       --node-count=node-count \
       --node-prefix=YOUR_QR_ID \
       [--reserved|--spot]

   Descripciones de las marcas de comandos

   YOUR_QR_ID

   Es una cadena definida por el usuario que identifica la solicitud de código QR.

   accelerator-type

   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.

   runtime-version

   La versión de software de Cloud TPU.

   node-count

   Es la cantidad de rebanadas que se crearán.

   node-prefix

   Es el prefijo que se usa para generar nombres para cada fragmento. Se agrega un número al prefijo de cada fragmento. Por ejemplo, si configuras node-prefix como mySlice, las porciones se denominan mySlice-0, mySlice-1 y continúan numéricamente para cada porción.

   reserved

   Usa una reserva cuando crees las porciones.

   spot

   Usa VMs Spot cuando crees las porciones.

   Console

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

      Ve 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 se crea 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á. Para obtener más información, consulta Imágenes de VM de TPU.

   7. Haz clic en el botón de activación Habilitar colas.

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

   9. Haz clic en la casilla de verificación Convertirlo en una TPU de porciones múltiples.

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

   11. Haz clic en Crear para crear la 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 a que el recurso en fila esté en el estado ACTIVE. Puedes verificar el estado de una solicitud de recurso en fila con gcloud CLI o la consola de Google Cloud:

   gcloud

   $ gcloud compute tpus queued-resources list \
       --filter=YOUR_QR_ID

   Esto debería generar un resultado similar al siguiente:

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

   Console

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

      Ve a TPUs

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

   3. Haz clic en el nombre de tu solicitud de recursos en cola.

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

7. Instala las dependencias.

   $ python3 multihost_runner.py \
       --TPU_PREFIX=YOUR_QR_ID \
       --COMMAND="bash setup.sh"

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

   $ python3 multihost_runner.py \
       --TPU_PREFIX=YOUR_QR_ID \
       --COMMAND="python3 pedagogical_examples/shardings.py \
       --dcn_data_parallelism $SLICE_COUNT \
       --ici_fsdp_parallelism 8 \
       --batch_size 131072 \
       --embedding_dimension 2048"

   Verás aproximadamente 230 TFLOP por segundo de rendimiento en los archivos de registro.

9. Limpia las TPU y el recurso en cola cuando termines. 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 --async opcional.

Escala 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 cambia a Multislice. Para lograr un alto rendimiento, la DCN debe asignarse a ejes paralelos de datos, de datos completamente fragmentados o de canalización. Las consideraciones de 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 pueda 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 porciones de v4-16, tienes ocho chips por porción × 4 porciones, por lo que len(jax.devices()) debería mostrar 32.

Cómo elegir los tamaños de las porciones para entornos de varias porciones

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 v4-512, Multislice logrará aproximadamente el doble del rendimiento agregando una segunda porción 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.

Ejecuta 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 experimentación, multihost_runner.py
2. Usa 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 y ejecuta tu comando en cada host, vuelve a copiar tus registros y realiza un seguimiento del estado de error de cada comando. La secuencia de comandos multihost_runner.py se documenta en el archivo README de MaxText.

Debido a que multihost_runner.py mantiene conexiones SSH persistentes, solo es adecuada para la experimentación de tamaño modesto y de ejecución relativamente corta. Puedes adaptar los pasos del instructivo multihost_runner.py a tu carga de trabajo y configuración de hardware.

Secuencia de comandos del ejecutor de producción

En el caso de las tareas de producción que necesitan resiliencia contra fallas de hardware y otras preempciones, lo mejor es realizar la integración directamente con la API de Create Queued Resource. Como ejemplo de trabajo, proporcionamos multihost_job.py, 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 la usurpación. La secuencia de comandos multihost_job.py se documenta en el archivo readme de MaxText.

Debido a que 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 Cómo administrar un entorno de varias capas.

Administra un entorno Multislice

Para aprovisionar y administrar códigos QR de forma manual sin usar las herramientas proporcionadas en el repositorio de MaxText, lee las siguientes secciones.

Crea recursos en cola

gcloud

1. Configura las siguientes variables de entorno antes de aprovisionar la capacidad:

   $ export YOUR_QR_ID=your-queued-resource-id
   $ export PROJECT=your-project-name
   $ export ZONE=us-central2-b
   $ export NETWORK_NAME=your-network-name
   $ export SUBNETWORK_NAME=your-subnetwork-name
   $ export RUNTIME_VERSION=tpu-ubuntu2204-base
   $ export ACCELERATOR_TYPE=v4-16
   $ export SLICE_COUNT=4
   $ export STARTUP_SCRIPT="#!/bin/bash\n ..."
   $ gcloud config set project project-name
   $ gcloud config set compute/zone zone

   Descripciones de las variables

   Entrada	Descripción
   YOUR_QR_ID	Es el ID asignado por el usuario al recurso en cola.
   PROYECTO	Nombre del proyecto de Google Cloud
   ZONA	Especifica la zona en la que se crearán los recursos.
   NETWORK_NAME	Es el 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 …	Etiquetas que se usan para identificar fuentes o destinos válidos para firewalls de red
   SLICE_COUNT	Cantidad de rebanadas Se limita a un máximo de 256 rebanadas.
   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 una solicitud de recursos en cola con el siguiente comando:

   $ gcloud compute tpus queued-resources \
       create ${YOUR_QR_ID} \
       --project your-project-id \
       --zone your-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]

   Descripciones de las marcas de comandos

   YOUR_QR_ID

   Es una cadena definida por el usuario que identifica la solicitud de recursos en cola.

   project

   El proyecto de Google Cloud en el que creas la solicitud de recursos en cola.

   zone

   La zona de Google Cloud en la que se creará el recurso en cola.

   node-count

   Es la cantidad de rebanadas que se crearán.

   accelerator-type

   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.

   runtime-version

   La versión de software de Cloud TPU.

   network

   Es el nombre de una red de VPC a la que se conectará el recurso de TPU.

   subnetwork

   Es el nombre de una subred de VPC a la que se adjuntará el recurso de TPU.

   reserved

   Usa una reserva cuando crees las porciones.

   spot

   Usa VMs Spot cuando crees las porciones.

Asegúrate de tener la cuota correspondiente antes de seleccionar --reserved, --spot o la cuota on demand predeterminada. Si deseas obtener más 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 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: Opcional, es una red a la que se conectará el recurso en fila.
   • your-subnetwork-name: Es opcional. Es una subred a la que se adjuntará el recurso en cola.
   • example-tag-1: Opcional, es 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_QR_ID: Es el ID que proporcionó el usuario para el recurso en fila.

   Para obtener más información, consulta la documentación de la API de REST Queued Resource para ver todas las opciones disponibles.

   Nota: Aunque node_spec es una lista, solo se admite una entrada node_spec.

   Para usar la capacidad de Spot, reemplaza lo siguiente:

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

   Quita la línea para usar la capacidad on demand predeterminada.

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_QR_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_QR_ID: Es el ID que proporcionó el usuario para el recurso en fila.

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 GUID al final del valor de cadena para el atributo name para obtener información sobre la solicitud de recursos en cola.

Console

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

   Ve 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 para 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 se crea 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á. Para obtener más información, consulta Imágenes de VM de TPU.

7. Haz clic en el botón de activación Habilitar colas.

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

9. Haz clic en la casilla de verificación Convertirlo en una TPU de porciones múltiples.

10. En el campo Cantidad de divisiones, ingresa la cantidad de divisiones 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 ${YOUR_QR_ID}

Para un recurso en cola que está en el estado ACTIVE, el resultado se ve de la siguiente manera:

...
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 está en el estado ACTIVE, el resultado se ve de la siguiente manera:

{
"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:

   Ve a TPUs

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

3. Haz clic en el nombre de tu solicitud de recursos en cola.

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

En casos excepcionales, es posible que encuentres tu recurso en cola en el estado FAILED mientras que algunas slices son ACTIVE. Si esto sucede, borra los recursos que se crearon y vuelve a intentarlo en unos minutos, o comunícate con Asistencia de Google Cloud.

SSH e instala dependencias

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

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

Este comando gcloud envía el comando especificado a todos los trabajadores y nodos en el QR a través de SSH. El comando se agrupa en grupos 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 se produce una falla con 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 rebanadas que usas supera el límite de subprocesos de tu computadora local (también llamado límite de lotes), se producirá un interbloqueo. A modo de ejemplo, asumamos que el límite de lotes en tu máquina local es 64. Si intentas ejecutar una secuencia de comandos de entrenamiento en más de 64 rebanadas, por ejemplo, 100, el comando SSH dividirá las rebanadas en lotes. Ejecutará la secuencia de comandos de entrenamiento en el primer lote de 64 fragmentos y esperará a que se completen las secuencias de comandos antes de ejecutarlas en el lote restante de 36 fragmentos. Sin embargo, el primer lote de 64 rebanadas no se puede completar hasta que las 36 rebanadas restantes comiencen a ejecutar la secuencia de comandos, lo que genera un interbloqueo.

Para evitar esta situación, puedes ejecutar la secuencia de comandos de entrenamiento en segundo plano en cada VM. Para ello, agrega un signo & (&) al comando de la secuencia de comandos que especifiques con la marca --command. Cuando lo hagas, después de iniciar la secuencia de comandos de entrenamiento en el primer lote de rebanadas, 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 rebanadas. Deberás canalizar tus flujos de stdout y stderr de forma adecuada cuando ejecutes los comandos en segundo plano. Para aumentar el paralelismo dentro del mismo código QR, puedes seleccionar rebanadas específicas con el parámetro --node.

Configuración de red

Ejecuta los siguientes 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 pod de TPU. Confirma que len(jax.devices()) sea igual a la cantidad de chips en tu entorno de Multislice. Para ello, en cada fragmento, ejecuta lo siguiente:

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

Si ejecutas este código en cuatro rebanadas de v4-16, hay ocho chips por rebanada y cuatro rebanadas, por lo que jax.devices() debería mostrar un total de 32 chips (dispositivos).

Cómo enumerar recursos en cola

gcloud

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

$ gcloud compute tpus queued-resources list

El resultado es similar al siguiente:

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

Console

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

   Ve a TPUs

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

Inicia tu trabajo en un entorno aprovisionado

Para ejecutar cargas de trabajo de forma manual, conéctate a todos los hosts de cada fragmento a través de SSH y ejecuta el siguiente comando en todos los hosts.

$ gcloud compute tpus tpu-vm ssh YOUR_QR_ID \
    --zone=your-zone \
    --worker=all \
    --node=all \
    --command="command-to-run"

Restablecimiento de los códigos QR

La API de ResetQueuedResource se puede usar para restablecer todas las VMs en un código QR de ACTIVE. Cuando se restablecen las VMs, se borra de manera forzosa 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 de un restablecimiento. La API de ResetQueuedResource puede ser útil cuando deseas reiniciar todas las TPU. Por ejemplo, cuando el entrenamiento se bloquea y restablecer todas las VMs es más fácil que depurar.

Los restablecimientos de todas las VMs se realizan en paralelo, y una operación ResetQueuedResource tarda entre uno y dos minutos en completarse. Para invocar la API, usa el siguiente comando:

$ gcloud compute tpus queued-resources reset YOUR_QR_ID

Borra recursos en cola

Para liberar recursos al final de la sesión de entrenamiento, borra 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.

gcloud

$ gcloud compute tpus queued-resources \
    delete YOUR_QR_ID --force (--async)

Console

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

   Ve a TPUs

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

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

4. Haz clic en Borrar.

Recuperación automática de fallas

En caso de interrupción, Multislice ofrece la reparación sin intervención de la porción afectada y el restablecimiento de todas las porciones después. La porción afectada se reemplaza por una nueva, y las porciones restantes que están en buen estado se restablecen. Si no hay capacidad disponible para asignar un fragmento 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 busque y cargue los últimos puntos de control guardados. Tu secuencia de comandos de inicio se ejecuta automáticamente cada vez que se reasignan las porciones 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 códigos QR) 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 MaxText/train.py 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

La generación de perfiles es la misma en los entornos de una sola porción y de varias porciones. Para obtener más información, consulta Cómo generar perfiles de programas de JAX.

Optimiza las capacitaciones

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

Para lograr el máximo rendimiento en entornos de Multislice, debes considerar cómo dividir los datos en las varias porciones. Por lo general, hay tres opciones (paralelismo de datos, paralelismo de datos completamente fragmentado y paralelismo de canalizaciones). No recomendamos particionar las activaciones en las dimensiones del modelo (a veces llamado paralelismo de tensores) porque requiere demasiado ancho de banda entre las 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 fragmentados es útil para liberar el uso de memoria. El inconveniente es que la comunicación entre las porciones usa la red de DCN y ralentiza la 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 que tengas una carga de trabajo que se ejecuta bien, pero deseas mejorar su rendimiento escalando en varias porciones.

Para lograr un escalamiento sólido en varias porciones, la cantidad de tiempo requerida para realizar la reducción total en DCN debe ser menor que la cantidad de tiempo requerida para realizar un pase hacia atrás. La DCN se usa para la comunicación entre porciones y es un factor limitante en la capacidad de procesamiento de cargas de trabajo.

Cada chip TPU v4 tiene un rendimiento máximo de 275 × 10¹² FLOPS por segundo.

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

Eso 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 DCN para cada parámetro por paso. Si usas dos rebanadas, tu modelo usará 32 bits de ancho de banda de DCN. Si usas más de dos rebanadas, el compilador realizará una operación de reducción completa de todos los shuffles 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 necesarios para un pase hacia adelante y uno 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 por paso: Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

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

Si usas más de dos rebanadas, usa 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. Debido a 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 paso hacia atrás por fragmento (ya sea cronometrándolo con un generador de perfiles o contando FLOPS). Luego, puedes compararlo con el tiempo de ejecución esperado para reducir todo en DCN y obtener una buena estimación de si el paralelismo de datos será útil 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 entre nodos) con la fragmentación de los pesos entre nodos. Para cada operación en los pases hacia adelante y hacia atrás, los pesos se reúnen para que cada fragmento tenga los pesos que necesita. En lugar de sincronizar los gradientes con la reducción completa, los gradientes se reducen y se dispersan a medida que se producen. De esta manera, cada fragmento solo obtiene los gradientes de los pesos de los que es responsable.

Al igual que el paralelismo de datos, la FSDP requerirá escalar el tamaño de lote global linalmente con la cantidad de rebanadas. El FSDP disminuirá la presión de memoria a medida que aumentes la cantidad de rebanadas. Esto se debe a que la cantidad de pesos y el estado del optimizador por fragmento disminuyen, pero lo hace a costa de un aumento del tráfico de red y una mayor posibilidad de bloqueo debido a una colectiva retrasada.

En la práctica, el FSDP en las porciones es mejor si aumentas el lote por porción, almacenas más activaciones para minimizar la rematerializació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, así 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 canalizaciones tiene dos aspectos negativos importantes:

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

El paralelismo de canalización solo debe usarse si las otras estrategias de paralelismo requieren un tamaño de 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 en 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 ser la mejor opción. La mayoría de las cargas de trabajo pueden tolerar tamaños de lotes 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, te recomendamos que lo combines 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 se convierta en un factor menos importante en la capacidad de procesamiento. Específicamente, si tienes N rebanadas, 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 los colectivos de 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 la canalización y, al mismo tiempo, te permitirá escalar más allá del límite de tamaño del lote global.

Prácticas recomendadas para las divisiones múltiples

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 entre los hosts para evitar que las TPUs 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 reescalado dentro del modelo. Además, MaxText aún no ofrece instantáneas deterministas, lo que permitiría que el iterador de datos cargue los mismos datos antes y después de la usurpación.

Controles

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

Configuraciones admitidas

Formas

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

Organización

La orquestació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 que los usuarios prueben Multislice con el paralelismo de datos. Para obtener más información sobre cómo implementar el paralelismo de canalización con Multislice, comunícate con tu representante de cuenta de Google Cloud.

Asistencia y comentarios

Todos los comentarios son bienvenidos. Para compartir comentarios o solicitar asistencia, comunícate con nosotros mediante el formulario de asistencia o comentarios de Cloud TPU.