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. En los trabajos de entrenamiento que requieren menos de 4,096 chips, una sola porción puede ofrecer el mejor rendimiento. Sin embargo, es más fácil disponer de varias porciones más pequeñas, lo que permite un tiempo de inicio más rápido cuando Multislice se usa con porciones más pequeñas.

Cuando se implementan en configuraciones de Multislice, 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).

No es necesario que los desarrolladores escriban código para implementar la comunicación de DCN entre secciones. 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 incluye una multislice. Cada porción en una solicitud de varias porciones es del mismo tipo de acelerador. Un tipo de acelerador consiste en un tipo de TPU (v4 o v5e) seguido de la cantidad de TensorCores. Por ejemplo, v4-128 especifica una TPU v4 con 128 TensorCores.

Reparación automática

Cuando una porción encuentra un evento de mantenimiento, una interrupción o una falla de hardware, Cloud TPU creará una porción nueva. En el caso poco frecuente de que no haya recursos suficientes para crear una porción nueva, la creación no se completará hasta que el hardware esté disponible. Una vez que se cree la nueva porción, se reiniciarán todas las demás secciones en el entorno Multislice para que el entrenamiento pueda continuar.Con una secuencia de comandos de inicio configurada de forma correcta, la secuencia de comandos de entrenamiento puede reiniciarse automáticamente sin intervención del usuario, cargar y reanudar desde el último punto de control.

Conjunto de datos

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

Herramientas de 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, 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 VM a la vez. Cada VM tiene una TPU dedicada.

Inferencia

Cargar un modelo de aprendizaje automático previamente entrenado en un host y hacer predicciones a partir de los datos

Interchip Interconnect (ICI)

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

Multislice

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

Node

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

Pod

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

Recurso en cola (QR)

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

Secuencia de comandos de inicio

Una secuencia de comandos de inicio de Compute Engine estándar que se ejecuta cada vez que se inicia o reinicia una VM. En el caso de Multislice, se especifica en la solicitud de creación 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í mediante la red ICI.

VM de TPU

Es 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, a una VM de TPU la llamamos trabajador.

Tensor

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

Unidad de procesamiento tensorial (TPU)

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

Tipos de capacidad de Cloud TPU

Las TPU se pueden crear a partir de diferentes tipos de capacidad (consulta las opciones de uso en Cómo funcionan los precios de 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.
• Spot: Se orienta a la cuota interrumpibles mediante VMs Spot. Tus recursos pueden interrumpirse para hacer lugar a las solicitudes de un trabajo de mayor prioridad. Usa la marca --spot cuando crees tus recursos.
• A pedido: Se orienta a las cuotas según demanda, que no necesitan una reserva y no se interrumpirán. La solicitud de TPU se pondrá en cola en una cola de cuotas según demanda que ofrece Cloud TPU; la disponibilidad de los recursos no está garantizada. 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 deseas conocer los límites de tamaño de porción para las solicitudes según demanda, consulta Tamaños predeterminados de los segmentos para QR a pedido.

Comenzar

Si no has usado TPU antes, comienza por instalar Google Cloud CLI y configurar tu entorno de Cloud TPU. Para usar Multislice, los recursos TPU deben administrarse 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 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, arbitrariamente escalable, de código abierto y comprobado que está escrito en Python y Jax. MaxText se diseñó para entrenarse de manera eficiente en Cloud TPU.

El código en shardings.py está diseñado para ayudarte a comenzar a experimentar con diferentes opciones de paralelismo. Por ejemplo, paralelismo de datos, paralelismo de datos completamente fragmentados (FSDP) y paralelismo de tensor. El código escala de entornos de una sola porción a entornos multislice.

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 de cada método de paralelismo.

Estas entradas se deben restringir 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 usuario para el paralelismo de ICI para los cuatro chips disponibles en v4-8:

	ici_data_parallelism	ici_fsdp_parallelism	ici_tensor_parallelism
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 debe dejarse como 1 en la mayoría de los casos, ya que la red ICI es lo suficientemente rápida como para preferir casi siempre FSDP en lugar de 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. 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 QR.

   accelerator-type

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

   runtime-version

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

   node-id

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

   reserved

   Usa la cuota reservada cuando creas las porciones.

   best-effort

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

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

4. Espera hasta que el código 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 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) en 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 lograr alrededor de 260 TFLOP por segundo 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 alto ancho de banda (HBM) de la TPU.

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 la porción de QR y 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 rastrea por separado de la cuota de TPU.

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

Fragmentación de varias porciones con el paralelismo de DCN

La secuencia de comandos shardings.py toma tres parámetros que especifican el paralelismo de DCN, que corresponden a la cantidad de fragmentos de cada tipo de paralelismo de datos:

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

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

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

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

dcn_tensor_parallelism siempre debe estar configurado como 1 porque el DCN no es adecuado para esa fragmentación. Para cargas de trabajo típicas de LLM en chips v4, dcn_fsdp_parallelism también se debe configurar 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 constantes), aumentas la cantidad de paralelismo de datos.

Ejecuta shardings.py en un entorno multislice

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 de MaxText, excepto que aquí ejecutamos shardings.py en lugar del LLM más complejo en train.py.

La herramienta multihost_runner.py está optimizada para experimentos rápidos, 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 en cualquier VM de Compute Engine en el mismo proyecto que tus porciones. No es posible ejecutar multihost_runner.py en un trabajador.

multihost_runner.py se conecta de forma automática a los trabajadores TPU a través de SSH.

En este ejemplo, ejecutamos shardings.py en dos porciones 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 de ejecutor.

2. Ve al directorio del repositorio.

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

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

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

     $ export SLICE_COUNT=2
     

5. Crea un entorno de Multislice con 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 la Política de cuotas.

   En el siguiente comando, se muestra cómo crear una TPU de Multislice v4. Para usar v5e, especifica accelerator-type v5e (por ejemplo, v5litepod-16) y runtime-version 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 QR.

   accelerator-type

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

   runtime-version

   La versión de software de Cloud TPU.

   node-count

   La cantidad de segmentos que se crearán.

   node-prefix

   El prefijo que se usa con el objetivo de 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, el nombre de las secciones será mySlice-0, mySlice-1 y así sucesivamente.

   reserved

   Usa la cuota reservada cuando creas las porciones.

   best-effort

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

6. Cuando se inicia el aprovisionamiento del código QR, puede tardar hasta cinco minutos en completarse, según el tamaño del código QR. 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 está en el estado WAITING_FOR_RESOURCES o PROVISIONING durante más de 15 minutos.

7. Instala las dependencias.

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

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

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

   Verás aproximadamente 230 TFLOPs 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 varias porciones, realiza los siguientes cambios de 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 mueva a Multislice. Para lograr un alto rendimiento, el DCN debe asignarse a ejes paralelos de datos, datos fragmentados por completo 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 Multislice para obtener el máximo rendimiento.

Para validar que tu código pueda acceder a todos los dispositivos, puedes confirmar que len(jax.devices()) sea igual a la cantidad de chips del 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 porción para entornos multislice

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

Ejecuta el trabajo en varias secciones

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

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 experimentación

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

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

Secuencia de comandos del ejecutor de producción

Para trabajos de producción que necesitan resiliencia frente a fallas de hardware y otras interrupciones, es mejor integrarlos directamente en la API de Create Queued Resource. Como ejemplo de funcionamiento, 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 caso de interrupción. La secuencia de comandos multihost_job.py se documenta en el archivo README de MaxText.

Debido a que multihost_job.py debe aprovisionar recursos para cada ejecución, no proporciona un ciclo de iteración tan rápido como multihost_runner.py.

Enfoque manual

Te recomendamos usar o adaptar multihost_runner.py o multihost_job.py para ejecutar tu 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 de varias porciones.

Administra un entorno de Multislice

Para aprovisionar y administrar manualmente los códigos QR sin usar las herramientas que se proporcionan en el repositorio de MaxText, consulta las siguientes secciones.

Crea códigos 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	Es el ID del 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. Limitado a un máximo de 256 slices.
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 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 QR.

project

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

zone

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

node-count

La cantidad de segmentos que se crearán.

accelerator-type

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

runtime-version

La versión de software de Cloud TPU.

network

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

subnetwork

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

reserved

Usa la cuota reservada cuando creas las porciones.

spot

Usa la cuota de VMs Spot cuando crees las porciones.

Asegúrate de tener la cuota correspondiente antes de seleccionar --reserved, --spot o la cuota según demanda predeterminada. Si deseas obtener información sobre los tipos de cuotas, consulta la 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: El número de tu proyecto de Google Cloud
• your-zone: Es la zona en la que deseas crear el código QR.
• accelerator-type: Es la versión y el tamaño de una sola porción.
• tpu-vm-runtime-version: Son las versiones del entorno de ejecución de VM de TPU.
• your-network-name: Es una red a la que se conectará el código QR (opcional).
• your-subnetwork-name: Es una subred a la que se conectará el QR (opcional).
• example-tag-1: Es una cadena de etiqueta arbitraria (opcional).
• your-startup-script: Es una secuencia de comandos de inicio que se ejecutará cuando se asigne el código QR.
• slice-count: Es la cantidad de porciones de TPU en el entorno multislice.
• your-qr-id: Es el ID que proporcionó el usuario para el QR.

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: 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 a pedido predeterminada.

Envía la solicitud de creación de código 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: El ID del proyecto de Google Cloud
• your-zone: Es la zona en la que deseas crear el código QR.
• your-qr-id: Es el ID que proporcionó el usuario para el 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 del atributo name para obtener información sobre la solicitud QR.

Recupera 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: El ID del proyecto de Google Cloud
• your-zone: Es la zona en la que se creará el código QR.
• your-qr-guid: Es el GUID que sigue a name en el resultado de la solicitud de creación de 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 creó correctamente ("done = true"), el estado dentro del campo response será WAITING_FOR_RESOURCES o FAILED. Si el código QR está en estado WAITING_FOR_RESOURCES, significa que se puso en cola y comenzará a aprovisionarse cuando haya suficientes recursos. Si el código QR tiene el 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 se complete la operación, usa la función describir QR para supervisar las etapas del código QR.

En una situación poco común, es posible que encuentres el código QR en el estado FAILED, mientras que algunas partes 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 las 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 con SSH en una sola porción. Para conectarte a todas las VM de TPU en el entorno de Multislice a través de SSH y, además, instalar las dependencias, usa el siguiente comando gcloud:

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

Este comando gcloud envía el comando especificado a todos los trabajadores y nodos de 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 hay una falla 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. Si la cantidad de porciones que usas supera el límite de subprocesos de tu computadora local (también llamado límite de lotes), se producirá 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, por ejemplo, 100, el comando SSH dividirá las secciones en lotes. Ejecutará la secuencia de comandos de entrenamiento en el primer lote de 64 slices y esperará a que se completen las secuencias de comandos antes de ejecutar la secuencia de comandos en el lote restante de 36 slices. Sin embargo, el primer lote de 64 partes no se puede completar hasta que las 36 partes restantes comiencen a ejecutar la secuencia de comandos, lo que causa 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 especificas con la marca --command. Cuando hagas esto, después de iniciar la secuencia de comandos de entrenamiento en el primer lote de porciones, el control regresará 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 slices. Deberás canalizar las transmisiones de stdout y stderr de manera adecuada cuando ejecutes los comandos en segundo plano. Para aumentar el paralelismo dentro del mismo QR, puedes seleccionar segmentos específicos con el parámetro --node.

Configuración de red

Sigue estos pasos para asegurarte de que las porciones de TPU puedan comunicarse entre sí. Instala JAX en cada una de las partes. Para obtener más información, consulta Ejecuta 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 hacer esto, en cada porción, ejecuta lo siguiente:

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

Si ejecutas este código en cuatro partes de la v4-16, hay ocho chips por porción y cuatro porciones, jax.devices() debe devolver un total de 32 chips (dispositivos).

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

Describir los QR

Para ver la configuración detallada y el estado de un código QR, usa la API de describe 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 código 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

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

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

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 la memoria de la máquina de manera forzosa 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 restablecer todas las VM es más fácil que la depuración.

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 códigos 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 el caso de una interrupción, Multislice ofrece una reparación sin intervención de la porción afectada y restablecer todas las secciones posteriormente. La porción afectada se reemplaza con una nueva, y las partes restantes que se encuentran 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 de forma automática 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 Create QR request.

La siguiente secuencia de comandos de inicio (que se usa en Crear QR) te permite recuperarte automáticamente de 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 probarlo.

Generación de perfiles y depuración

La creación de perfiles es la misma en entornos de una sola porción y de multislices. Para obtener más información, consulta Cómo crear perfiles de programas JAX.

Optimiza las capacitaciones

Cómo fragmentar con Multislice para obtener el máximo rendimiento

Lograr el máximo rendimiento en entornos de Multislice requiere que consideres cómo fragmentar varias secciones. 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 llamadas paralelismo de tensor) porque requiere demasiado ancho de banda entre porciones. Para todas estas estrategias, puedes mantener la misma estrategia de fragmentación dentro de una porción que te haya funcionado en el pasado.

Recomendamos comenzar con paralelismo de datos puros. El uso de paralelismo de datos completamente fragmentados es útil para liberar el uso de memoria. La desventaja es que la comunicación entre partes usa la red DCN y ralentizará tu carga de trabajo. Usa el paralelismo de canalizaciones 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 puro funcionará bien en los casos en los que tienes una carga de trabajo que se ejecuta bien, pero deseas mejorar su rendimiento mediante el escalamiento en varias porciones.

Con el fin de lograr un escalamiento sólido en varias porciones, la cantidad de tiempo necesaria para realizar Allreduce en DCN debe ser menor que la cantidad de tiempo necesaria para realizar un pase hacia atrás. El DCN se usa para la comunicación entre porciones y es un factor limitante en la capacidad de procesamiento de la carga de trabajo.

Cada chip 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 / 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 DCN. Si usas más de dos segmentos, el compilador realizará una operación de All-Reduce en Shuffle completo y usarás hasta 64 bits de ancho de banda de DCN para cada parámetro por paso. La cantidad de FLOPS necesaria para cada parámetro variará según el modelo. Específicamente, para los modelos de lenguaje basados en Transformer, la cantidad de FLOPS necesaria para un avance y un retroceso es de aproximadamente 6 * B * P, donde:

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

La cantidad de FLOPS por parámetro es 6 * B y la cantidad de FLOPS por parámetro durante el pase hacia atrás es de 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 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

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

Para otras arquitecturas del modelo, deberás estimar el tiempo de ejecución del 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 para reducir todo en DCN y obtener una buena estimación de si el paralelismo de datos tiene 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 los nodos) con la fragmentación de los pesos entre nodos. Para cada operación en las pasadas y hacia delante, los pesos se reúnen todos para que cada porción tenga los pesos que necesita. En lugar de sincronizar los gradientes con Allreduce, los gradientes se dispersan 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 el paralelismo de datos, la FSDP requerirá escalar el tamaño del lote global de manera lineal con la cantidad de porciones. El FSDP disminuirá la presión de la memoria a medida que aumentes la cantidad de porciones. Esto se debe a que la cantidad de ponderaciones y el estado del optimizador por porción disminuye, pero lo hace según el precio de aumento del tráfico de red y la mayor posibilidad de bloqueo debido a un grupo retrasado.

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

Las operaciones de recopilación y reducción total 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 máximo de lote preferido. El paralelismo de canalizaciones permite que las porciones que componen una canalización “compartan” un lote. Sin embargo, el paralelismo de canalizaciones tiene dos desventajas importantes:

1. Genera el “burbuja de canalización”, en el 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 canalizaciones debe usarse solo 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 una FSDP de alto rendimiento. El FSDP tiende a lograr un mayor uso de FLOP del 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 canalizaciones, pero tu carga de trabajo puede ser diferente.

Si el paralelismo de canalizaciones es necesario, recomendamos combinarlo con paralelismo de datos o FSDP. Esto te permitirá minimizar la profundidad de la canalización mientras aumentas el tamaño de lote por canalización hasta que la latencia de DCN se vuelva menos importante en la capacidad de procesamiento. Concretamente, si tienes N porciones, considera canalizaciones de profundidad 2 y réplicas N/2 de paralelismo de datos, luego canalizaciones de profundidad 4 y réplicas 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 de DCN puedan ocultarse detrás de la aritmética. Esto minimizará la demora que presenta 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 multislice

Carga de datos

Durante el entrenamiento, cargamos repetidamente los lotes de un conjunto de datos para alimentar el modelo. Es importante contar con un cargador de datos asíncrono y eficiente que fragmente el lote entre los hosts para evitar la falta de trabajo de las TPU. El cargador de datos actual de 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 puntos de control de Orbax proporciona primitivas para controlar JAX PyTrees en el almacenamiento local o en el almacenamiento de Google Cloud. Proporcionamos una integración de referencia con punto de control síncrono en MaxText en checkpointing.py.

Configuraciones admitidas

Formas

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

Organización

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

Frameworks

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

Paralelismo

Recomendamos a los usuarios que prueben Multislice con paralelismo de datos. Si deseas 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

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