Descripción general de Cloud TPU Multislice [versión preliminar pública]

Vista previa

Este producto o función está sujeto a las "Condiciones de las Ofertas de Fase Previa a la DG" de la sección Condiciones Generales del Servicio de las Condiciones específicas del Servicio. Los productos y las funciones que se encuentran en la fase previa a la DG están disponibles “tal cual” y podrían tener asistencia limitada. Para obtener más información, consulta las descripciones de la etapa de lanzamiento.

Cloud TPU Multislice es una tecnología de escalamiento de rendimiento de pila completa que permite que un trabajo de entrenamiento use porciones de TPU múltiples dentro de un solo pod o en fragmentos en 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 los trabajos de entrenamiento que requieren menos de 4,096 chips, una sola división puede ofrecer el mejor rendimiento. Sin embargo, varias secciones más pequeñas están más disponibles, lo que permite un tiempo de inicio más rápido cuando Multislice se usa con secciones más pequeñas.

Cuando se implementa en las configuraciones de Multislice, los chips TPU de cada porción se comunican a través de la interconexión entre chips (ICI). Los chips TPU en diferentes porciones se comunican mediante la transferencia de datos a CPU (hosts), que, a su vez, transmiten los datos a través de la red del centro de datos (DCN).

Los desarrolladores no escriben código para implementar la comunicación DCN entre segmentos. El compilador de XLA genera ese código por ti y superpone la comunicación con el cálculo para obtener el máximo rendimiento.

Conceptos

AcceleratorType

Forma de cada porción de TPU que incluye un Multislice. Cada porción en una solicitud de varias porciones es del mismo AcceleratorType. AcceleratorType consta de un tipo de TPU (v2, v3, v4, v5e y uno posterior) seguido de la cantidad de TensorCores. Por ejemplo, v4-128 especifica una TPU v4 con 128 TensorCores.

Reparación automática

Cuando una divisió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 no haya suficientes recursos para crear una porción nueva, la creación no se completará hasta que el hardware esté disponible. Luego, se reiniciarán todas las demás porciones del entorno Multislice para continuar con el entrenamiento. Con una secuencia de comandos de inicio configurada de forma correcta, la secuencia de comandos de entrenamiento puede reiniciarse de forma automática sin la intervención del usuario, y sin cargarse y reanudarse desde el último punto de control.

Conjunto de datos

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

Redes de centros de datos (DCN)

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

Programación de pandillas

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

Host

Un host es una computadora física que ejecuta VMs. Un host puede ejecutar como máximo cuatro VM a la vez. Cada VM tiene una TPU dedicada.

Inferencia

Cargar un modelo de aprendizaje automático previamente entrenado en un host y realizar predicciones sobre los datos

Interchip Interconnect (ICI)

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

Porciones múltiples

Dos o más porciones de chip TPU que se pueden comunicar a través de DCN.

Nodo

En el contexto de varias porciones, 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. Este te permite distribuir la carga de procesamiento entre varias TPU.

Recurso en cola (QR)

Representación de los recursos TPU que se usa a fin de poner en cola y administrar una solicitud para un entorno de TPU de una sola porción o de varias partes.

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 se reinicia una VM. Para Multislice, se especifica en la solicitud de creación de recursos en cola. Para obtener más información sobre las secuencias de comandos de inicio de Cloud TPU, consulta Administra recursos TPU.

Porción de TPU

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í mediante la red ICI.

VM de TPU

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

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 el procesamiento más rápido y con mayor eficiencia energética 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 tres tipos de capacidad (consulta Opciones de uso en Cómo funcionan los precios de TPU) :

• Reserva: Segmenta la cuota reservada. Para usar la cuota reservada, debes tener un acuerdo de reserva con Google. Usa la marca --reserved cuando crees tus recursos.
• Interrumpible: Segmenta la cuota interrumpible. Tus recursos pueden ser interrumpidos para hacer lugar a las solicitudes de un trabajo de mayor prioridad. Usa la marca --best-effort cuando crees tus recursos.
• A pedido: Se orienta a 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 cuotas según demanda que ofrece Cloud TPU. Se selecciona de forma predeterminada, no se necesitan marcas.

Nota: Cloud TPU limita el tamaño de los fragmentos de todos los tipos de acelerador. Si quieres conocer los límites de tamaño de los segmentos según demanda, consulta Tamaños predeterminados de los segmentos para los recursos en cola según demanda.

Comenzar

Si no has usado TPU antes, primero instala los componentes de gcloud y habilita la API de Cloud TPU. Luego, configura el comando gcloud para establecer tu proyecto. Sigue todas las instrucciones en Configura un proyecto y detente cuando llegues a la sección Administra las TPU.

Para usar Multislice, los recursos TPU se deben administrar como recursos en cola. Si eres un usuario existente de TPU v4 y tienes una reserva, es posible que debas migrarla a un sistema de reservas nuevo. Para obtener más información, comunícate con el equipo de Cloud TPU. En la etapa de disponibilidad general, la solicitud de creación de recursos en cola podrá acceder a todas las reservas, y los usuarios no necesitarán transferir ninguna capacidad antes de usar recursos múltiples o en cola.

Ejemplo introductorio

En este instructivo, se usa código en el repositorio de GitHub de MaxText. shardings.py (dentro del repositorio de MaxText) está diseñado para facilitar la experimentación con diferentes opciones de paralelismo. Es el mejor lugar para comenzar si deseas explorar el paralelismo de datos, el paralelismo de datos fragmentados por completo (FSDP) o el paralelismo de tensor. Te permite extender con facilidad el código a la configuración de una sola Slice y Multislice.

Esta opción no es específica para Multislice y debería estar en otro lugar.

Fragmentación de una sola sección con paralelismo de ICI

La fragmentación de ICI corresponde a la fragmentación dentro de cada porción. shardings.py proporciona tres entradas 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 de cada método de paralelismo.

Estas entradas deben estar restringidas para que ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism sea igual a la cantidad de chips de la porción.

En la siguiente tabla, se muestran ejemplos de entradas de usuarios para el paralelismo de ICI en los cuatro chips disponibles en la versión v4-8.

	ici_data_paralelismo	ici_fsdp_paralelismo	ici_tensor_paralelismo
FSDP de 4 direcciones	1	4	1
Paralelismo de tensor de 4 vías	1	1	4
FSDP de 2 vías + paralelismo de tensor de 2 vías	1	2	2

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

En este ejemplo, se supone que estás familiarizado con 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. Este ejemplo se adaptó de esa guía para ejecutar shardings.py en una sola Slice.

1. Configura el entorno:

   $ QUEUED_RESOURCE_ID=user-defined-queued-resource-id
   PROJECT=your-gcp-project-name
   ZONE=us-central2-b
   RUNTIME_VERSION=tpu-ubuntu2204-base
   
   gcloud auth login
   gcloud config set project ${PROJECT}
   gcloud config set compute/zone ${ZONE}
   

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

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

3. Aprovisiona las TPU con el siguiente comando:

   $ gcloud alpha compute tpus queued-resources create $QUEUED_RESOURCE_ID \
   --accelerator-type v4-8 \
   --runtime-version tpu-ubuntu2204-base \
   --node-id $QUEUED_RESOURCE_ID \
   [--reserved |--best-effort]
   

   Nota: Por el momento, la CLI de gcloud no admite todas las opciones de creación de recursos en cola, como las etiquetas. Para obtener más información, consulta Crea recursos en cola.

4. Espera hasta que el QR tenga el estado ACTIVE, lo que significa que los nodos trabajadores están en el estado READY. Una vez que comienza el aprovisionamiento del código QR, puede tardar entre uno y cinco minutos en completarse, según el tamaño del código QR. Puedes verificar el estado mediante el siguiente comando:

   $ gcloud alpha compute tpus queued-resources \
   list --filter=$QUEUED_RESOURCE_ID
   

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

   $ gcloud compute tpus tpu-vm ssh $QUEUED_RESOURCE_ID
   

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

7. Dentro del directorio del repositorio de MaxText, ejecuta la secuencia de comandos de configuración para instalar JAX y otras dependencias en tu porción de TPU.

   $ bash setup.sh # These installations will take
   a few minutes.
   

8. Ejecuta el siguiente comando para ejecutar shardings.py en la 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 lograr alrededor de 260 TFLOP/s o una utilización impresionante 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).

9. Puedes explorar otras estrategias de fragmentación a través de 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 código QR 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 y se puede ejecutar en segundo plano con la marca opcional --async.

    Importante: Los códigos QR se deben borrar después de usarlos. De lo contrario, seguirán consumiendo tu cuota de QR, que se registra por separado de la cuota de TPU.

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

Fragmentación de varias partes 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_paralelismo
• dcn_fsdp_paralelismo
• dcn_tensor_paralelismo

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

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

	dcn_data_paralelismo	dcn_fsdp_paralelismo	dcn_tensor_paralelismo	Cant. de porciones
Paralelismo de datos bidireccional	2	1	1	2

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

A medida que aumentas la cantidad de porciones (suponiendo que mantienes el tamaño de las porciones y el lote por porción constantes), aumentas la cantidad de paralelismo de datos. Estos pasos también se enumeran en el archivo README de MaxText.

Ejecuta shardings.py en un entorno de varias partes

Puedes ejecutar shardings.py en un entorno de varias porciones con multihost_runner.py o mediante la ejecución de 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 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 reutilizan repetidamente las mismas TPU. Debido a 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, horas o días), te recomendamos usar multihost_job.py.

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

multihost_runner.py se conecta de forma automática a los trabajadores TPU mediante SSH.

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

Configure su entorno

1. Clona MaxText en la máquina del ejecutor.
2. Ve al directorio del repositorio.
3. Configura el entorno:

   $ export QUEUED_RESOURCE_ID=user-defined-queued-resource-id
   $ export PROJECT=your-gcp-project-name
   $ export ZONE=us-central2-b
   $ export RUNTIME_VERSION=tpu-ubuntu2204-base
   $ gcloud config set project ${PROJECT}
   $ gcloud config set compute/zone ${ZONE}
   

4. Crea llaves 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 solicita que el archivo google_compute_engine ya existe, selecciona no conservar tu versión existente.

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

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

   $ export SLICE_COUNT=2
   

6. Crea un entorno de Multislice con queued-resources create.

   Nota: Asegúrate de tener la cuota correspondiente antes de seleccionar --reserved, --best_effort, o la cuota según demanda predeterminada. Si deseas 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. TPU v5e también es compatible con la versión preliminar pública de Multislice. Para usar v5e, especifica una versión accelerator-type v5e (por ejemplo, v5e-16) y la v5e runtime-version (v2-alpha-tpuv5-lite).

   $ gcloud alpha compute tpus queued-resources \
   create $QUEUED_RESOURCE_ID \
   --accelerator-type=v4-16 \
   --runtime-version=tpu-ubuntu2204-base \
   --node-count=$SLICE_COUNT \
   --node-prefix=$QUEUED_RESOURCE_ID \
   [--reserved|--best-effort]
   

   Descripciones de las marcas de comandos

   QUEUED_RESOURCE_ID

   Una cadena definida por el usuario que identifica la solicitud de recurso en cola.

   accelerator-type

   El tipo de Cloud TPU que se creará.

   runtime-version

   La [versión del software de Cloud TPU](/tpu/docs/supported-tpu-configurations#tpu_software_versions).

   node-count

   La cantidad de porciones que se crearán.

   node-prefix

   Es el prefijo que se usa para generar nombres para cada porción. Se agrega un número al prefijo de cada porción. Por ejemplo, si configuras node-prefix como mySlice, las partes se denominan mySlice-0, mySlice-1, etcétera.

   reserved

   Usa la cuota reservada cuando crees las porciones.

   best-effort

   Usa la cuota del mejor esfuerzo cuando crees las porciones (opción predeterminada).

7. Espera hasta que el recurso en cola (QR) esté en el estado ACTIVE, lo que significa que los nodos trabajadores están en el estado READY. Una vez que se inicia el aprovisionamiento del código QR, puede tardar entre uno y cinco minutos en completarse, según el tamaño del código QR. Puedes verificar el estado con el siguiente comando:

   $ gcloud alpha compute tpus queued-resources list \
   --filter=$QUEUED_RESOURCE_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
   ...
   

   Comunícate con nosotros si el estado del código QR está en el estado ACCEPTED (WAITING_FOR_RESOURCES) o PROVISIONING durante más de 15 minutos.

8. Instala las dependencias.

   $ pip3 install absl-py
   $ python3 multihost_runner.py \
     --TPU_PREFIX=$QUEUED_RESOURCE_ID \
     --COMMAND="bash setup.sh"
   

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

   $ python3 multihost_runner.py \
     --TPU_PREFIX=$QUEUED_RESOURCE_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/s de rendimiento en los archivos de registro.

10. Limpia las TPU y el código QR cuando termines. La eliminación tardará entre dos y cinco minutos en completarse y se puede ejecutar en segundo plano con la marca opcional --async.

MaxText: Un LLM escalable, de alto rendimiento y de varias partes

Para facilitar comenzar, proporcionamos train.py, una implementación de LLM escalable, legible y de alto rendimiento que funciona de inmediato. Te recomendamos que lo ejecutes por tu cuenta y uses el código en tus propias implementaciones. MaxText contiene sus propios resultados de rendimiento y README.

Se incluyen los resultados de rendimiento de MaxText, copiados del archivo README de MaxText. MaxText también contiene todo lo necesario para reproducir los resultados. Aquí suponemos FLOPS en 6 * BATCH * Parameters.

Modelo de 22,000 millones de parámetros

Hardware	TFLOP/s/chip	Uso de FLOP del modelo
1 v4-128	165	60%
2x v4-128	158	57,4%
4x v4-128	157	57,1%

Modelo de 52,000 millones de parámetros

Hardware	TFLOP/s/chip	Uso de FLOP del modelo (MFU)
1 v4-384	154	56%
2x v4-384	162	58,9%

Las mejoras de 1x v4-384 a 2x v4-384 están impulsadas por una particularidad del compilador. Es posible que se muestre en diferentes tamaños, pero, por lo general, notarás que la MFU disminuye a medida que aumentamos la cantidad de divisiones.

Escala una carga de trabajo a Multislice

Cambios de código

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

• Cuando crees la malla, usa jax.experimental.mesh_utils.create_hybrid_device_mesh en lugar de jax.experimental.mesh_utils.create_device_mesh.

• Configura algunas variables de entorno:

  os.environ["JAX_USE_PJRT_C_API_ON_TPU"] = "1"
  os.environ["LIBTPU_INIT_ARGS"] = "--xla_tpu_enable_megascale_barrier=true"
  

  Nota: LIBTPU_INIT_ARGS es una lista separada por espacios, por ejemplo, LIBTPU_INIT_ARGS="--flag1=val1 --flag2=var2".

Estos deberían ser los únicos cambios de código necesarios cuando se pase a Multislice. Para lograr un rendimiento alto, el DCN debe asignarse a ejes paralelos de datos fragmentados por completo o a ejes paralelos 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 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 porciones de la versión 4-16, tienes ocho chips por porción × 4 porciones, por lo que len(jax.devices()) debería mostrar 32.

Elige tamaños de porción para entornos de Multislice

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

Ejecuta el objeto Job en varias partes

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

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

Secuencia de comandos del ejecutor de experimentación

Para una experimentación rápida, proporcionamos multihost_runner.py (alojado en MaxText), una secuencia de comandos simple para distribuir código a un entorno Multislice existente, ejecutar tu comando en cada host, copiar tus registros y hacer un seguimiento del estado de error de cada comando. La secuencia de comandos multihost_runner.py está más detalladamente documentada en el archivo README de MaxText.

Debido a que multihost_runner.py mantiene conexiones SSH persistentes, solo es adecuado para experimentos de tamaño moderado y de ejecución relativamente corta. Adapta 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 contra fallas de hardware y otras interrupciones, es mejor realizar la integración directamente con la API de Create Queued Resource. Como ejemplo de funcionamiento, proporcionamos multihost_job.py (alojado en MaxText), que activa la llamada a la API de Created Queued Resource con la secuencia de comandos de inicio adecuada para ejecutar el entrenamiento y reanudarlo en caso de interrupción. La secuencia de comandos multihost_job.py está documentada con más detalle 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 usar o adaptar multihost_runner.py o multihost_job.py para ejecutar tu carga de trabajo personalizada en la configuración de Multislice que desees. Sin embargo, si prefieres aprovisionar y administrar tu entorno directamente mediante comandos de recursos en cola, consulta Administra un entorno de varias partes.

Administra un entorno de varias partes

Para aprovisionar y administrar manualmente los recursos en cola sin usar las herramientas que se proporcionan en el repositorio de MaxText, lee las siguientes secciones.

Crea recursos en cola

Configura las siguientes variables de entorno antes de aprovisionar capacidad:

  $ export QUEUED_RESOURCE_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

Entrada	Descripción
QUEUED_RESOURCE_ID	El ID asignado por el usuario del recurso en cola.
PROYECTO	Nombre del proyecto de Google Cloud
ZONA	us-central2-b
RED_NAME	Nombre de las redes de VPC.
SUBNETWORK_NAME	Nombre de la subred en las redes de VPC
RUNTIME_VERSION	tpu-ubuntu2204-base
TIPO_ACCELERATOR	Versión 4-16
EXAMPLE_TAG_1, EXAMPLE_TAG_2 ...	Etiquetas que se usan con el fin de identificar objetivos o fuentes válidos para firewalls de red
SLICE_COUNT	Cantidad de porciones. Limitado a un máximo de 32 fragmentos, pero se puede aumentar durante la versión preliminar pública.
STARTUP_SCRIPT	Si se agrega a la solicitud de creación, se puede ejecutar una secuencia de comandos de inicio cada vez que se aprovisiona o reinicia una porción de TPU y si esta se repara o restablece.

Crea una solicitud de recurso en cola con gcloud

$ gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_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|--best-effort]
  

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

Crea una solicitud de recurso en cola con curl

Crea un archivo llamado queuedresourcereq.json y copia el siguiente JSON en él.

{
  "guaranteed": { "reserved": true },
  "tpu": {
    "node_spec": [
      {
        "parent": "projects/<PROJECT>/locations/<ZONE>",
        "node": {
          "accelerator_type": "<ACCELERATOR_TYPE>",
          "runtime_version": "<RUNTIME_VERSION>",
          "network_config": {
            "network": "<NETWORK_NAME>",
            "subnetwork": "<SUBNETWORK_NAME>",
            "enable_external_ips": true
          },
     "tags" : ["<EXAMPLE_TAG_1>","<EXAMPLE_TAG_2>"]
          "metadata": {
            "startup-script": "<STARTUP_SCRIPT>"
          }
        },
        "multi_node_params": {
          "node_count": <SLICE_COUNT>,
          "node_id_prefix": "<QUEUED_RESOURCE_ID>"
        }
      }
    ]
  }
}

Reemplaza los elementos de "< >" por los valores adecuados. Para obtener más información, consulta la documentación de la API de recursos en cola de REST para ver todas las opciones disponibles.

Nota: Si bien node_spec es una lista, en este momento solo se admite una entrada node_spec.

Para usar la capacidad prioritaria, reemplaza lo siguiente:

"guaranteed": { "reserved": true }

con

"best_effort": {}

También puedes quitar la línea para usar la capacidad predeterminada según demanda.

Envía la solicitud de creación de QR con la carga útil de 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/${PROJECT}/locations/${ZONE}/queuedResources\?queued_resource_id\=${QUEUED_RESOURCE_ID}

Deberías ver un resultado similar al siguiente, con el nombre de la operación para la solicitud de creación del código QR:

{
  "name": "projects/123456789012/locations/us-central2-b/operations/operation-0000000000000-0000000000000-00000000-00000000",
  "metadata": {
    "@type": "type.googleapis.com/google.cloud.common.OperationMetadata",
    "createTime": "2023-03-09T07:35:11.142588587Z",
    "target": "projects/123456789012/locations/us-central2-b/queuedResources/your-queued-resource-id",
    "verb": "create",
    "cancelRequested": false,
    "apiVersion": "v2alpha1"
  },
  "done": false
}

Puedes usar la operación name para obtener el estado del comando QR de creación.

Obtener operación

A fin de examinar los detalles de la operación para crear un QR, usa lo siguiente:

  $ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2alpha1/projects/$PROJECT/locations/$ZONE/operations/operation-00000000000000-0000000000000-00000000-00000000

La respuesta de este comando contiene el estado de la operación:

{
  "name": "projects/project/locations/zone/operations/operation-00000000000000-0000000000000-00000000-00000000",
  "metadata": {...},
  "done": true,
  "response": {
    "@type": "type.googleapis.com/google.cloud.tpu.v2alpha1.QueuedResource",
    ...
    "state": {
      "state": "WAITING_FOR_RESOURCES"
    }
  }
}

Si se creó correctamente el QR ("done = true"), el estado dentro del campo response será ACCEPTED (WAITING_FOR_RESOURCES) o FAILED. Si el código QR está en el estado ACCEPTED (WAITING_FOR_RESOURCES), significa que se puso en cola y comenzará a aprovisionarse cuando haya suficientes recursos. Si el QR está en estado FAILED, el motivo de la falla aparecerá en el resultado. Para obtener más información sobre otros estados posibles, consulta la Guía del usuario de recursos en cola.

Una vez que finalice la operación, usa los recursos en cola de descripción para supervisar las etapas del código QR.

En una situación poco frecuente, es posible que encuentres el QR en estado FAILED mientras que algunas partes tienen el estado ACTIVE. Si esto sucede, borra los recursos que se crearon y vuelve a intentarlo en unos minutos o comunícate con el equipo de Cloud TPU para resolver el problema.

Instala dependencias y SSH

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

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

Este comando gcloud envía el comando especificado a todos los trabajadores y nodos en QR mediante SSH. El comando se agrupa en grupos de cuatro y se envía en simultáneo. El siguiente lote de comandos se envía cuando el lote actual completa la ejecución. Si se produce un error en uno de los comandos, el procesamiento se detiene y no se envían más lotes. Para obtener más información, consulta la referencia de la API de recursos en cola.

Configuración de red

Ejecuta los pasos que se indican a continuación para asegurarte de que las porciones de TPU se puedan comunicar entre sí.

Instala JAX en cada una de las porciones. 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 Multislice. Para ello, ejecuta lo siguiente en cada porción:

export JAX_USE_PJRT_C_API_ON_TPU=1 && python3 -c 'import jax; print(jax.devices())'

Si ejecutas este código en cuatro partes de la versión 4-16, habrá ocho chips por porción y cuatro porciones, jax.devices() debe mostrar un total de 32 chips (dispositivos).

Obtén una lista de los recursos en cola

Puedes ver el estado de los códigos QR con el comando queued-resources list:

$ gcloud alpha compute tpus queued-resources list

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

Describir recursos en cola

Para ver la configuración detallada y el estado de un código QR, usa la API de QR de descripción. Puedes llamar a esta API con gcloud o curl.

Usa gcloud de la siguiente manera:

$ gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}
...state:
 state: ACTIVE
...

Usa curl de la siguiente manera:

$ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2alpha1/projects/${PROJECT}/locations/${ZONE}/queuedResources/${QUEUED_RESOURCE_ID}
{
  "name": your-queued-res,
  "tpu": {
    "nodeSpec": [
      {
        ... // node 1
      },
      {
        ... // node 2
      },
      ...
    ]
  },
  ...
  "state": "ACTIVE"
  }
}

"state" representa el estado de un código QR. Los estados se definen de la siguiente manera:

ACCEPTED (obsoleto, se reemplazó por WAITING_FOR_RESOURCES)	Se recibió el comando QR create y comenzará el aprovisionamiento en cuanto los recursos estén disponibles.
PROVISIONING	Se están aprovisionando los fragmentos de TPU.
ACTIVE	Todas las TPU se aprovisionan y están listas para usarse. Si se proporciona una secuencia de comandos de inicio, esta comenzará a ejecutarse en todas las TPU, cuando el estado QR pase a ACTIVE. Ten en cuenta que, en este momento, hay un pequeño error en el que el estado QR sigue mostrándose como ACTIVE, mientras se borran algunas secciones. En este caso, el código QR no funcionará como se espera y se deben limpiar todos los recursos.
FAILED	No se pudieron aprovisionar los segmentos.
SUSPENDING	Se están borrando uno o más segmentos.
SUSPENDED	Se borran todas las secciones subyacentes, pero el código QR permanece intacto hasta que se borra explícitamente. En este momento, un código QR suspendido no se puede reanudar y se debe borrar.
DELETING	Se está borrando el código QR.

Inicia tu trabajo en un entorno aprovisionado

Para ejecutar cargas de trabajo de forma manual, conéctate a todos los hosts en cada fragmento a través de SSH y ejecuta el siguiente comando en todos los hosts. En general, debes ejecutar un comando como el siguiente una vez para cada porción.

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

Restablece recursos en cola

La API de ResetQueuedResource se puede usar para restablecer todas las VM en un código 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. Todos los datos almacenados de forma local permanecerán intactos y la secuencia de comandos de inicio se invocará después de un restablecimiento. La API de ResetQueuedResource puede ser útil cuando deseas reiniciar todas las TPU. Por ejemplo, cuando el entrenamiento está atascado y el restablecimiento de todas las VMs es más fácil que la depuración.

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

$ gcloud alpha compute tpus queued-resources reset $QUEUED_RESOURCE_ID

Borra recursos en cola

Para liberar recursos al final de tu sesión de entrenamiento, borra el recurso en cola con la marca --force. La eliminación tardará entre dos y cinco minutos en completarse y se puede ejecutar en segundo plano con la marca opcional --async.

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

Recuperación automática de fallas

En el caso de una interrupción, Multislice ofrece reparación sin intervención de la parte afectada y restablecimiento de todas las secciones posteriormente. La porción afectada se reemplaza por una nueva, y las partes restantes, que en otro estado serían en buen estado, se restablecerán. Si no hay capacidad disponible para asignar una porción de reemplazo, el entrenamiento se detiene.

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 de forma automática cada vez que se reasigna una parte o se restablece una VM. Debes especificar una secuencia de comandos de inicio en la carga útil de JSON que envías a la API de creación de solicitud de QR.

La siguiente secuencia de comandos de inicio (que se usa en Crea recursos en cola) te permite recuperarte de las fallas de forma automática 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 probarlo.

Generación de perfiles y depuración

La creación de perfiles es la misma en entornos de una sola Slice y Multislice. Para obtener más información, consulta Genera perfiles de programas JAX.

Optimiza el entrenamiento

Fragmenta con Multislice para obtener el máximo rendimiento

Para lograr el máximo rendimiento en entornos de varias porciones, es necesario considerar cómo fragmentar las varias porciones. Por lo general, hay tres opciones (paralelismo de datos, paralelismo de datos fragmentados por completo y paralelismo de canalización). No recomendamos fragmentar activaciones en las dimensiones del modelo (a veces llamadas paralelismo de tensor) porque requiere demasiado ancho de banda entre las secciones. Para todas estas estrategias, puedes mantener la misma estrategia de fragmentación dentro de una porción que te haya funcionado en el pasado.

Recomendamos comenzar con un paralelismo de datos puros. El uso de paralelismo de datos fragmentados por completo es útil para liberar el uso de memoria. La desventaja es que la comunicación entre segmentos usa la red DCN y ralentiza la carga de trabajo. Usa el paralelismo de la canalización solo cuando sea necesario en función del tamaño del lote (como se analiza a continuación).

Cuándo usar el paralelismo de datos

El paralelismo puro de datos funcionará bien en los casos en los que tengas una carga de trabajo que se ejecute correctamente, pero desees mejorar su rendimiento mediante el escalamiento en varias porciones.

A fin de lograr un escalamiento sólido en varias porciones, la cantidad de tiempo necesaria para realizar la reducción completa a través de DCN debe ser menor que la cantidad de tiempo necesaria para realizar un pase hacia atrás. DCN se usa para la comunicación entre fragmentos y es un factor limitante en la capacidad de procesamiento de la carga de trabajo.

Cada chip de TPU v4 tiene un rendimiento 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. Eso significa que la intensidad aritmética es 4 * 275 * 10¹² FLOPS / 50 Gbps = 22,000 FLOPS / bit.

Tu modelo usará de 32 a 64 bits de ancho de banda DCN para cada parámetro y por paso. Si usas dos porciones, tu modelo usará 32 bits de ancho de banda de DCN. Si usas más de dos segmentos, el compilador realizará una operación de redistribución completa de Shuffle y usarás hasta 64 bits de ancho de banda de DCN para cada parámetro por paso.

La cantidad de FLOPS que se necesita 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 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 de 6 * B y la cantidad de FLOPS por parámetro durante la pasada es de 4 * B.

Para garantizar un escalamiento sólido en varias partes, debes asegurarte 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 la transmisión 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 porciones:

Operational intensity = 4 * B  / 32

O bien, si usas más de dos porciones:

Operational intensity = 4 * B  / 64

Esto sugiere un tamaño mínimo de lote 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 significativo de error, implementando 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 del modelo, deberás estimar el entorno de ejecución de tu pase hacia atrás por porción (ya sea mediante el tiempo con un generador de perfiles o mediante el recuento de FLOPS). Luego, puedes comparar eso con el tiempo de ejecución esperado para reducir todo a través de DCN y obtener una buena estimación de si el paralelismo de datos tendrá sentido en tu caso.

Cuándo usar el paralelismo de datos completamente fragmentados (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 los nodos. Para cada operación en los pases hacia adelante y hacia atrás, los pesos se recopilan para que cada porción tenga los pesos que necesita. En lugar de sincronizar los gradientes con All-Reduce, los gradientes se reducen a medida que se producen. De esta manera, cada porción solo obtiene los gradientes de los pesos de los que es responsable.

Al igual que en el paralelismo de datos, el FSDP requerirá escalar el tamaño del lote global de forma lineal con la cantidad de porciones. La FSDP disminuirá la presión de la memoria a medida que aumentes la cantidad de divisiones. Esto se debe a que la cantidad de ponderaciones y estados del optimizador por porción disminuye, pero lo hace al precio de un aumento del tráfico de red y a la mayor posibilidad de bloqueo debido a un colectivo retrasado.

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

Las operaciones de recopilación y reducción total de 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 DCN de la misma manera que se describió en la sección anterior.

Cuándo usar el paralelismo de canalizaciones

El paralelismo de la canalización se vuelve relevante cuando se alcanza un alto rendimiento con otras estrategias de paralelismo que requieren un tamaño de lote global mayor que el tamaño máximo de lote preferido. El paralelismo de canalizaciones permite que las porciones que comprenden una canalización “compartan” un lote. Sin embargo, el paralelismo de la canalización tiene dos desventajas importantes:

1. Se genera una “burbuja de canalización”, en la que los chips están inactivos porque están esperando datos.
2. Requiere microlotes, que disminuyen el tamaño efectivo del lote, la intensidad aritmética y, en última instancia, el uso de FLOP del modelo.

El paralelismo de la canalización debe usarse solo si las otras estrategias de paralelismo requieren un tamaño de lote global demasiado grande. Antes de probar el paralelismo de la 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 una FSDP de alto rendimiento. El FSDP suele lograr una utilización más alta de FLOP en el modelo, pero si la convergencia por muestra disminuye a medida que aumenta el tamaño del lote, el paralelismo de la 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 la canalización, pero tu carga de trabajo puede ser diferente.

Si es necesario el paralelismo de la canalización, recomendamos combinarlo con el paralelismo de datos o FSDP. Esto te permitirá minimizar la profundidad de la canalización y aumentar el tamaño del lote por canalización hasta que la latencia de DCN se convierta en un factor menor en la capacidad de procesamiento. En concreto, si tienes N porciones, considera las canalizaciones de profundidad 2 y N/2 réplicas de paralelismo de datos, luego las canalizaciones de profundidad 4 y las réplicas de N/4 de paralelismo de datos, y así sucesivamente, hasta que el lote por canalización sea lo suficientemente grande como para que los colectivos DCN puedan ocultarse detrás de la aritmética en el retroceso. Esto minimizará la demora que presente el paralelismo de la canalización, a la vez que te permitirá escalar más allá del límite global del tamaño del lote.

Prácticas recomendadas de varias partes

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 fragmente el lote 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 texto, pero requiere una refragmentación 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 interrupción.

Controles

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

Optimiza el entrenamiento

Fragmenta con Multislice para obtener el máximo rendimiento

Para lograr el máximo rendimiento en entornos de varias porciones, es necesario considerar cómo fragmentar las varias porciones. Por lo general, hay tres opciones (paralelismo de datos, paralelismo de datos fragmentados por completo y paralelismo de canalización). No recomendamos fragmentar activaciones en las dimensiones del modelo (a veces llamadas paralelismo de tensor) porque requiere demasiado ancho de banda entre las secciones. Para todas estas estrategias, puedes mantener la misma estrategia de fragmentación dentro de una porción que te haya funcionado en el pasado.

Recomendamos comenzar con un paralelismo de datos puros. El uso de paralelismo de datos fragmentados por completo es útil para liberar el uso de memoria. La desventaja es que la comunicación entre segmentos usa la red DCN y ralentizará la carga de trabajo. Usa el paralelismo de la canalización solo cuando sea necesario en función del tamaño del lote (como se analiza a continuación).

¿Cuándo usar el paralelismo de datos?

El paralelismo puro de datos funcionará bien en los casos en los que tengas una carga de trabajo que se ejecute correctamente, pero desees mejorar su rendimiento mediante el escalamiento en varias porciones.

A fin de lograr un escalamiento sólido en varias porciones, la cantidad de tiempo necesaria para realizar la reducción completa a través de DCN debe ser menor que la cantidad de tiempo necesaria para realizar un pase hacia atrás. DCN se usa para la comunicación entre fragmentos y es un factor limitante en la capacidad de procesamiento de la carga de trabajo.

Cada chip de TPU v4 tiene un rendimiento 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. Eso significa que la intensidad aritmética es 4 * 275 * 10¹² FLOPS / 50 Gbps = 22,000 FLOPS / bit.

Tu modelo usará de 32 a 64 bits de ancho de banda DCN para cada parámetro y por paso. Si usas dos porciones, tu modelo usará 32 bits de ancho de banda de DCN. Si usas más de dos segmentos, el compilador realizará una operación de redistribución completa de Shuffle y usarás hasta 64 bits de ancho de banda de DCN para cada parámetro por paso.

La cantidad de FLOPS que se necesita 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 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 de 6 * B y la cantidad de FLOPS por parámetro durante la pasada es de 4 * B.

Para garantizar un escalamiento sólido en varias partes, debes asegurarte de que la intensidad operativa supere la intensidad aritmética del hardware de la TPU. Calculamos la intensidad aritmética al comienzo de esta sección. Para calcular la intensidad operativa, dividimos la cantidad de FLOPS por parámetro durante la transmisión 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 usamos dos porciones:

Operational intensity = 4 * B  / 32

O, si usamos más de dos porciones:

Operational intensity = 4 * B  / 64

Esto sugiere un tamaño mínimo de lote 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 significativo de error y 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 del modelo, deberás estimar el entorno de ejecución de tu pase hacia atrás por porción (ya sea mediante el tiempo con un generador de perfiles o mediante el recuento de FLOPS). Luego, puedes comparar eso con el tiempo de ejecución esperado para reducir todo a través de DCN y obtener una buena estimación de si el paralelismo de datos tendrá sentido en tu caso.

Parámetros de configuración compatibles

Formas

Todos los segmentos deben tener la misma forma (por ejemplo, el mismo AcceleratorType). Por el momento, no se admiten formas heterogéneas de porciones.

Generaciones de TPU

En la actualidad, Multislice solo es compatible con los pods v4.

Organización

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

Frameworks

La vista previa de Multislice en v4 solo admite cargas de trabajo JAX.

Paralelismo

Recomendamos a los usuarios que prueben Multislice con paralelismo de datos en esta versión preliminar. Comunícate con nosotros si deseas obtener más información sobre cómo implementar el paralelismo de canalizaciones con Multislice.

Asistencia y comentarios

Los comentarios son bienvenidos. Para compartir comentarios o solicitar asistencia, comunícate con nosotros aquí o envía un correo electrónico a [e-mail].

Condiciones

Toda la información que Google te proporcionó con respecto a esta vista previa privada es información confidencial de Google y está sujeta a las disposiciones de confidencialidad de las Condiciones del Servicio de Google Cloud Platform (o cualquier otro acuerdo que rija tu uso de Google Cloud).