TPU v5e

En este documento, se describen la arquitectura y las configuraciones compatibles de la Cloud TPU v5e.

La TPU v5e admite el entrenamiento de host único y de varios hosts, y la inferencia de host único. La inferencia de varios hosts se admite con Sax. Para obtener más información, consulta Inferencias en Cloud TPU.

Arquitectura del sistema

Cada chip v5e contiene un TensorCore. Cada TensorCore tiene cuatro unidades de multiplicación de matrices (MXU), una unidad vectorial y una unidad escalar.

En el siguiente diagrama, se ilustra un chip de TPU v5e.

Diagrama de un chip v5e

En la siguiente tabla, se muestran las especificaciones clave del chip y sus valores para la versión 5e.

Especificaciones clave del chip Valores de v5e
Procesamiento máximo por chip (bf16) 197 TFLOPS
Capacidad y ancho de banda de HBM2 16 GB, 819 GBps
BW de interconexión entre chips 1600 Gbps

En la siguiente tabla, se muestran las especificaciones de los Pods y sus valores para la versión 5e.

Especificaciones clave del Pod Valores de v5e
Tamaño del pod de TPU 256 chips
Topología de interconexión Toroide 2D
Procesamiento máximo por Pod 100 PetaOps (Int8)
Ancho de banda de reducción total por Pod 51.2 TB/s
Ancho de banda de bisección por Pod 1.6 TB/s
Ancho de banda de la red del centro de datos por Pod 6.4 Tbps

Configuraciones

La Cloud TPU v5e es un producto combinado de entrenamiento e inferencia (publicación). Para diferenciar un entorno de entrenamiento de uno de inferencia, usa el parámetro AcceleratorType con la API de TPU o la marca --machine-type cuando crees un grupo de nodos de Google Kubernetes Engine (GKE).

Los trabajos de entrenamiento están optimizados para la capacidad de procesamiento y la disponibilidad, mientras que los trabajos de entrega están optimizados para la latencia. Un trabajo de entrenamiento en TPU aprovisionadas para la entrega podría tener una disponibilidad más baja y, de manera similar, un trabajo de entrega ejecutado en TPU aprovisionadas para el entrenamiento podría tener una latencia más alta.

Usas AcceleratorType para especificar la cantidad de TensorCores que deseas usar. Especificas AcceleratorType cuando creas una TPU con gcloud CLI o la Google Cloud consola. El valor que especifiques para AcceleratorType es una cadena con el formato v$VERSION_NUMBER-$CHIP_COUNT.

Se admiten las siguientes formas de corte 2D para la versión 5e:

Topología Cantidad de chips TPU Cantidad de hosts
1x1 1 1/8
2x2 4 1/2
2x4 8 1
4x4 16 2
4x8 32 4
8x8 64 8
8x16 128 16
16x16 256 32

Tipos de VM

Cada VM de TPU en una porción de TPU v5e contiene 1, 4 u 8 chips. En las porciones de 4 chips o menos, todos los chips de TPU comparten el mismo nodo de acceso a la memoria no uniforme (NUMA).

En el caso de las VMs de TPU v5e con 8 chips, la comunicación entre la CPU y la TPU será más eficiente dentro de las particiones NUMA. Por ejemplo, en la siguiente figura, la comunicación CPU0-Chip0 será más rápida que la comunicación CPU0-Chip4.

Comunicación de nodos de NUMA

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

Tipo de VM Tipo de máquina (API de GKE) Cantidad de CPU virtuales por VM RAM (GB) por VM Cantidad de nodos de NUMA por VM
VM con 1 chip ct5lp-hightpu-1t 24 48 1
VM con 4 chips ct5lp-hightpu-4t 112 192 1
VM de 8 chips ct5lp-hightpu-8t 224 384 2

Tipos de Cloud TPU v5e para la publicación

Se admite la publicación en un solo host para hasta 8 chips de v5e. Se admiten las siguientes configuraciones: segmentos de 1 x 1, 2 x 2 y 2 x 4. Cada porción tiene 1, 4 y 8 chips, respectivamente.

Configuraciones de TPU v5e que admiten la entrega: 1x1, 2x2 y 2x4.

Para aprovisionar TPUs para un trabajo de procesamiento, usa uno de los siguientes tipos de aceleradores en tu solicitud de creación de TPU de la CLI o la API:

AcceleratorType (API de TPU) Tipo de máquina (API de GKE)
v5litepod-1 ct5lp-hightpu-1t
v5litepod-4 ct5lp-hightpu-4t
v5litepod-8 ct5lp-hightpu-8t

El siguiente comando crea un segmento de TPU v5e con 8 chips v5e para la inferencia:

$ gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=us-central1-a \
    --accelerator-type=v5litepod-8 \
    --version=v2-alpha-tpuv5-lite

Para obtener más información sobre la administración de TPU, consulta Administra TPUs. Para obtener más información sobre la arquitectura del sistema de Cloud TPU, consulta Arquitectura del sistema.

La entrega en más de 8 chips v5e, también llamada entrega con varios hosts, se admite con Sax. Para obtener más información, consulta Inferencias en Cloud TPU.

Tipos de Cloud TPU v5e para el entrenamiento

Se admite el entrenamiento para hasta 256 chips.

Para aprovisionar TPUs para un trabajo de entrenamiento con v5e, usa uno de los siguientes tipos de aceleradores en tu solicitud de creación de TPU de la CLI o la API:

AcceleratorType (API de TPU) Tipo de máquina (API de GKE) Topología
v5litepod-16 ct5lp-hightpu-4t 4x4
v5litepod-32 ct5lp-hightpu-4t 4x8
v5litepod-64 ct5lp-hightpu-4t 8x8
v5litepod-128 ct5lp-hightpu-4t 8x16
v5litepod-256 ct5lp-hightpu-4t 16x16

El siguiente comando crea un segmento de TPU v5e con 256 chips v5e para el entrenamiento:

$ gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=us-east5-a \
    --accelerator-type=v5litepod-256 \
    --version=v2-alpha-tpuv5-lite

Para obtener más información sobre la administración de TPU, consulta Administra TPUs. Para obtener más información sobre la arquitectura del sistema de Cloud TPU, consulta Arquitectura del sistema.