Descripción general de Cloud TPU Multislice

Cloud TPU Multislice es una tecnología de escalamiento de rendimiento de pila completa que permite que una tarea de entrenamiento use varias porciones de TPU dentro de un solo pod o en porciones de varios pods con un paralelismo de datos simple. Con los chips TPU v4, esto significa que los trabajos de entrenamiento pueden usar más de 4,096 chips en una sola ejecución. Para entrenamiento que requieren menos de 4,096 chips, una sola porción puede ofrecer la mejor rendimiento. Sin embargo, las porciones más pequeñas son más fáciles de encontrar lo que permite un tiempo de inicio más rápido cuando se usa Multislice con menor porciones.

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 de TPU en diferentes porciones se comunican transfiriendo datos a las CPU (hosts), que a su vez transmiten los datos a través de la red del centro de datos (DCN).

Los desarrolladores no tienen que escribir código para implementar la comunicación DCN entre porciones. El compilador de XLA genera ese código por ti y superpone la comunicación con y procesamiento para maximizar el rendimiento.

Conceptos

Tipo de acelerador

La forma de cada porción de TPU que comprende una Multislice. Cada fragmento de una solicitud de varias porciones es del mismo tipo de acelerador. Un 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 TensorCores.

Reparación automática

Cuando una porción encuentra un evento de mantenimiento, una interrupción o una falla de hardware, Cloud TPU creará una porción nueva. En el caso poco frecuente de que haya recursos insuficientes para crear una nueva porción, la creación no se completará hasta que el hardware esté disponible. Después de crear la nueva porción en el entorno de varias porciones se reiniciarán, por lo que el entrenamiento Con una secuencia de comandos de inicio configurada correctamente, la secuencia de comandos puede reiniciarse automáticamente sin intervención del usuario, y cargar y reanudar del punto de control más reciente.

Conjunto de datos

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

Herramientas de redes de centros de datos (DCN)

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

Programación 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 correctamente.

Host

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

Inferencia

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

Interconexión entre chips (ICI)

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

Multislice

Dos o más porciones de chip de 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. R El Pod te permite distribuir la carga de procesamiento en varias TPU.

Recurso en cola (QR)

Es una representación de los recursos de TPU que se usa para poner en cola y administrar una solicitud de un entorno de TPU de una sola porción o de 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. Para Multislice, se especifica en la solicitud de creación del código QR. 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 en un slice 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. En el caso de las TPU v4, cada VM de TPU tiene acceso directo a cuatro chips. A veces, llamamos a una VM de TPU un trabajador.

Tensor

Una estructura de datos que se usa para representar datos multidimensionales en una máquina de aprendizaje automático.

Unidad de procesamiento tensorial (TPU)

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

Tipos de capacidad de Cloud TPU

Las TPU se pueden crear a partir de diferentes tipos de capacidad (consulta las opciones de uso en Cómo funcionan los precios de las TPU):

• Reserva: Se segmenta para la cuota reservada. Para usar la cuota reservada, debes tener un de reservación con Google. Usa la marca --reserved cuando crees tus recursos.
• Spot: Se orienta a la cuota interrumpible con VMs Spot. Es posible que se usurpen tus recursos para dejar espacio para solicitudes de un trabajo de prioridad más alta. Usa la marca --spot cuando crees tus recursos.
• A pedido: se orienta a la cuota según demanda, que no necesita una reserva. y no se interrumpirán. La solicitud de TPU se pondrá en cola en una cola de cuotas on demand que ofrece Cloud TPU. No se garantiza la disponibilidad de los recursos. Seleccionada de forma predeterminada, no se necesitan marcas.

Nota: Cloud TPU limita el tamaño de la porción para todos los tipos de aceleradores. Para Slice y límites de tamaño para las solicitudes según demanda, consulta Tamaños predeterminados de las porciones para los QR según demanda.

Comenzar

Si no has usado TPU anteriormente, primero instala Google Cloud CLI, y configura tu entorno de Cloud TPU. Para usar Multislice, tus recursos de TPU deben administrarse como Recursos en cola.

Si ya eres usuario de TPU v4 y tienes una reserva, es posible que debas migrar tu reserva a un nuevo sistema de reservas. 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 eficientemente Cloud TPU

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

Paralelismo de ICI

ICI se refiere a la interconexión de alta velocidad que conecta las TPU en un solo porción. El fragmentación de ICI corresponde a la fragmentación dentro de una porción. shardings.py proporciona tres parámetros de paralelismo de ICI:

• ici_data_parallelism
• ici_fsdp_parallelism
• ici_tensor_parallelism

Los valores que especifiques para estos parámetros determinan la cantidad de fragmentos para cada método de paralelismo.

Estas entradas deben restringirse para que ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism es igual a al número de chips en la porción.

En la siguiente tabla, se muestran ejemplos de entradas del usuario para el paralelismo de ICI para los cuatro chips disponibles en las versiones 4 a 8:

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

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

En este ejemplo, se supone que estás familiarizado con la ejecución de código en una sola porción de TPU como en Ejecuta un cálculo en una VM de Cloud TPU con JAX. En este ejemplo, se muestra cómo ejecutar shardings.py en una sola porción.

1. Configura el entorno:

   $ gcloud auth login
   $ gcloud config set project your-project-id
   $ gcloud config set compute/zone your-zone

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

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

3. Aprovisiona tus TPUs con el siguiente comando:

   $ gcloud compute tpus queued-resources \
   create your-qr-id \
   --accelerator-type your-accelerator-type \
   --runtime-version tpu-ubuntu2204-base \
   --node-id qr-id \
   [--reserved |--spot]

   Descripciones de las marcas de comandos

   your-qr-id

   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. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.

   runtime-version

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

   node-id

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

   reserved

   Usar la cuota reservada cuando crees las porciones.

   spot

   Usar la cuota de VMs Spot cuando se creen los segmentos

   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 código QR esté en el estado ACTIVE, lo que significa que los nodos de trabajo están en el estado READY. Una vez que se inicia el aprovisionamiento del código QR, pueden transcurrir de uno a cinco minutos para completarse según el tamaño del código QR. Puedes consultar 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) a tu VM de TPU.

7. En el directorio del repositorio de MaxText, ejecuta la secuencia de comandos de configuración para instalar JAX y otras dependencias en tu porción de TPU. La secuencia de comandos de configuración tarda unos minutos en ejecutarse.

   $ bash setup.sh

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

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

   Puedes ver los resultados en los registros. Tus TPU deberían alcanzar alrededor de 260 TFLOP o una impresionante utilización de más del 90%de FLOOP. En este caso, hemos seleccionar aproximadamente el lote máximo que se ajusta a la longitud de la TPU alta Memoria de ancho de banda (HBM).

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

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

10. Borra el QR y la porción de TPU cuando termines. Debes ejecutar estos pasos de limpieza desde el entorno en el que configuraste la porción (primero, ejecuta exit para salir de la sesión de SSH). La eliminación tardará entre dos y cinco minutos en completarse y se puede ejecutar en segundo plano con la marca opcional --async.

    Importante: Los códigos QR se deben borrar después de usarlos. De lo contrario, se seguirá consumiendo tu 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 de varias porciones con paralelismo de DCN

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

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

Los valores de estos parámetros deben limitarse para que dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism es igual a el número de porciones.

Como ejemplo de dos rebanadas, 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 se debe establecer en 1 porque el DCN es deficiente. la opción adecuada para dicha fragmentación. Para cargas de trabajo de LLM típicas en chips v4, dcn_fsdp_parallelism también debe establecerse en 1 y, por lo tanto, dcn_data_parallelism debe establecerse en la cantidad de rebanadas, pero esto depende de la aplicación.

A medida que aumentes la cantidad de porciones (si mantienes el tamaño de las porciones y el lote constante por porción), 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 pasos siguientes son muy similares a los Primeros pasos: Experimentos rápidos en varias secciones del repositorio de MaxText, excepto que, aquí, ejecutamos shardings.py en lugar del un LLM más complejo en train.py.

La herramienta multihost_runner.py está optimizada para experimentos rápidos y reutiliza las mismas TPU de forma reiterada. Debido a que la secuencia de comandos multihost_runner.py depende de conexiones SSH de larga duración, no la recomendamos para trabajos de larga duración. Si quieres ejecutar un trabajo más largo (por ejemplo, horas o días), te recomendamos usa multihost_job.py.

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

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

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

Configure su entorno

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

2. Ve al directorio del repositorio.

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

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

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

     $ export SLICE_COUNT=2
     

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

   Nota: Asegúrate de tener la cuota correspondiente antes de seleccionar --reserved, --spot o la cuota predeterminada a pedido. Para obtener información sobre la cuota 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 el runtime-version de v5e (v2-alpha-tpuv5-lite).

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

   Descripciones de las marcas de comandos

   your-qr-id

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

   accelerator-type

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

   runtime-version

   La versión de software de Cloud TPU.

   node-count

   Es la cantidad de rebanadas que se crearán.

   node-prefix

   Es el prefijo que se usa para generar nombres para cada fragmento. Se agrega un número al prefijo de cada porción. Por ejemplo, si estableces node-prefix a mySlice, las secciones se denominan mySlice-0, mySlice-1, continuando de forma numérica para cada porción.

   reserved

   Usa la cuota reservada cuando crees las porciones.

   spot

   Usar la cuota de VMs Spot cuando se creen los segmentos

6. Cuando se inicia el aprovisionamiento del código QR, es posible que tarde hasta cinco minutos en completarse según el tamaño del código QR. Esperar hasta que el recurso en cola (QR) se encuentre en el estado ACTIVE Puedes verificar el estado de una solicitud de código 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 de QR se encuentra en Los estados WAITING_FOR_RESOURCES o PROVISIONING durante más de 15 minutos.

7. Instala las dependencias.

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

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

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

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

9. Limpia las TPU y el 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 función --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 cambia a Multislice. Para lograr un alto rendimiento, la DCN debe asignarse a ejes paralelos de datos, de datos completamente fragmentados o de canalización. Las consideraciones de rendimiento las estrategias de fragmentación se analizan con más detalle en Fragmentación con multislice para un rendimiento máximo.

Para validar que tu código pueda acceder a todos los dispositivos, puedes afirmar que len(jax.devices()) es igual a la cantidad de chips en tu entorno de Multislice. Por ejemplo, si usas cuatro rebanadas de v4-16, tienes ocho chips por rebanada × 4 rebanadas, por lo que len(jax.devices()) debería mostrar 32.

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

Para obtener una velocidad lineal, agrega porciones nuevas del mismo tamaño que las del porción. Por ejemplo, si usas una porción v4-512, Multislice lograr aproximadamente el doble de rendimiento agregando una segunda porción de v4-512 y duplicando tu tamaño de lote global. Para 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 tu carga de trabajo personalizada en un entorno de Multislice:

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

Secuencia de comandos del ejecutor de experimentación

El multihost_runner.py La secuencia de comandos distribuye el código a un entorno de Multislice existente y lo ejecuta tu comando en cada host, copia tus registros y realiza un seguimiento de los errores de cada comando estado. La secuencia de comandos multihost_runner.py se documenta en Archivo README de MaxText

Debido a que multihost_runner.py mantiene conexiones SSH persistentes, solo adecuado para experimentación de tamaño moderado y duració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 es mejor realizar la integración directamente con el recurso API de gcloud. Como ejemplo de trabajo, proporcionamos multihost_job.py, que activa la llamada a la API de Created Queued Resource con la secuencia de comandos de inicio adecuada para ejecutar tu entrenamiento y reanudarlo en la usurpación. El multihost_job.py de comandos está documentado en la Archivo README de MaxText

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

Enfoque manual

Te recomendamos que uses o adaptes multihost_runner.py. o multihost_job.py para ejecutar tu carga de trabajo personalizada tu configuración de Multislice. Sin embargo, si prefieres aprovisionar y administrar tu entorno directamente con comandos de QR, consulta Cómo administrar un entorno de varias capas.

Administra un entorno de varias porciones

Para aprovisionar y administrar manualmente los QR sin usar las herramientas que se proporcionan en el repositorio MaxText, consulta secciones a continuación.

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	El ID asignado por el usuario del código QR
PROYECTO	Nombre del proyecto de Google Cloud
ZONA	us-central2-b
NETWORK_NAME	Es el 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 que se usan para identificar fuentes o destinos válidos para firewalls de red
SLICE_COUNT	Cantidad de porciones. Se limita a un máximo de 256 rebanadas.
STARTUP_SCRIPT	Si se agrega a la solicitud de creación, una secuencia de comandos de inicio puede ejecutarse 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 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. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.

runtime-version

La versión de software de Cloud TPU.

network

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

subnetwork

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

reserved

Usa 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 según demanda predeterminada. Para obtener información sobre los tipos de cuotas, consulta Política de cuotas.

Crea una solicitud de código 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 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: 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 cadena 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 tu entorno de Multislice.
• your-qr-id: Es el ID que proporcionó el usuario para el código QR.

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

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

Para usar la capacidad de Spot, reemplaza lo siguiente:

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

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

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 ser similar a la siguiente:

{
  "name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qr-guid>",
  "metadata": {
    "@type": "type.googleapis.com/google.cloud.common.OperationMetadata",
    "createTime": "2023-11-01T00:17:05.742546311Z",
    "target": "projects/<your-project-id>/locations/<your-zone>/queuedResources/<your-qa-id>",
    "verb": "create",
    "cancelRequested": false,
    "apiVersion": "v2alpha1"
  },
  "done": false
}

Usa el valor GUID al final del valor de cadena para el atributo name para obtener información sobre la solicitud de código QR.

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

Para obtener el estado de la solicitud de código 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 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 crea correctamente ("done = true"), el estado dentro del campo response será WAITING_FOR_RESOURCES o FAILED. Si el código QR tiene el estado WAITING_FOR_RESOURCES, significa que el código QR se pone en cola y comenzará el aprovisionamiento cuando haya suficientes recursos. Si el código QR está en el estado FAILED, el motivo de la falla aparecerá en el resultado. Para ver 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 el método describe QRs para supervisar las etapas del QR.

En casos excepcionales, es posible que encuentres tu código QR en el estado FAILED mientras que algunas porciones son ACTIVE. Si esto sucede, borra los recursos que se crearon y vuelve a intentarlo en unos minutos, o bien comunícate con nosotros. al equipo de Cloud TPU para resolver el problema.

Instala dependencias y SSH

Ejecuta el código JAX en porciones de pod de TPU se describe cómo conectarte a las VMs de TPU con SSH en una sola porción. Para conectarte a todas las VMs de TPU en tu entorno de Multislice mediante SSH y instala 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 en el QR a través de SSH. El comando se divide en lotes en grupos de cuatro y se envía al mismo tiempo. El siguiente lote de comandos se envía cuando finaliza la ejecución del lote actual. Si hay una falla con uno de los comandos, procesar paradas y no se envían más lotes. Para obtener más información, consulta la referencia de la API de recursos en cola. Si la cantidad de rebanadas que usas supera el límite de subprocesos de tu computadora local (también llamado límite de lotes), se producirá un interbloqueo. A modo de ejemplo, asumamos que el límite de lotes en tu máquina local es 64. Si intentas ejecutar una secuencia de comandos de entrenamiento en más de 64 rebanadas, por ejemplo, 100, el comando SSH dividirá las rebanadas en lotes. Ejecutará la secuencia de comandos de entrenamiento en el primer lote de 64 fragmentos y esperará a que se completen las secuencias de comandos antes de ejecutarlas en el lote restante de 36 fragmentos. Sin embargo, el primer lote de 64 rebanadas no se puede completar hasta que las 36 rebanadas restantes comiencen a ejecutar la secuencia de comandos, lo que genera un interbloqueo.

Para evitar esta situación, puedes ejecutar la secuencia de comandos de entrenamiento en segundo plano en para cada VM agregando 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 rebanadas, el control volverá de inmediato al comando SSH. Entonces, el comando SSH puede comenzar a ejecutar la secuencia de comandos de entrenamiento en el lote restante de 36 porciones. Deberás canalizar tus flujos de stdout y stderr de forma adecuada cuando ejecutes los comandos en segundo plano. Para aumentar el paralelismo dentro del mismo código QR, puedes seleccionar rebanadas específicas con el parámetro --node.

Configuración de red

Ejecuta los siguientes pasos para asegurarte de que las porciones de TPU puedan comunicarse entre sí. Instala JAX en cada una de las porciones. Para obtener más información, consulta Ejecuta el código JAX en porciones de pod de TPU. Confirma que len(jax.devices()) sea igual a la cantidad de chips en tu entorno de Multislice. Para ello, en cada fragmento, ejecuta lo siguiente:

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

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

Lista de códigos 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 el describir la API 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 código QR. Para obtener más información sobre los posibles estados de los códigos 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 fragmento a través de SSH y ejecuta el siguiente comando en todos los hosts.

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

Restablecimiento de los códigos QR

Se puede usar la API de ResetQueuedResource para restablecer todas las VMs en un QR de ACTIVE. Cuando se restablecen las VMs, se borra de manera forzosa la memoria de la máquina y se restablece la VM a su estado inicial. Todos los datos almacenados localmente permanezcan intactas y la secuencia de comandos de inicio se invocará tras el restablecimiento. La API de ResetQueuedResource puede ser útil cuando deseas reiniciar todas las TPU. Para Por ejemplo, cuando el entrenamiento se detiene y restablecer todas las VMs es más fácil que realizar una depuración.

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

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

Cómo borrar códigos QR

Para liberar recursos al final de la sesión de entrenamiento, borra los recursos en cola recurso con la marca --force. La eliminación tardará entre dos y cinco minutos en que se complete 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 reparar la porción afectada y restablecer todas las porciones posteriormente. Usuarios afectados se reemplaza por una nueva y las restantes, en caso contrario, están 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 automáticamente cada vez que se reasigna una porción o se restablece una VM. Especificas una startup 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 Crear códigos QR) te permite recuperarte automáticamente de las fallas y reanudar el entrenamiento desde los puntos de control almacenados en un bucket de Cloud Storage durante el entrenamiento de MaxText:

{
 "tpu": {
   "node_spec": [
     {
      ...
         "metadata": {
               "startup-script": "#! /bin/bash \n pwd \n runuser -l user1 -c 'cd /home/user1/MaxText && python3 MaxText/train.py MaxText/configs/base.yml run_name=run_test_failure_recovery dcn_data_parallelism=4 ici_fsdp_parallelism=8 steps=10000 save_period=10 base_output_directory='gs://user1-us-central2'' EOF"
         }
     ...
     }
   ]
 }
}

Clona el repositorio de MaxText antes de probar esto y sale de ella.

Generación de perfiles y depuración

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

Optimiza las capacitaciones

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

Lograr el máximo rendimiento en entornos de Multislice requiere teniendo en cuenta cómo fragmentar en múltiples porciones. Por lo general, hay tres de datos (paralelismo de datos, paralelismo de datos completamente fragmentados y paralelismo de canalizaciones). No recomendamos el particionamiento de activaciones en las dimensiones del modelo (a veces llamado paralelismo de tensores) porque requiere demasiado ancho de banda entre las porciones. Para todas estas estrategias, puedes mantener la misma estrategia de fragmentación dentro de una porción. que haya funcionado para ti en el pasado.

Te recomendamos que comiences con el paralelismo de datos puro. El uso del paralelismo de datos completamente fragmentados es útil para liberar el uso de memoria. El inconveniente es que la comunicación entre las porciones usa la red de DCN y ralentiza la carga de trabajo. Usa el paralelismo de 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 puros funcionará bien en los casos en los que tengas una carga de trabajo funciona bien, pero quieres mejorar su rendimiento escalando varias porciones.

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

Cada chip de TPU v4 rinde 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 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 rebanadas, tu modelo usará 32 bits de ancho de banda de DCN. Si usas más de dos rebanadas, el compilador realizará una operación de reducción completa de todos los shuffles y usarás hasta 64 bits de ancho de banda de DCN para cada parámetro por paso. La cantidad de FLOPS necesarios para cada parámetro variará según tu modelo. Específicamente, para los modelos de lenguaje basados en Transformer, la cantidad de FLOPS necesarios para avanzar y retroceder son aproximadamente 6 * B * P, donde:

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

La cantidad de FLOPS por parámetro es 6 * B y la cantidad de FLOPS por parámetro durante el pase hacia atrás es 4 * B.

Para garantizar un escalamiento sólido en varias porciones, asegúrate de que la intensidad operativa supere la intensidad aritmética del hardware de la TPU. Para calcular la intensidad operativa, divide la cantidad de FLOPS por parámetro durante el pase hacia atrás por el ancho de banda de la red (en bits) por parámetro por paso: Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

Por lo tanto, para un modelo de lenguaje basado en Transformer, si usas dos porciones, haz lo siguiente: 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 Transformer de lenguaje extensos. Debido a que la red de DCN puede descartar paquetes brevemente, es mejor mantener un margen de error significativo y, luego, implementar el paralelismo de datos solo si el tamaño del lote por pod es de al menos 350,000 (dos pods) a 700,000 (muchos pods).

Para otras arquitecturas de modelos, deberás estimar el tiempo de ejecución de tu paso hacia atrás por fragmento (ya sea cronometrándolo con un generador de perfiles o contando FLOPS). Luego, puedes compararlo con el tiempo de ejecución esperado para reducir DCN y obtén una buena estimación de si el paralelismo de datos tendrá sentido para ti.

Cuándo usar el paralelismo de datos completamente fragmentado (FSDP)

El paralelismo de datos completamente fragmentado (FSDP) combina el paralelismo de datos (fragmentación de datos entre nodos) con la fragmentación de las ponderaciones entre los nodos. Para cada operación en los pases hacia adelante y hacia atrás, los pesos se reúnen para que cada fragmento tenga los pesos que necesita. En lugar de sincronizar los gradientes con la reducción completa, los gradientes se reducen y se dispersan a medida que se producen. De esta manera, cada fragmento solo obtiene los gradientes de los pesos de los que es responsable.

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

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

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

Cuándo usar el paralelismo de canalización

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

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

El paralelismo de canalizaciones debe usarse solo si las otras estrategias de paralelismo requieren un tamaño de lote global muy grande. Antes de probar el paralelismo de canalización, vale la pena experimentar para ver de forma empírica si la convergencia por muestra se ralentiza en el tamaño del lote necesario para lograr un FSDP de alto rendimiento. El FSDP suele alcanzar mayor utilización del modelo FLOP, pero si la convergencia por muestra se ralentiza a medida que el tamaño de los lotes crece, 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 canalización, pero tu carga de trabajo puede ser diferente.

Si el paralelismo de la canalización es necesario, recomendamos combinarlo con los datos paralelismo o FSDP. Esto te permitirá minimizar la profundidad de la canalización aumentando el tamaño del lote por canalización hasta que la latencia de DCN sea menor de un tiene en cuenta la capacidad de procesamiento. Específicamente, si tienes N rebanadas, considera canalizaciones de profundidad 2 y N/2 réplicas de paralelismo de datos, luego canalizaciones de profundidad 4 y N/4 réplicas de paralelismo de datos, y así sucesivamente, hasta que el lote por canalización sea lo suficientemente grande como para que los colectivos de DCN se puedan ocultar detrás de la aritmética en el pase hacia atrás. Esto minimizará la ralentización que introduce el paralelismo de la canalización y, al mismo tiempo, te permitirá escalar más allá del límite de tamaño del lote global.

Prácticas recomendadas para las divisiones múltiples

Carga de datos

Durante el entrenamiento, cargamos repetidamente lotes desde un conjunto de datos para alimentar el un modelo de responsabilidad compartida. Es importante tener un cargador de datos asíncrono y eficiente que divida el lote entre los hosts para evitar que las TPUs se queden sin trabajo. El cargador de datos actual en MaxText hace que cada host cargue un subconjunto igual de los ejemplos. Esta solución es adecuada para el texto, pero requiere un reescalado dentro del modelo. Además, MaxText aún no ofrece instantáneas deterministas, lo que permitiría que el iterador de datos cargue los mismos datos antes y después de la usurpación.

Controles

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

Configuraciones admitidas

Formas

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

Organización

La orquestación es compatible con GKE. Para obtener más información, consulta TPUs 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. Para obtener más información sobre la implementación del paralelismo de canalizaciones con Multislice, comuníquese con su Representante de cuenta de Google Cloud.

Asistencia y comentarios

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