Introducción a Multislice de la TPU de Cloud

Multislice de TPU de Cloud es una tecnología de escalado del rendimiento de pila completa que permite que un trabajo de entrenamiento use varios slices de TPU en un solo slice o en slices de varios pods con paralelismo de datos estándar. Con los chips TPU v4, los trabajos de entrenamiento pueden usar más de 4096 chips en una sola ejecución. En el caso de las tareas de entrenamiento que requieren menos de 4096 chips, un solo segmento puede ofrecer el mejor rendimiento. Sin embargo, hay más porciones pequeñas disponibles, lo que permite que el tiempo de inicio sea más rápido cuando se usa Multislice con porciones más pequeñas.

Varias particiones escalan el rendimiento de forma lineal

Cuando se implementan en configuraciones de Multislice, los chips de TPU de cada slice se comunican a través de la interconexión entre chips (ICI). Los chips de TPU de diferentes slices se comunican transfiriendo datos a las CPUs (hosts), que a su vez transmiten los datos a través de la red del centro de datos (DCN). Para obtener más información sobre cómo escalar con Multislice, consulta Cómo escalar el entrenamiento de la IA a decenas de miles de chips de TPU de Cloud con Multislice.

Flujo de datos multirebanada

Los desarrolladores no tienen que escribir código para implementar la comunicación DCN entre slices. El compilador XLA genera ese código y superpone la comunicación con el cálculo para obtener el máximo rendimiento.

Conceptos

Tipo de acelerador
La forma de cada slice de TPU que compone un Multislice. Cada segmento de una solicitud de varios segmentos es del mismo tipo de acelerador. Un tipo de acelerador consta de un tipo de TPU (versión 4 o posterior) seguido del número de Tensor Cores. Por ejemplo, v5litepod-128 especifica una TPU v5e con 128 Tensor Cores.
Reparación automática
Cuando una porción se enfrenta a un evento de mantenimiento, una expropiación o un fallo de hardware, TPU de Cloud creará una nueva porción. Si no hay suficientes recursos para crear una nueva partición, la creación no se completará hasta que haya hardware disponible. Una vez creada la nueva porción, se reiniciarán todas las demás porciones del entorno Multislice para que el entrenamiento pueda continuar. Con un script de inicio configurado correctamente, el script de entrenamiento se puede volver a iniciar automáticamente sin que el usuario tenga que hacer nada, cargando y reanudando el entrenamiento desde el último punto de control.
Redes de centros de datos (DCN)
Una red con mayor latencia y menor rendimiento (en comparación con ICI) que conecta sectores de TPU en una configuración Multislice.
Programación de bandas
Cuando se aprovisionan todas las rebanadas de TPU al mismo tiempo, se garantiza que se aprovisionan todas o ninguna correctamente.
Interconexión entre chips (ICI)
Enlaces internos de alta velocidad y baja latencia que conectan las TPUs de un pod de TPU.
Multicorte
Dos o más porciones de chips de TPU que pueden comunicarse a través de DCN.
Nodo
En el contexto de Multislice, el término "nodo" hace referencia a una sola slice de TPU. A cada porción de TPU de un Multislice se le asigna un ID de nodo.
Script de inicio
Una secuencia de comandos de inicio de Compute Engine estándar que se ejecuta cada vez que se inicia o reinicia una VM. En el caso de Multislice, se especifica en la solicitud de creación de QR. Para obtener más información sobre las secuencias de comandos de inicio de las TPU de Cloud, consulta Gestionar recursos de TPU.
Tensor
Estructura de datos que se usa para representar datos multidimensionales en un modelo de aprendizaje automático.
Tipos de capacidad de TPU de Cloud

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

  • Reserva: para consumir una reserva, debes tener un contrato de reserva con Google. Usa la marca --reserved al crear tus recursos.

  • Spot: se dirige a la cuota interrumpible mediante máquinas virtuales de acceso puntual. Es posible que tus recursos se expropien para dejar espacio a las solicitudes de un trabajo de mayor prioridad. Usa la marca --spot al crear tus recursos.

  • Bajo demanda: se dirige a la cuota bajo demanda, que no necesita una reserva y no se puede anticipar. La solicitud de TPU se pondrá en cola en una cola de cuota bajo demanda ofrecida por TPU de Cloud. No se garantiza la disponibilidad de los recursos. Seleccionado de forma predeterminada, no se necesitan marcas.

Empezar

  1. Configura tu entorno de TPU de Cloud.

  2. In the Google Cloud console, activate Cloud Shell.

    Activate Cloud Shell

    At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

    3. Para usar Multislice, tus recursos de TPU deben gestionarse como recursos en cola.

    Ejemplo introductorio

    En este tutorial 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 ha diseñado para entrenar de forma eficiente en las TPU de Cloud.

    El código de shardings.py se ha diseñado para ayudarte a empezar a experimentar con diferentes opciones de paralelismo. Por ejemplo, el paralelismo de datos, el paralelismo de datos totalmente fragmentado (FSDP) y el paralelismo de tensores. El código se adapta a entornos de un solo segmento y de varios segmentos.

    Paralelismo de ICI

    ICI hace referencia a la interconexión de alta velocidad que conecta las TPUs de un mismo slice. La fragmentación de ICI se corresponde con la fragmentación dentro de un segmento. shardings.py proporciona tres parámetros de paralelismo de ICI:

    • ici_data_parallelism
    • ici_fsdp_parallelism
    • ici_tensor_parallelism

    Los valores que especifique para estos parámetros determinan el número de particiones de cada método de paralelismo.

    Estas entradas deben estar restringidas de forma que ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism sea igual al número de chips del segmento.

    En la siguiente tabla se muestran ejemplos de entradas de usuario para el paralelismo de ICI de los cuatro chips disponibles en la versión 4-8:

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

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

    En este ejemplo, se da por hecho que sabes cómo ejecutar código en un solo segmento de TPU, como en Ejecutar un cálculo en una VM de TPU de Cloud con JAX. En este ejemplo se muestra cómo ejecutar shardings.py en una sola porción.

    1. Configura el entorno:

      $ gcloud auth login
$ export QR_ID=your-queued-resource-id
$ export TPU_NAME=your-tpu-name
$ export PROJECT=your-project-name
$ export ZONE=us-central1-a
$ export NETWORK_NAME=your-network-name
$ export SUBNETWORK_NAME=your-subnetwork-name
$ export RUNTIME_VERSION=v2-alpha-tpuv5-lite
$ export ACCELERATOR_TYPE=v5litepod-16
$ export EXAMPLE_TAG_1=your-tag-1
$ export EXAMPLE_TAG_2=your-tag-2
$ export SLICE_COUNT=4
$ export STARTUP_SCRIPT='#!/bin/bash\n'

      Descripciones de las variables

      • QR_ID: ID asignado por el usuario del recurso en cola.
      • TPU_NAME: nombre asignado por el usuario de tu TPU.
      • PROJECT: Google Cloud nombre del proyecto
      • ZONE: especifica la zona en la que se crearán los recursos.
      • NETWORK_NAME: nombre de la red de VPC.
      • SUBNETWORK_NAME: nombre de la subred de la red de VPC
      • RUNTIME_VERSION: la versión de software de TPU de Cloud.
      • ACCELERATOR_TYPE: el tipo de acelerador especifica la versión y el tamaño de la TPU de Cloud que quieres crear.
      • EXAMPLE_TAG_1, EXAMPLE_TAG_2 …: etiquetas que se usan para identificar orígenes o destinos válidos en los cortafuegos de la red.
      • SLICE_COUNT: número de porciones. Se limita a un máximo de 256 porciones.
      • STARTUP_SCRIPT: si especificas una secuencia de comandos de inicio, se ejecutará cuando se aprovisione o se reinicie la porción de TPU.

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

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

    3. Aprovisiona tus TPUs:

      gcloud

      $ gcloud compute tpus queued-resources \
    create ${QR_ID} \
    --accelerator-type=${ACCELERATOR_TYPE} \
    --runtime-version=${RUNTIME_VERSION} \
    --node-id=${TPU_NAME} \
    --zone=${ZONE} \
    [--reserved |--spot]

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

      Consola

      1. En la Google Cloud consola, ve a la página TPUs:

        Ir a TPUs

      2. Haz clic en Crear TPU.

      3. En el campo Nombre, introduce un nombre para tu TPU.

      4. En el cuadro Zona, selecciona la zona en la que quieras crear la TPU.