Descripción general de Cloud TPU Multislice

Cloud TPU Multislice es una tecnología de escalamiento de rendimiento de pila completa que permite que un trabajo 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 los trabajos de entrenamiento que requieren menos de 4,096 chips, una sola porción puede ofrecer el mejor rendimiento. Sin embargo, varios segmentos más pequeños tienen más disponibilidad, 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 varias porciones, los chips TPU en cada porción se comunican a través de una interconexión entre chips (ICI). Los chips TPU en diferentes porciones se comunican mediante la transferencia de 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 DCN entre porciones. El compilador de XLA genera ese código por ti y superpone la comunicación con el procesamiento para obtener el máximo rendimiento.

Conceptos

Tipo de acelerador

Es la forma de cada porción de TPU que comprende una Multislice. Cada porción en una solicitud de varias secciones es del mismo tipo de acelerador. Un tipo de acelerador consta de un tipo de TPU (v4 o v5e) seguido de la cantidad de TensorCores. Por ejemplo, v4-128 especifica una TPU v4 con 128 núcleos de tensor.

Reparación automática

Cuando una porción se encuentra con un evento de mantenimiento, una interrupción o una falla de hardware, Cloud TPU creará una nueva porción. En el caso poco frecuente de que no haya recursos suficientes 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 en el entorno de Multislice para que el entrenamiento pueda continuar.Con una secuencia de comandos de inicio configurada de manera correcta, la secuencia de comandos de entrenamiento puede reiniciarse automáticamente sin intervención del usuario y cargar y reanudar desde el último punto de control.

Conjunto de datos

Los datos que usa un modelo para entrenamiento o inferencia.

Herramientas de redes de centros de datos (DCN)

Una red de latencia más alta 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, se garantiza 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 VMs a la vez. Cada VM tiene una TPU dedicada.

Inferencia

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

Interchip Interconnect (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 chip de TPU que se pueden comunicar a través de DCN.

Node

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

Pod

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

Recurso en cola (QR)

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 sola porción o de Multislice.

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. Para Multislice, se especifica en la solicitud de creación del código QR. Para obtener más información sobre las secuencias de comandos de inicio de Cloud TPU, consulta Administra 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 que tiene acceso a las TPU subyacentes. Para 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ñadas con el objetivo de ofrecer un procesamiento rápido y de bajo consumo energético 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 Opciones de uso en Cómo funcionan los precios de TPU) :

• Reserva: se orienta a la cuota reservada. Para usar las cuotas reservadas, 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. Tus recursos pueden interrumpirse a fin de dejar lugar a solicitudes para un trabajo de mayor prioridad. 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 que no se interrumpirá. La solicitud de TPU se pondrá en cola a una cola de cuota según demanda que ofrece Cloud TPU; no se garantiza la disponibilidad de los recursos. Seleccionada de forma predeterminada, no se necesitan marcas.

Nota: Cloud TPU limita el tamaño de las porciones para todos los tipos de acelerador. Si deseas conocer los límites del tamaño de las porciones para las solicitudes a pedido, consulta Tamaños predeterminados de las porciones para los QR a pedido.

Comenzar

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

Si eres un usuario de TPU v4 existente y tienes una reserva, es posible que debas migrarla a un sistema de reservas 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, paralelismo de datos, paralelismo de datos completamente fragmentados (FSDP) y paralelismo de tensores. El código escala de entornos de una sola porción a entornos de varias secciones.

Paralelismo de ICI

ICI se refiere a la interconexión de alta velocidad que conecta las TPU en una sola porción. La 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 especificas para estos parámetros determinan la cantidad de fragmentos para 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 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 la versión 4-8:

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

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

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. 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 dejar 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, reemplaza la versión existente.

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

3. Aprovisiona tus TPU con el siguiente comando:

   $ gcloud alpha 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

   Una cadena definida por el usuario que identifica la solicitud de QR.

   accelerator-type

   El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Si deseas obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta las versiones de TPU.

   runtime-version

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

   node-id

   Es el ID de los recursos TPU que se crearán en respuesta a la solicitud de QR.

   reserved

   Usar la cuota reservada cuando crees las porciones.

   best-effort

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

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

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 se inicia el aprovisionamiento del QR, puede tardar de uno a cinco minutos en completarse según el tamaño del código QR. Puedes verificar el estado de una solicitud de QR con el siguiente comando:

   $ gcloud compute tpus queued-resources \
     list --filter=your-qr-id
   

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) a tu 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. 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 un impresionante uso de FLOP superior al 90%. En este caso, seleccionamos aproximadamente el lote máximo que se ajusta a la memoria de alto ancho de banda (HBM) de la TPU.

9. Siéntete libre de 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 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 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 deben borrarse después de usarlos. De lo contrario, seguirán consumiendo la cuota de QR, que se registra por separado de la cuota de TPU.

    $ gcloud compute tpus queued-resources
      delete your-qr-id --force (--async)
    

Fragmentación multislice mediante el paralelismo de DCN

La secuencia de comandos shardings.py toma tres parámetros que especifican el paralelismo de DCN, que corresponde al número 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 limitarse para que dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism sea igual al número de porciones.

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

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

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

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

Ejecuta shardings.py en un entorno de varias secciones

Puedes ejecutar shardings.py en un entorno de Multislice con multihost_runner.py o ejecutando shardings.py en cada VM de TPU. Aquí usamos multihost_runner.py. Los siguientes pasos son muy similares a los de Getting Started: Quick Experiments on Multiple slices del repositorio de MaxText, excepto que, en este caso, ejecutamos shardings.py en lugar del LLM más complejo en train.py.

La herramienta multihost_runner.py está optimizada para experimentos rápidos, en los que se reutilizan las mismas TPU de forma repetida. 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 un trabajo más largo (por ejemplo, horas o días), te recomendamos que uses 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 en cualquier VM de Compute Engine en el mismo proyecto que tus porciones. No se admite la ejecución de multihost_runner.py en un trabajador.

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

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

Configura tu entorno.

1. Clona MaxText en tu máquina del ejecutor.

2. Ve al directorio del repositorio.

3. Crea claves SSH para gcloud. Te recomendamos dejar 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
     

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

     $ export SLICE_COUNT=2
     

5. Crear un entorno de varias porciones con queued-resources create

   Nota: Asegúrate de tener la cuota correspondiente antes de seleccionar --reserved, --best_effort o la cuota a pedido predeterminada. Si deseas obtener más información sobre los tipos de cuotas, consulta Política de cuotas.

   En el siguiente comando, se muestra cómo crear una Multislice TPU v4. Para usar v5e, especifica un accelerator-type de v5e (por ejemplo, v5litepod-16) y un runtime-version de v5e (v2-alpha-tpuv5-lite).

     $ gcloud alpha 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

   Una cadena definida por el usuario que identifica la solicitud de QR.

   accelerator-type

   El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. Si deseas obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta las versiones de TPU.

   runtime-version

   La versión de software de Cloud TPU.

   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, los fragmentos se denominan mySlice-0, mySlice-1, etc.

   reserved

   Usar la cuota reservada cuando crees las porciones.

   best-effort

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

6. Cuando se inicia el aprovisionamiento del QR, es posible que tarde hasta cinco minutos en completarse, según su tamaño. Espera hasta que el recurso en cola (QR) esté en el estado ACTIVE. Puedes verificar el estado de una solicitud de QR con el siguiente comando:

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

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

7. Instala las dependencias.

   $ python3 multihost_runner.py \
     --TPU_PREFIX=your-qr-id \
     --COMMAND="bash setup.sh"
   

8. Ejecuta shardings.py en cada trabajador mediante 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 alrededor de 230 TFLOP por segundo de rendimiento en los archivos de registro.

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

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 la malla.

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

Para validar que tu código pueda acceder a todos los dispositivos, puedes confirmar 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.

Elige tamaños de porciones para entornos de varias porciones

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

Cómo ejecutar un trabajo en varias secciones

Existen tres enfoques diferentes para ejecutar la carga de trabajo personalizada en un entorno de varias secciones:

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. Usar un enfoque manual

Secuencia de comandos del ejecutor de experimentos

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

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

Secuencia de comandos del ejecutor de producción

Para los trabajos de producción que necesitan resiliencia frente a fallas de hardware y otras interrupciones, es mejor integrarlos directamente con la API de Create Queued Resource. Como ejemplo funcional, 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 el entrenamiento y reanudarlo en la interrupción. La secuencia de comandos multihost_job.py está documentada 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 la carga de trabajo personalizada en la configuración de Multislice. Sin embargo, si prefieres aprovisionar y administrar tu entorno directamente con comandos QR, consulta Administra un entorno multislice.

Administra un entorno de varias porciones

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

Crear QR

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

Entrada	Descripción
your-qr-id	El ID del código QR asignado por el usuario.
PROYECTO	Nombre del proyecto de Google Cloud
ZONA	us-central2-b
NETWORK_NAME	Nombre de las redes de VPC.
SUBNETWORK_NAME	Nombre de la subred en las redes de VPC
RUNTIME_VERSION	tpu-ubuntu2204-base
ACCELERATOR_TYPE	v4-16
EXAMPLE_TAG_1, EXAMPLE_TAG_2...	Etiquetas usadas para identificar objetivos o fuentes válidos para firewalls de red
SLICE_COUNT	Cantidad de porciones. Tiene un límite máximo de 256 porciones.
STARTUP_SCRIPT	Si se la 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 la porción de TPU se repara o restablece.

Crea una solicitud QR con gcloud

$ gcloud alpha 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

Una cadena definida por el usuario que identifica la solicitud de QR.

project

Una cadena definida por el usuario que identifica la solicitud de QR.

zone

La zona de Google Cloud en la que se crea el QR.

node-count

La cantidad de porciones 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. Si deseas obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta las 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 conectará el recurso de TPU.

reserved

Usar la cuota reservada cuando crees las porciones.

spot

Usar la cuota de VMs Spot cuando se creen los segmentos

Asegúrate de tener la cuota correspondiente antes de seleccionar --reserved, --spot o la cuota a pedido predeterminada. Si deseas obtener más información sobre los tipos de cuotas, consulta Política de cuotas.

Crea una solicitud QR con curl

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"
      }
    }
    ]
  }
}

Nota: Reemplaza los siguientes valores:

• your-project-number: Es el número de proyecto de Google Cloud.
• your-zone: Es la zona en la que deseas crear el QR.
• accelerator-type: Es la versión y el tamaño de una sola porción.
• tpu-vm-runtime-version: las versiones del entorno de ejecución de la VM de TPU
• your-network-name: Opcional, una red a la que se conectará el QR
• your-subnetwork-name: Es una subred a la que se conectará el QR (opcional).
• example-tag-1: Opcional. Una string de etiqueta arbitraria.
• your-startup-script: Es una secuencia de comandos de inicio que se ejecutará cuando se asigne el QR.
• slice-count: Es la cantidad de porciones de TPU en el entorno de Multislice.
• your-qr-id: Es el ID que proporcionó el usuario para el código QR.

Si deseas obtener más información, consulta la documentación de la API de recursos en cola de REST para 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 según demanda predeterminada.

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/your-project-id/locations/your-zone/queuedResources\?queued_resource_id\=your-qr-id

Nota: Reemplaza los siguientes valores:

• your-project-id: Es el ID del proyecto de Google Cloud.
• your-zone: Es la zona en la que deseas crear el QR.
• your-qr-id: Es el ID que proporcionó el usuario para el código QR.

La respuesta debería verse de la siguiente manera:

{
  "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 QR.

Cómo recuperar el estado de un código QR

Para obtener el estado de la solicitud de QR, usa el siguiente comando:

  $ 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/operations/operation-your-qr-guid

Nota: En el comando anterior, reemplaza lo siguiente:

• your-project-id: Es el ID del proyecto de Google Cloud.
• your-zone: Es la zona en la que se crea el QR.
• your-qr-guid: Es el GUID que sigue a name en el resultado de la solicitud de creación del QR.

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

{
  "name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qa-guid>,
  "metadata": {...},
  "done": true,
  "response": {
    "@type": "type.googleapis.com/google.cloud.tpu.v2.QueuedResource",
    ...
    "state": {
      "state": "WAITING_FOR_RESOURCES"
    }
  }
}

Si el código QR se crea correctamente ("done = true"), el estado dentro del campo response será WAITING_FOR_RESOURCES o FAILED. Si el código QR está en el estado WAITING_FOR_RESOURCES, significa que se agregó a la cola y que comenzará a aprovisionarse cuando haya suficientes recursos. Si el código QR está en estado FAILED, el motivo de la falla estará 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 QR de descripción para supervisar las etapas del código QR.

En una situación poco común, es posible que tu QR tenga el estado FAILED, mientras que algunas secciones son 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 código JAX en porciones de pod de TPU, se describe cómo conectarte a tus VM de TPU con SSH en una porción única. Para conectarte a todas las VM de TPU en tu entorno de Multislice mediante SSH y, además, instalar dependencias, usa el siguiente comando de 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 mediante SSH. El comando se agrupa en lotes en grupos de cuatro y se envía de forma simultánea. El siguiente lote de comandos se envía cuando el lote actual completa la ejecución. Si hay una falla con 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. Si la cantidad de segmentos que usas supera el límite de subprocesos de tu computadora local (también llamado límite de lotes), te encontrarás con un interbloqueo. Por ejemplo, supongamos 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 porciones, digamos 100, el comando SSH las dividirá en lotes. Ejecutará la secuencia de comandos de entrenamiento en el primer lote de 64 porciones y esperará a que se completen antes de ejecutarla en el lote restante de 36 porciones. Sin embargo, el primer lote de 64 fragmentos no se puede completar hasta que los 36 fragmentos 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 de unión (&) al comando de 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 porciones, el control volverá de inmediato al comando SSH. Luego, el comando SSH puede comenzar a ejecutar la secuencia de comandos de entrenamiento en el lote restante de 36 porciones. Deberás canalizar tus transmisiones de stdout y stderr de forma adecuada cuando ejecutes los comandos en segundo plano. Para aumentar el paralelismo dentro del mismo QR, puedes seleccionar porciones específicas con el parámetro --node.

Configuración de red

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()) es igual a la cantidad de chips en tu entorno de Multislice. Para hacerlo, en cada porción, ejecuta lo siguiente:

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

Si ejecutas este código en cuatro porciones de las versiones 4-16, hay ocho chips por porción y cuatro porciones, jax.devices() debería mostrar un total de 32 chips (dispositivos).

Enumerar QR

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

$ gcloud compute tpus queued-resources list

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

Describe los QR

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

Usa gcloud de la siguiente manera:

$ gcloud compute tpus queued-resources describe ${your-qr-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/v2/projects/your-project-id/locations/your-zone/queuedResources/${your-qr-id}
{
  "name": your-queued-res,
  "tpu": {
    "nodeSpec": [
      {
        ... // node 1
      },
      {
        ... // node 2
      },
      ...
    ]
  },
  ...
  "state": "ACTIVE"
}

state representa el estado de un QR. Para obtener más información sobre los estados posibles de los QR, consulta Recursos en cola.

Inicia tu trabajo en un entorno aprovisionado

Puedes ejecutar cargas de trabajo de forma manual si te conectas a todos los hosts de cada porción a través de SSH y si ejecutas 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"

Restableciendo QR

La API de ResetQueuedResource se puede usar para restablecer todas las VM en un QR de ACTIVE. Cuando se restablecen las VM, 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 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 se detiene y restablecer todas las VM es más fácil que depurar.

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

$ gcloud compute tpus queued-resources reset your-qr-id

Borrando QR

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 compute tpus queued-resources \
delete your-qr-id --force (--async)

Recuperación automática de fallas

En caso de una interrupción, Multislice ofrece una reparación sin intervención de la porción afectada y el restablecimiento de todas las porciones posteriormente. La porción afectada se reemplaza por una nueva, y las otras en buen estado se restablecen. 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 busque y cargue los últimos puntos de control guardados. La secuencia de comandos de inicio se ejecuta de forma automática cada vez que se reasigna una porción o se restablece una VM. Debes especificar una secuencia de comandos de inicio en la carga útil de JSON que envías a la API de creación de solicitudes de QR.

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

{
 "tpu": {
   "node_spec": [
     {
      ...
         "metadata": {
               "startup-script": "#! /bin/bash \n pwd \n runuser -l user1 -c 'cd /home/user1/MaxText && python3 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 creación de perfiles es la misma en entornos de una sola porción y de varias porciones. Para obtener más información, consulta Cómo crear perfiles de programas de JAX.

Optimiza las capacitaciones

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

Lograr el máximo rendimiento en entornos de varias porciones requiere considerar cómo fragmentar en las múltiples porciones. Por lo general, hay tres opciones (paralelismo de datos, paralelismo de datos completamente fragmentados y paralelismo de canalizaciones). No recomendamos fragmentar activaciones en las dimensiones del modelo (a veces llamado paralelismo de tensor), ya que se 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 el paralelismo de datos puros. El uso del paralelismo de datos fragmentado por completo es útil para liberar el uso de memoria. La desventaja es que la comunicación entre porciones 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 de datos puros funcionará bien en los casos en los que tengas una carga de trabajo que se ejecute bien, pero te gustaría mejorar su rendimiento mediante el escalamiento en varias secciones.

Para lograr un escalamiento sólido en varias porciones, la cantidad de tiempo necesaria para realizar la reducción completa en 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 porciones y es un factor limitante en la capacidad de procesamiento de la carga de trabajo.

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

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

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

Tu modelo usará entre 32 y 64 bits de ancho de banda de DCN para cada parámetro 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 Allreduce de Shuffle y utilizará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 el modelo. Específicamente, para los modelos de lenguaje basados en Transformer, la cantidad de FLOPS necesarios para avanzar y retroceder 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 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 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

Si usas más de dos segmentos: 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, implementar el paralelismo de datos solo si el tamaño del lote por Pod es de al menos 350,000 (dos Pods) a 700,000 (muchos Pods).

Para otras arquitecturas de modelos, deberás estimar el tiempo de ejecución de tu pase hacia atrás por porción (ya sea cronometrando con un generador de perfiles o contando FLOPS). Luego, puedes comparar eso con el tiempo de ejecución esperado a fin de reducir todo con DCN y obtener una buena estimación de si el paralelismo de datos es adecuado para ti.

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 datos entre nodos) con la fragmentación de las ponderaciones entre los nodos. Para cada operación de avance y retroceso, los pesos se juntan, de modo que cada porción tenga los pesos que necesita. En lugar de sincronizar los gradientes con Allreduce, 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.

De manera similar al paralelismo de datos, FSDP requerirá el escalamiento lineal del tamaño del lote global con la cantidad de porciones. FSDP disminuirá la presión de memoria a medida que aumentes la cantidad de porciones. Esto se debe a que la cantidad de pesos y el estado del optimizador por porción disminuye, pero lo hace a precio de un mayor tráfico de red y una mayor posibilidad de bloqueo por un colectivo retrasado.

En la práctica, la FSDP entre porciones es mejor si se aumenta el lote por porción y se almacenan más activaciones para minimizar la rematerialización durante el pase hacia atrás o si se aumenta la cantidad de parámetros en la red neuronal.

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

Cuándo usar el paralelismo de canalizaciones

El paralelismo de canalizaciones 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 canalizaciones permite que las partes de una canalización “compartan” un lote. Sin embargo, el paralelismo de la canalización tiene dos desventajas importantes:

1. Genera el “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 efectivo del lote, la intensidad aritmética y, en última instancia, el uso del modelo de FLOP.

El paralelismo de canalizaciones solo debe usarse si las otras estrategias de paralelismo requieren un tamaño de lote global demasiado grande. Antes de probar el paralelismo de canalizaciones, 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 un mayor uso de FLOP del modelo, pero si la convergencia por muestra se ralentiza a medida que aumenta el tamaño del lote, el paralelismo de la canalización aún puede ser la mejor opción. La mayoría de las cargas de trabajo pueden tolerar tamaños de lote lo suficientemente grandes como para no beneficiarse del paralelismo de la canalización, pero tu carga de trabajo puede ser diferente.

Si el paralelismo de la canalización es necesario, te recomendamos combinarlo con el paralelismo de datos o FSDP. Esto te permitirá minimizar la profundidad de la canalización mientras aumentas el tamaño del lote por canalización hasta que la latencia de DCN se vuelva menos importante en la capacidad de procesamiento. De manera concreta, si tienes N porciones, considera canalizaciones de profundidad 2 y N/2 réplicas del paralelismo de datos, luego canalizaciones de profundidad 4 y N/4 réplicas del paralelismo de datos, y así sucesivamente, hasta que el lote por canalización sea lo suficientemente grande como para que los colectivos de DCN puedan ocultarse detrás de la aritmética en el pase hacia atrás. Esto minimizará la demora que ingresa el paralelismo de la canalización, a la vez que te permitirá escalar más allá del límite de tamaño del lote global.

Prácticas recomendadas de multislice

Carga de datos

Durante el entrenamiento, cargamos lotes de forma repetida desde un conjunto de datos para alimentar el modelo. Tener un cargador de datos asíncrono y eficiente que fragmenta el lote en los hosts es importante para evitar que las TPU no tengan 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 al iterador de datos cargar los mismos datos antes y después de la interrupción.

Controles

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

Configuraciones admitidas

Formas

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

Organización

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

Frameworks

Multislice solo es compatible con cargas de trabajo de JAX y PyTorch.

Paralelismo

Recomendamos a los usuarios probar Multislice con el paralelismo de datos. Si quieres obtener más información para implementar el paralelismo de canalizaciones 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 a través del formulario de comentarios o asistencia de Cloud TPU.