TPU v5e

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

TPU v5e admite el entrenamiento de uno y varios hosts, y la inferencia de uno solo. La inferencia de varios hosts se admite mediante Sax. Para obtener más información, consulta Entrega de modelos de lenguaje grande.

Arquitectura del sistema

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

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

Diagrama de un chip v5e

En la siguiente tabla, se muestran las especificaciones de los chips de claves y sus valores para la v5e.

Especificaciones de los chips de claves Valores de v5e
Procesamiento máximo por chip (bf16) 197 TFLOP
Capacidad y ancho de banda de HBM2 16 GB, 819 GBps
Ancho de banda de interconexión entre chips 1,600 Gbps

En la siguiente tabla, se muestran las especificaciones de Pods y sus valores para v5e.

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

Configuraciones

Cloud TPU v5e es un producto combinado de inferencia (entrega) y entrenamiento. Para diferenciar entre un entorno de entrenamiento y uno de inferencia, usa las marcas AcceleratorType o AcceleratorConfig con la API de TPU o la marca --machine-type cuando crees un grupo de nodos de GKE.

Los trabajos de entrenamiento están optimizados para la capacidad de procesamiento y la disponibilidad, mientras que los trabajos de publicación 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.

Usa AcceleratorType para especificar la cantidad de tensores de red que deseas usar. Debes especificar el AcceleratorType cuando creas una TPU con el gcloud CLI o la consola de Google Cloud. El valor que especificas para AcceleratorType es una cadena con el formato v$VERSION_NUMBER-$CHIP_COUNT.

También puedes usar AcceleratorConfig para especificar la cantidad de TensorCores que quieres usar. Sin embargo, como no hay variantes de topología 2D personalizadas para la TPU v5e, no hay diferencia entre usar AcceleratorConfig y AcceleratorType.

Para configurar una TPU v5e con AcceleratorConfig, usa --version y el --topology. Establece --version en la versión de TPU que deseas usar. --topology a la disposición física de los chips TPU en la porción. El valor que especificas para AcceleratorConfig es una cadena con el formato AxB, en la que A y B son los recuentos de chips en cada dirección.

Las siguientes formas de porciones en 2D son compatibles con la v5e:

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

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

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

Comunicación de nodos de NUMA

Tipos de Cloud TPU v5e para entregar

La entrega para un solo host es compatible con hasta 8 chips v5e. Se admiten las siguientes configuraciones: rebanadas de 1 × 1, 2 × 2 y 2 × 4. Cada porción tiene 1, 4 y 8 chips respectivamente.

Parámetros de configuración de TPU v5e que admiten la publicación: 1 × 1, 2 × 2 y 2 × 4.

Si quieres aprovisionar TPU para un trabajo de entrega, usa uno de los siguientes tipos de acelerador en tu solicitud de creación de TPU en 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

La entrega en más de 8 chips v5e, también llamada entrega de varios hosts, es compatible con Sax. Para obtener más información, consulta Entrega de modelos de lenguaje grande.

Tipos de Cloud TPU v5e para el entrenamiento

El entrenamiento es compatible con hasta 256 chips.

Si quieres aprovisionar TPU para un trabajo de entrenamiento v5e, usa uno de los siguientes aceleradores en tu solicitud de creación de TPU de CLI o 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

Comparación de los tipos de VM de TPU v5e:

Tipo de VM n2d-48-24-v5lite-tpu n2d-192-112-v5lite-tpu n2d-384-224-v5lite-tpu
Cantidad de chips v5e 1 4 8
Cantidad de CPUs virtuales 24 112 224
RAM (GB) 48 192 384
Cantidad de nodos de NUMA 1 1 2
Se aplica a v5litepod-1 v5litepod-4 v5litepod-8
Interrupción Alta Medio Low (Baja)

Para liberar espacio para las cargas de trabajo que requieren más chips, los programadores pueden interrumpir las VMs con menos chips. Por lo tanto, es probable que las VMs de 8 chips interrumpan a las de 1 y 4 chips.