TPU v4

En este documento se describen la arquitectura y las configuraciones admitidas de la versión 4 de TPU de Cloud.

Nota: Todas las solicitudes de cuota en us-central2-b (actualmente, solo TPU v4) deben ser aprobadas manualmente por Google. No se concede ninguna cuota predeterminada. Si necesitas cuota en us-central2-b, consulta Ver y solicitar cuota adicional.

Arquitectura del sistema

Cada chip de la versión 4 de TPU contiene dos TensorCores. Cada Tensor Core tiene cuatro unidades de multiplicación de matrices (MXUs), una unidad vectorial y una unidad escalar. En la tabla siguiente se muestran las especificaciones clave de un pod de TPU v4.

Especificaciones principales	Valores de los pods de la versión 4
Rendimiento máximo por chip	275 teraflops (bf16 o int8)
Capacidad y ancho de banda de HBM2	32 GiB, 1200 GB/s
Potencia mínima, media y máxima medidas	90/170/192 W
Tamaño del pod de TPUs	4096 chips
Topología de interconexión	Malla 3D
Computación máxima por pod	1,1 exaflops (bf16 o int8)
Ancho de banda de All-reduce por pod	1,1 PB/s
Ancho de banda de bisección por pod	24 TB/s

En el siguiente diagrama se muestra un chip TPU v4.

Para obtener más información sobre los detalles de la arquitectura y las características de rendimiento de la TPU v4, consulta TPU v4: An Optically Reconfigurable Supercomputer for Machine Learning with Hardware Support for Embeddings (TPU v4: un superordenador ópticamente reconfigurable para el aprendizaje automático con asistencia de hardware para las inserciones).

Malla 3D y toro 3D

Las TPUs v4 tienen una conexión directa con los chips vecinos más cercanos en 3 dimensiones, lo que da como resultado una malla 3D de conexiones de red. Las conexiones se pueden configurar como un toroide 3D en las secciones en las que la topología, AxBxC, es 2A=B=C o 2A=2B=C, donde cada dimensión es un múltiplo de 4. Por ejemplo, 4x4x8, 4x8x8 o 12x12x24. En general, el rendimiento de una configuración de toroide 3D será mejor que el de una configuración de malla 3D. Para obtener más información, consulta el artículo sobre topologías de toros retorcidos.

Ventajas de rendimiento de la versión 4 de TPU con respecto a la versión 3

En esta sección se muestra una forma eficiente de usar la memoria para ejecutar un script de entrenamiento de ejemplo en la versión 4 de TPU, así como las mejoras de rendimiento de la versión 4 de TPU en comparación con la versión 3.

Sistema de memoria

El acceso a memoria no uniforme (NUMA) es una arquitectura de memoria de ordenador para máquinas que tienen varias CPUs. Cada CPU tiene acceso directo a un bloque de memoria de alta velocidad. Una CPU y su memoria se denominan nodo NUMA. Los nodos NUMA están conectados a otros nodos NUMA que son adyacentes entre sí. Una CPU de un nodo NUMA puede acceder a la memoria de otro nodo NUMA, pero este acceso es más lento que el acceso a la memoria de un nodo NUMA.

El software que se ejecuta en una máquina con varias CPUs puede colocar los datos que necesita una CPU en su nodo NUMA, lo que aumenta el rendimiento de la memoria. Para obtener más información sobre NUMA, consulta Acceso a memoria no uniforme en Wikipedia.

Puedes aprovechar las ventajas de la localidad NUMA vinculando tu secuencia de comandos de entrenamiento al nodo NUMA 0.

Para habilitar la vinculación de nodos NUMA, haz lo siguiente:

1. Crea una TPU v4.

2. Conéctate a la TPU mediante SSH.

3. Instala la herramienta de línea de comandos numactl. numactl te permite ejecutar procesos con una política de programación o de colocación de memoria NUMA específica.

    $ sudo apt-get update
    $ sudo apt-get install numactl

4. Vincula el código de tu secuencia de comandos al nodo NUMA 0. Sustituye your-training-script por la ruta a tu secuencia de comandos de entrenamiento.

    $ numactl --cpunodebind=0 python3 your-training-script

Habilita la vinculación de nodos NUMA si se cumplen las siguientes condiciones:

• Si tu carga de trabajo depende en gran medida de las cargas de trabajo de la CPU (por ejemplo, clasificación de imágenes o cargas de trabajo de recomendaciones), independientemente del framework.
• Si usas una versión del entorno de ejecución de TPU sin el sufijo -pod (por ejemplo, tpu-vm-tf-2.10.0-v4).

Otras diferencias del sistema de memoria:

• Los chips de TPU v4 tienen un espacio de memoria HBM unificado de 32 GiB en todo el chip, lo que permite una mejor coordinación entre los dos TensorCores del chip.
• Rendimiento de HBM mejorado con los estándares y las velocidades de memoria más recientes.
• Perfil de rendimiento de DMA mejorado con compatibilidad integrada para el avance de alto rendimiento con granularidades de 512 B.

TensorCores

• El doble de unidades de ejecución de matriz y una frecuencia de reloj más alta, lo que da como resultado 275 TFLOPS como máximo.
• Ancho de banda de transposición y permutación 2x.
• Modelo de acceso a memoria de carga y almacenamiento para memoria común (Cmem).
• Ancho de banda de carga de peso de MXU más rápido y compatibilidad con el modo de 8 bits para permitir tamaños de lote más pequeños y una latencia de inferencia mejorada.

Interconexión entre chips

Seis enlaces de interconexión por chip para habilitar topologías de red con diámetros de red más pequeños.

Otro

• Interfaz PCIE gen3 x16 para el host (conexión directa).
• Modelo de seguridad mejorado.
• Mejora de la eficiencia energética.

Configuraciones

Un pod de TPU v4 está compuesto por 4096 chips interconectados con enlaces de alta velocidad reconfigurables. La red flexible de la TPU v4 te permite conectar los chips de una porción del mismo tamaño de varias formas. Cuando creas una porción de TPU, especificas la versión de TPU y el número de recursos de TPU que necesitas. Cuando creas un segmento de TPU v4, puedes especificar su tipo y tamaño de dos formas: AcceleratorType y AccleratorConfig.

Nota: AcceleratorType y AccleratorConfig son mutuamente excluyentes. Solo puedes usar uno de ellos al crear una TPU.

Estás usando AcceleratorType

Usa AcceleratorType cuando no especifiques una topología. Para configurar las TPUs v4 con AcceleratorType, usa la marca --accelerator-type al crear tu sector de TPU. Asigna a --accelerator-type una cadena que contenga la versión de la TPU y el número de Tensor Cores que quieras usar. Por ejemplo, para crear un segmento de v4 con 32 Tensor Cores, usarías --accelerator-type=v4-32.

Usa el comando gcloud compute tpus tpu-vm create para crear un segmento de TPU v4 con 512 Tensor Cores mediante la marca --accelerator-type:

  $ gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=us-central2-b \
    --accelerator-type=v4-512 \
    --version=tpu-ubuntu2204-base

El número que aparece después de la versión de la TPU (v4) indica el número de Tensor Cores. Una TPU v4 tiene dos Tensor Cores, por lo que el número de chips de TPU sería 512/2 = 256.

Para obtener más información sobre cómo gestionar las TPUs, consulta Gestionar TPUs. Para obtener más información sobre la arquitectura del sistema de las TPU de Cloud, consulta Arquitectura del sistema.

Estás usando AcceleratorConfig

Usa AcceleratorConfig cuando quieras personalizar la topología física de tu segmento de TPU. Por lo general, es necesario para optimizar el rendimiento con segmentos de más de 256 chips.

Para configurar las TPUs v4 con AcceleratorConfig, usa las marcas --type y --topology. Asigna a --type la versión de TPU que quieras usar y a --topology la disposición física de los chips de TPU en el slice.

Para especificar una topología de TPU, se usa una tupla de 3 elementos, AxBxC, donde A<=B<=C y A, B y C son todos <= 4 o todos múltiplos enteros de 4. Los valores A, B y C son el número de chips de cada una de las tres dimensiones. Por ejemplo, para crear una porción de la versión 4 con 16 chips, debes definir --type=v4 y --topology=2x2x4.

Usa el gcloud compute tpus tpu-vm create comando para crear un segmento de TPU v4 con 128 chips de TPU dispuestos en una matriz de 4x4x8:

  $ gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=us-central2-b \
    --type=v4 \
    --topology=4x4x8 \
    --version=tpu-ubuntu2204-base

Las topologías en las que 2A=B=C o 2A=2B=C también tienen variantes de topología optimizadas para la comunicación de todos con todos, como 4×4×8, 8×8×16 y 12×12×24. Se conocen como topologías de toros retorcidos.

En las siguientes ilustraciones se muestran algunas topologías habituales de TPU v4.

Las porciones más grandes se pueden crear a partir de uno o varios "cubos" de chips de 4x4x4.

Para obtener más información sobre cómo gestionar las TPUs, consulta Gestionar TPUs. Para obtener más información sobre la arquitectura del sistema de las TPU de Cloud, consulta Arquitectura del sistema.

Topologías de toros retorcidos

Algunas formas de sectores de toro en 3D de la versión 4 tienen la opción de usar lo que se conoce como topología de toro retorcido. Por ejemplo, dos cubos de la versión 4 se pueden organizar como un sector de 4x4x8 o como 4x4x8_twisted. Las topologías retorcidas ofrecen un ancho de banda de bisección significativamente mayor. Por ejemplo, una porción con la topología 4x4x8_twisted proporciona un aumento teórico del 70 % en el ancho de banda de bisección con respecto a una porción 4x4x8 sin torsión. El aumento del ancho de banda de bisección es útil para las cargas de trabajo que usan patrones de comunicación globales. Las topologías retorcidas pueden mejorar el rendimiento de la mayoría de los modelos, y las cargas de trabajo de inserción de TPU de gran tamaño son las que más se benefician.

En las cargas de trabajo que usan el paralelismo de datos como única estrategia de paralelismo, las topologías retorcidas pueden ofrecer un rendimiento ligeramente mejor. En el caso de los LLMs, el rendimiento con una topología retorcida puede variar en función del tipo de paralelismo (DP, MP, etc.). Lo más recomendable es entrenar tu LLM con y sin una topología retorcida para determinar cuál ofrece el mejor rendimiento para tu modelo. En algunos experimentos con el modelo FSDP MaxText se han observado mejoras de 1-2 MFU con una topología retorcida.

La principal ventaja de las topologías retorcidas es que transforman una topología de toro asimétrica (por ejemplo, 4×4×8) en una topología simétrica estrechamente relacionada. La topología simétrica tiene muchas ventajas:

• Balanceo de carga mejorado
• Mayor ancho de banda de bisección
• Rutas de paquetes más cortas

En última instancia, estas ventajas se traducen en un mejor rendimiento para muchos patrones de comunicación globales.

El software de TPU admite toros retorcidos en las porciones en las que el tamaño de cada dimensión es igual o el doble del tamaño de la dimensión más pequeña. Por ejemplo, 4x4x8, 4×8×8 o 12x12x24.

Por ejemplo, considera esta topología de toro de 4×2 con TPUs etiquetadas con sus coordenadas (X,Y) en el slice:

Los bordes de este gráfico de topología se muestran como bordes no dirigidos para mayor claridad. En la práctica, cada arista es una conexión bidireccional entre las TPUs. Nos referimos a los bordes entre un lado de esta cuadrícula y el lado opuesto como bordes envolventes, tal como se indica en el diagrama.

Al retorcer esta topología, obtenemos una topología de toro retorcido 4×2 completamente simétrica:

Lo único que ha cambiado entre este diagrama y el anterior son los bordes envolventes de Y. En lugar de conectarse a otra TPU con la misma coordenada X, se han desplazado para conectarse a la TPU con la coordenada X+2 mod 4.

La misma idea se generaliza a diferentes tamaños de dimensión y a diferentes números de dimensiones. La red resultante es simétrica, siempre que cada dimensión sea igual o el doble de la dimensión más pequeña.

Consulta Uso de AcceleratorConfig para obtener información sobre cómo especificar una configuración de toroide retorcido al crear una TPU de Cloud.

En la siguiente tabla se muestran las topologías trenzadas admitidas y el aumento teórico del ancho de banda de bisección con ellas en comparación con las topologías sin trenzar.

Topología retorcida	Aumento teórico de la bisección del ancho de banda en comparación con un toro no retorcido
4×4×8_twisted	~70%
8x8x16_twisted
12×12×24_twisted
4×8×8_twisted	~40%
8×16×16_twisted

Nota: La topología retorcida no se usará automáticamente a menos que se especifique el sufijo _twisted con la marca --topology. Por ejemplo, para especificar una topología retorcida de 4x4x8, se debe especificar la marca topology de la siguiente manera: 4x4x8_twisted.

Variantes de topología de TPU v4

Algunas topologías que contienen el mismo número de chips se pueden organizar de diferentes formas. Por ejemplo, un slice de TPU con 512 chips (1024 TensorCores) se puede configurar con las siguientes topologías: 4x4x32, 4x8x16 u 8x8x8. Una porción de TPU con 2048 chips (4096 TensorCores) ofrece aún más opciones de topología: 4x4x128, 4x8x64, 4x16x32 y 8x16x16.

La topología predeterminada asociada a un número de chips determinado es la que más se parece a un cubo. Esta forma es probablemente la mejor opción para el entrenamiento de aprendizaje automático paralelo de datos. Otras topologías pueden ser útiles para cargas de trabajo con varios tipos de paralelismo (por ejemplo, paralelismo de modelos y de datos, o partición espacial de una simulación). Estas cargas de trabajo funcionan mejor si la topología coincide con el paralelismo utilizado. Por ejemplo, si se aplica un paralelismo de modelo de 4 vías en la dimensión X y un paralelismo de datos de 256 vías en las dimensiones Y y Z, se obtiene una topología de 4x16x16.

Los modelos con varias dimensiones de paralelismo funcionan mejor si sus dimensiones de paralelismo se asignan a las dimensiones de la topología de la TPU. Normalmente, se trata de modelos de lenguaje extensos (LLMs) paralelos de datos y modelos. Por ejemplo, en un segmento de TPU v4 con una topología de 8x16x16, las dimensiones de la topología de TPU son 8, 16 y 16. Es más eficiente usar el paralelismo de modelos de 8 o 16 vías (asignado a una de las dimensiones de la topología física de la TPU). Un paralelismo de modelo de 4 vías no sería óptimo con esta topología, ya que no se alinea con ninguna de las dimensiones de la topología de TPU, pero sí lo sería con una topología de 4x16x32 en el mismo número de chips.

Las configuraciones de TPU v4 se dividen en dos grupos: las que tienen topologías de menos de 64 chips (topologías pequeñas) y las que tienen topologías de más de 64 chips (topologías grandes).

Topologías pequeñas de v4

Cloud TPU admite los siguientes sectores de TPU v4 con menos de 64 chips, un cubo de 4x4x4. Puedes crear estas pequeñas topologías v4 usando su nombre basado en Tensor Core (por ejemplo, v4-32) o su topología (por ejemplo, 2x2x4):

Nombre (basado en el número de Tensor Cores)	Número de fichas	Topología
v4-8	4	2x2x1
v4-16	8	2x2x2
v4-32	16	2x2x4
v4-64	32	2x4x4

Topologías grandes de versión 4

Las slices de TPU v4 están disponibles en incrementos de 64 chips, con formas que son múltiplos de 4 en las tres dimensiones. Las dimensiones deben estar en orden ascendente. En la siguiente tabla se muestran varios ejemplos. Algunas de estas topologías son topologías personalizadas que solo se pueden crear con las marcas --type y --topology, ya que hay más de una forma de organizar los chips.

Nombre (basado en el número de Tensor Cores)	Número de fichas	Topología
v4-128	64	4x4x4
v4-256	128	4x4x8
v4-512	256	4x8x8
Topología personalizada: debe usar las marcas --type y --topology	256	4x4x16
v4-1024	512	8x8x8
v4-1536	17.	8x8x12
v4-2048	1024	8x8x16
Topología personalizada: debe usar las marcas --type y --topology	1024	4x16x16
v4-4096	2048	8x16x16
…	…	…