TPU v6e

En este documento, se describen la arquitectura y los parámetros de configuración compatibles de la Cloud TPU v6e (Trillium).

Trillium es el acelerador de IA de última generación de Cloud TPU. En todas las plataformas técnicas, como la API y los registros, y en todo este documento, se hará referencia a Trillium como v6e.

Con una capacidad de 256 chips por Pod, v6e comparte muchas similitudes con v5e. Este sistema está optimizado para ser el producto de mayor valor para el entrenamiento, el ajuste y la entrega de transformadores, modelos de texto a imagen y redes neuronales convolucionales (CNN).

Arquitectura del sistema

Cada chip v6e contiene un TensorCore. Cada TensorCore tiene 2 unidades de multiplicación de matrices (MXU), una unidad vectorial y una unidad escalar. En la siguiente tabla, se muestran las especificaciones clave y sus valores para la TPU v6e en comparación con la TPU v5e.

Especificación v5e v6e
Rendimiento/costo total de propiedad (TCO) (esperado) 0.65x 1
Procesamiento máximo por chip (bf16) 197 TFLOPS 918 TFLOPS
Procesamiento máximo por chip (Int8) 393 TOPS 1,836 TOPS
Capacidad de HBM por chip 16 GB 32 GB
Ancho de banda de HBM por chip 800 Gbps 1,600 Gbps
Ancho de banda de interconexión entre chips (ICI) 1,600 Gbps 3,200 Gbps
Puertos de ICI por chip 4 4
DRAM por host 512 GiB 1,536 GiB
Chips por host 8 8
Tamaño del pod de TPU 256 chips 256 chips
Topología de interconexión Toro 2D Toro 2D
Cómputo máximo en BF16 por Pod 50.63 PFLOPS 234.9 PFLOPS
Ancho de banda total reducido por Pod 51.2 TB/s 102.4 TB/s
Ancho de banda de bisección por Pod 1.6 TB/s 3.2 TB/s
Configuración de NIC por host NIC de 2 x 100 Gbps NIC de 4 x 200 Gbps
Ancho de banda de la red del centro de datos por Pod 6.4 Tbps 25.6 Tbps
Características especiales - SparseCore

Parámetros de configuración admitidos

En la siguiente tabla, se muestran las formas de porción 2D que se admiten para v6e:

Topología Chips TPU Hosts VMs Tipo de acelerador (API de TPU) Tipo de máquina (API de GKE) Alcance
1x1 1 1/8 1 v6e-1 ct6e-standard-1t Subhost
2x2 4 1/2 1 v6e-4 ct6e-standard-4t Subhost
2x4 8 1 1 v6e-8 ct6e-standard-8t Único host
2x4 8 1 2 - ct6e-standard-4t Único host
4x4 16 2 4 v6e-16 ct6e-standard-4t Varios hosts
4x8 32 4 8 v6e-32 ct6e-standard-4t Varios hosts
8x8 64 8 16 v6e-64 ct6e-standard-4t Varios hosts
8x16 128 16 32 v6e-128 ct6e-standard-4t Varios hosts
16x16 256 32 64 v6e-256 ct6e-standard-4t Varios hosts

Los segmentos con 8 chips (v6e-8) conectados a una sola VM están optimizados para la inferencia, lo que permite que los 8 chips se usen en una sola carga de trabajo de entrega. Puedes realizar inferencias con varios hosts usando rutas de aprendizaje en Cloud. Para obtener más información, consulta Cómo realizar inferencias en varios hosts con rutas de aprendizaje.

Para obtener información sobre la cantidad de VMs para cada topología, consulta Tipos de VMs.

Tipos de VMs

Cada VM de TPU v6e puede contener 1, 4 o incluso 8 chips. Las porciones de 4 chips y más pequeñas tienen el mismo nodo de acceso a la memoria no uniforme (NUMA). Para obtener más información sobre los nodos de NUMA, consulta Acceso no uniforme a la memoria en Wikipedia.

Diagrama de un host de v6e

Las porciones de v6e se crean con VMs de medio host, cada una con 4 chips TPU. Existen dos excepciones a esta regla:

  • v6e-1: Es una VM con un solo chip, diseñada principalmente para pruebas.
  • v6e-8: Es una VM de host completo que se optimizó para un caso de uso de inferencia con los 8 chips conectados a una sola VM.

En la siguiente tabla, se muestra una comparación de los tipos de VM de TPU v6e:

Tipo de VM Cantidad de CPU virtuales por VM RAM (GB) por VM Cantidad de nodos de NUMA por VM
VM con 1 chip 44 176 1
VM con 4 chips 180 720 1
VM con 8 chips 180 1440 2

Especifica la configuración de v6e

Cuando asignas una porción de TPU v6e con la API de TPU, especificas su tamaño y forma con el parámetro AcceleratorType.

Si usas GKE, usa la marca --machine-type para especificar un tipo de máquina que admita la TPU que deseas usar. Para obtener más información, consulta Planifica las TPU en GKE en la documentación de GKE.

Utiliza AcceleratorType

Cuando asignas recursos TPU, usas AcceleratorType para especificar la cantidad de TensorCores en una porción. El valor que especifiques para AcceleratorType es una cadena con el formato v$VERSION-$TENSORCORE_COUNT. Por ejemplo, v6e-8 especifica una porción de TPU v6e con 8 TensorCores.

En el siguiente ejemplo, se muestra cómo crear una porción de TPU v6e con 32 TensorCores usando AcceleratorType:

gcloud

    $ gcloud compute tpus tpu-vm create tpu-name \
        --zone=zone \
        --accelerator-type=v6e-32 \
        --version=v2-alpha-tpuv6e

Consola

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

    Ir a TPU

  2. Haz clic en Crear una TPU.

  3. En el campo Nombre, ingresa el nombre de tu TPU.

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

  5. En el cuadro Tipo de TPU, selecciona v6e-32.

  6. En el cuadro Versión de software de TPU, selecciona v2-alpha-tpuv6e. Cuando creas una VM de Cloud TPU, la versión de software de TPU especifica la versión del entorno de ejecución de TPU que se instalará. Para obtener más información, consulta Imágenes de VM de TPU.

  7. Haz clic en el botón para activar o desactivar la opción Habilitar las colas.

  8. En el campo Nombre del recurso en cola, ingresa un nombre para tu solicitud de recurso en cola.

  9. Haz clic en Crear.

¿Qué sigue?