Descripción general de Cloud TPU Multislice

Cloud TPU Multislice es una tecnología que escala el rendimiento de pila completa que permite que un trabajo de entrenamiento use varias porciones de TPU en un solo Pod o en porciones en múltiples 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 TPU en diferentes se comunican mediante la transferencia de datos a las CPU (hosts) que, a su vez, transmiten los datos en 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

Es la forma de cada porción de TPU que comprende una Multislice. Cada en una solicitud de varias secciones 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 núcleos de tensor.

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 excepcional 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

Los datos que usa un modelo para entrenamiento o inferencia.

Herramientas de redes de centros de datos (DCN)

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

Programación de pandillas

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

Host

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

Inferencia

Cargar un modelo de aprendizaje automático previamente entrenado en un host y hacer predicciones sobre con datos valiosos y útiles.

Interchip Interconnect (ICI)

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

Multislice

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

Node

En el contexto de varias porciones, el nodo hace referencia a una sola porción de TPU. Cada En una porción de TPU, se le asigna un ID de nodo.

Pod

Es una colección de chips TPU conectados por interfaces de red dedicadas de ICI. R El 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 para poner en cola y administrar una solicitud de un entorno de TPU de una sola porción o de multislice.

Secuencia de comandos de inicio

Una secuencia de comandos de inicio de Compute Engine estándar. que se ejecuta cada vez que se inicia o reinicia una VM. Para Multislice, se especifica en la solicitud de creación del código QR. 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. Para TPU v4, cada VM de TPU tiene acceso directo a cuatro chips. A veces llamamos TPU La VM es 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 ofrece un procesamiento rápido y de bajo consumo para tareas clave de aprendizaje automático, como la multiplicación de matrices.

Tipos de capacidad de Cloud TPU

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

• Reserva: se orienta a la cuota reservada. Para usar 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. Tu recursos pueden ser interrumpibles para dar lugar a solicitudes de trabajo prioritario. Usa la marca --spot cuando crees tus recursos.
• Según demanda: Se orienta a la cuota según demanda, que no necesita una reserva. y no se interrumpirán. La solicitud de TPU se pondrá en cola en un cola de cuota que ofrece Cloud TPU, la disponibilidad de recursos no garantizada. Seleccionada de forma predeterminada, no se necesitan marcas.

Comenzar

Si no has usado TPU anteriormente, primero instala Google Cloud CLI, y configura tu entorno de Cloud TPU. Para usar Con multislice, tus recursos TPU se deben administrar 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 una herramienta de código abierto, de alto rendimiento, escalable de manera arbitraria y comprobada un LLM básico 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 escala de una sola porción a Multislice. entornos de prueba.

Paralelismo de ICI

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

• ici_data_parallelism
• ici_fsdp_parallelism
• ici_tensor_parallelism

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

Estas entradas deben 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 tensorial de 4 vías	1	1	4
FSDP bidireccional + paralelismo tensorial bidireccional	1	2	2

Ten en cuenta que ici_data_parallelism debe dejarse como 1 en la mayoría de los casos, 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. Recomendamos dejar una contraseña en blanco (presiona ingresa dos veces después de ejecutar el siguiente comando). Si se te solicita que El archivo google_compute_engine ya existe; reemplaza la versión existente.

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

3. Aprovisiona tus TPU con el siguiente comando:

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

   Descripciones de las marcas de comandos

   your-qr-id

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

   accelerator-type

   El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. 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

   Es el ID de los recursos TPU que se crearán en respuesta al Solicitud de QR.

   reserved

   Usar la cuota reservada cuando crees las porciones.

   best-effort

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

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

4. Espera hasta que el QR esté en el estado ACTIVE, lo que significa que los nodos trabajadores están en el estado READY. Una vez que se inicia el aprovisionamiento del código QR, 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 algunos minutos 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. Si lo deseas, puedes explorar otras estrategias de fragmentación sobre ICI, por ejemplo, podrías 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 estas limpiezas pasos del entorno en el que configuraste la porción (primero ejecuta exit hasta salga 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.

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

Fragmentación multislice mediante el paralelismo de DCN

La secuencia de comandos shardings.py toma tres parámetros que especifican el paralelismo de DCN. correspondiente al número de fragmentos de cada tipo de paralelismo de datos:

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

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

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

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

dcn_tensor_parallelism siempre debe configurarse en 1 porque el DCN es deficiente. la opción adecuada para dicha fragmentación. Para cargas de trabajo típicas de LLM en chips v4, dcn_fsdp_parallelism también debe establecerse en 1, por lo que dcn_data_parallelism se debe establecer en la cantidad de porciones, pero esta es 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 varias porciones 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 realizar experimentos rápidos de forma reiterada. volver a usar las mismas TPU. Debido a que la secuencia de comandos multihost_runner.py depende de conexiones SSH de larga duración, no lo 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 runner para indicar la máquina en la que Ejecuta la secuencia de comandos multihost_runner.py. Usamos el término trabajadores para indicar la las VMs de TPU que conforman las porciones. Puedes ejecutar multihost_runner.py en una aplicación local Compute Engine o cualquier VM de Compute Engine en el mismo proyecto que tus slices. Correr multihost_runner.py en un trabajador no es compatible.

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

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

Configure su entorno

1. Clona MaxText en tu ejecutor máquina.

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 aparece el mensaje el archivo google_compute_engine ya existe, selecciona no conservar tu versión existente.

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

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

     $ export SLICE_COUNT=2
     

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

   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 versión 5e runtime-version (v2-alpha-tpuv5-lite).

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

   Descripciones de las marcas de comandos

   your-qr-id

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

   accelerator-type

   El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas crear. 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

   La cantidad de porciones que se crearán.

   node-prefix

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

   reserved

   Usar la cuota reservada cuando crees las porciones.

   best-effort

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

6. Cuando se inicia el aprovisionamiento del 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 mediante multihost_runner.py.

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

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

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, haz que la 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 migra a Multislice. Para lograr un alto rendimiento, la DCN debe asignarse a datos paralelos, de datos fragmentados o en ejes paralelos de canalización. Las consideraciones de rendimiento las estrategias de fragmentación se analizan con más detalle en Fragmentación con varias porciones para obtener el máximo rendimiento.

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

Elige tamaños de porciones para entornos de varias porciones

Para obtener una velocidad lineal, agrega porciones nuevas del mismo tamaño que 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 varias porciones:

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

Secuencia de comandos del ejecutor de experimentos

La 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 adapta los pasos del instructivo de multihost_runner.py. a tu carga de trabajo y configuración de hardware.

Secuencia de comandos del ejecutor de producción

Para 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. A modo de ejemplo, brindamos multihost_job.py, lo que activa la llamada a la API de Created Queued Resource con el inicio para ejecutar el entrenamiento y reanudarlo en la interrupció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 QR, consulta Administra un entorno de varias porciones.

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.

Crear QR

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

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

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

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

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

Para obtener más información, consulta la API de recursos en cola de REST. documentación para todas las opciones disponibles.

Para usar la capacidad de Spot, reemplaza lo siguiente:

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

Quita la línea para usar la capacidad según demanda predeterminada.

Envía la solicitud de creación de QR con la carga útil de JSON:

  $ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" -d @queuedresourcereq.json https://tpu.googleapis.com/v2alpha1/projects/your-project-id/locations/your-zone/queuedResources\?queued_resource_id\=your-qr-id

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

La respuesta debería verse de la siguiente manera:

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

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

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

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

  $ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2/projects/your-project-id/locations/your-zone/operations/operation-your-qr-guid

• 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 creó correctamente ("done = true"), el estado dentro del El 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 estará 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 una situación poco común, es posible que el código QR tenga el estado FAILED, mientras que 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 de QR con 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 el lote actual completa la ejecución. 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 segmentos que usas supera el subproceso de tu computadora local límite (también llamado límite de lotes), te encontrarás con un interbloqueo. A modo de ejemplo, supongamos que el límite de lotes en tu máquina local es 64. Si intentas ejecutar un secuencia de comandos de entrenamiento en más de 64 porciones, digamos 100, el comando SSH las porciones en lotes. Ejecutará la secuencia de comandos de entrenamiento en el primer lote de 64. porciones y esperar a que se completen las secuencias de comandos antes de ejecutarlas en la lo que queda de 36 porciones. Sin embargo, el primer lote de 64 porciones no puede hasta que las 36 fracciones restantes comiencen a ejecutar la secuencia de comandos. 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. Después de iniciar la secuencia de comandos de entrenamiento, en el primer lote de porciones, el control volverá inmediatamente a el 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 tu stdout y stderr de forma adecuada cuando se ejecutan los comandos en segundo plano. Para aumentar paralelismo dentro del mismo código QR, puedes seleccionar porciones específicas con el --node parámetro.

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

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

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

Enumerar QR

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

$ gcloud compute tpus queued-resources list

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

Describe los QR

Para ver la configuración detallada y el estado de un código QR, usa 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 QR. Para obtener más información sobre los posibles estados de QR, consulta Recursos en cola.

Inicia tu trabajo en un entorno aprovisionado

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

Se puede usar la API de ResetQueuedResource para restablecer todas las VMs en un QR de ACTIVE. Si restableces las VMs por la fuerza, se borra la memoria y 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. El 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 ResetQueuedResource la operación tarda de uno a dos minutos en completarse. Para invocar la API, usa el siguiente comando: :

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

Borrando QR

Para liberar recursos al final de 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 verifique y carga 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 Cómo crear código QR) te permite recuperarte automáticamente de las fallas y reanudar el entrenamiento de 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 creación de perfiles es la misma en 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 fragmentar activaciones en las dimensiones del modelo (a veces, llamado paralelismo de tensores), ya que requiere demasiado ancho de banda entre segmentos. 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.

Recomendamos comenzar con el paralelismo de datos puros. Usa datos completamente fragmentados el paralelismo es útil para liberar el uso de memoria. La desventaja es que la comunicación entre porciones usa la red DCN y ralentiza 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 múltiples segmentos, la cantidad de tiempo necesaria para realizar Allreduce a DCN debe ser menor que la cantidad de tiempo requerida para hacer un pase hacia atrás. DCN se usa para la comunicación entre porciones y es un factor limitante en la capacidad de procesamiento de la carga de trabajo.

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

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

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

Tu modelo usará entre 32 y 64 bits de ancho de banda de DCN para cada parámetro por paso. Si usas dos porciones, tu modelo usará 32 bits de ancho de banda de DCN. Si Usa más de dos porciones y el compilador realizará un Allreduce de Shuffle completo y usa 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 model. 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 múltiples porciones, asegúrate de que los procesos supera la intensidad aritmética del hardware de la TPU. Para calcular intensidad operativa, divide la cantidad de FLOPS por parámetro durante la Retroceder 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. Como la red DCN puede descartar paquetes brevemente, es mejor mantener un margen de error significativo, implementando solo el paralelismo de datos 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, necesitas estimar el tiempo de ejecución de tu pase hacia atrás por porción (ya sea cronometrando con un generador de perfiles o contando FLOPS). Luego, puedes 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 fragmentados (FSDP)

El paralelismo de datos completamente fragmentado (FSDP) combina el paralelismo de datos (fragmentación de datos entre nodos) con la fragmentación de las ponderaciones entre los nodos. Para cada operación en hacia adelante y hacia atrás, los pesos se juntan para que cada porción tiene los pesos que necesita. En lugar de sincronizar los gradientes con Allreduce, los gradientes se dispersan por su reducción a medida que se producen. De esta manera, cada porción solo obtiene los gradientes de los pesos de los que es responsable.

Al igual que en el paralelismo de datos, FSDP requerirá escalar el tamaño del lote global de forma lineal con el número de porciones. FSDP disminuirá la presión de memoria a medida que aumentas 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 con un aumento del tráfico de red y la mayor posibilidad de bloqueo debido a colectiva.

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 el retroceder o aumentar la cantidad de parámetros en la red neuronal.

Las operaciones de recopilación y Allreduce en la FSDP funcionan de forma similar a las de la DP para determinar si tu carga de trabajo FSDP está limitada por el rendimiento de DCN en de la misma manera que se describe 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 de paralelismo que requieren un tamaño de lote global mayor que el tamaño de lote máximo preferido. El paralelismo de canalizaciones permite que las slices compuesta por una canalización para “compartir” por lotes. Sin embargo, el paralelismo de la canalización tiene dos importantes desventajas:

1. Genera el “cuadro de canalización”. en la que los chips están inactivos porque están esperando para los 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 canalizaciones, vale la pena experimentar para ver empíricamente 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. Más probable las cargas de trabajo pueden tolerar tamaños de lotes lo suficientemente grandes como para no beneficiarse paralelismo de la canalización, pero la 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. Concretamente, si tienes N porciones, considera canalizaciones de réplicas de profundidad 2 y N/2 del paralelismo de datos, luego, canalizaciones de profundidad 4 y N/4. réplicas de paralelismo de datos, etc., hasta que el lote por canalización aumente suficiente que los colectivos DCN pueden estar ocultos detrás de la aritmética en la y retroceder. Esto minimizará la demora que introduce la canalización paralelismo, a la vez que permite escalar más allá del límite de tamaño del lote global.

Prácticas recomendadas de multislice

Carga de datos

Durante el entrenamiento, cargamos repetidamente lotes desde un conjunto de datos para alimentar el model. Tener un cargador de datos asíncrono y eficiente que fragmenta el lote de los hosts es importante para no privar a las TPU de trabajo. El cargador de datos actual en MaxText hace que cada host cargue un subconjunto igual de los ejemplos. Esta solución es adecuada para texto, pero requiere refragmentación dentro del modelo. Además, MaxText todavía no ofrece una instantánea determinista, lo que permitiría al iterador de datos los mismos datos antes y después de la interrupción.

Controles

La biblioteca de control de Orbax proporciona para el control de JAX PyTrees en el almacenamiento local o el de 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 organización es compatible con GKE. Para obtener más información, consulta TPU en GKE.

Frameworks

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

Paralelismo

Recomendamos a los usuarios probar Multislice con el paralelismo de datos. 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 con el formulario de comentarios o asistencia de Cloud TPU.