TPU v4

En este documento, se describen la arquitectura y los parámetros de configuración admitidos de Cloud TPU v4.

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

Arquitectura del sistema

Cada chip TPU v4 contiene dos TensorCores. Cada TensorCore tiene cuatro unidades de multiplicación de matrices (MXU), una unidad vectorial y una unidad escalar. En la siguiente tabla, se muestran las especificaciones clave del pod de TPU v4.

Especificaciones clave	Valores de Pods v4
Procesamiento máximo por chip	275 teraflops (bf16 o int8)
Capacidad y ancho de banda de HBM2	32 GiB, 1,200 GBps
Potencia mínima, media y máxima medida	90/170/192 W
Tamaño del pod de TPU	4,096 chips
Topología de interconexión	Malla 3D
Procesamiento máximo por Pod	1.1 exaflops (bf16 o int8)
Ancho de banda total reducido por Pod	1.1 PB/s
Ancho de banda de bisección por Pod	24 TB/s

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

Para obtener más información sobre los detalles de la arquitectura y las características de rendimiento de TPU v4, consulta TPU v4: Una supercomputadora reconfigurable ópticamente orientada al aprendizaje automático y la compatibilidad a nivel de hardware para las embeddings.

Malla y toro 3D

Las TPU v4 tienen una conexión directa con los chips vecinos más cercanos en 3 dimensiones, lo que genera una malla 3D de conexiones de red. Las conexiones se pueden configurar como un toro 3D en porciones en las que la topología, AxBxC, es 2A=B=C o 2A=2B=C, en 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 toro 3D será mejor que el de una de malla 3D. Para obtener más información, consulta Topologías de toros retorcidos.

Beneficios de rendimiento de TPU v4 en comparación con v3

En esta sección, se muestra una forma eficiente en el uso de memoria para ejecutar una secuencia de comandos de entrenamiento de ejemplo en TPU v4, así como las mejoras de rendimiento de TPU v4 en comparación con TPU v3.

Sistema de memoria

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

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

Puedes aprovechar los beneficios de la localidad de NUMA a través de la vinculación de tu secuencia de comandos de entrenamiento al nodo de NUMA 0.

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

1. Crea una TPU v4.

2. Conéctate a la TPU con SSH.

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

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

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

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

Habilita la vinculación de nodos de NUMA en los siguientes casos:

• Si tu carga de trabajo depende en gran medida de las cargas de trabajo de la CPU (por ejemplo, clasificación de imágenes y 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 TPU v4 tienen un espacio de memoria de HBM unificado de 32 GiB en todo el chip, lo que permite una mejor coordinación entre los dos TensorCores integrados en el chip
• Se mejoró el rendimiento del HBM con los estándares y las velocidades de memoria más recientes
• Se mejoró el perfil de rendimiento de la DMA con compatibilidad integrada para el avance de alto rendimiento en granularidades de 512B

TensorCores

• El doble de unidades MXU y una mayor frecuencia de reloj, que ofrecen un máximo de 275 TFLOPS
• El doble de ancho de banda para operaciones de transposición y permutación
• Modelo de acceso a la memoria de carga y almacenamiento para la 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 lotes más pequeños y una latencia de inferencia mejorada

Interconexión entre chips

Seis vínculos 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)
• Se mejoró el modelo de seguridad
• Se mejoró la eficiencia energética

Parámetros de configuración

El pod de TPU v4 se compone de 4,096 chips interconectados con vínculos 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 maneras. Cuando creas una porción de TPU, debes especificar la versión de TPU y la cantidad de recursos TPU que necesitas. Cuando creas una porción de TPU v4, puedes especificar su tipo y tamaño de una de las siguientes dos maneras: AcceleratorType y AccleratorConfig.

Nota: AcceleratorType y AccleratorConfig son mutuamente excluyentes; solo puedes usar uno cuando creas una TPU.

Usa AcceleratorType

Usa AcceleratorType cuando no especifiques una topología. Para configurar TPU v4 con AcceleratorType, usa la marca --accelerator-type cuando crees tu porción de TPU. Establece --accelerator-type en una cadena que contenga la versión de la TPU y la cantidad de TensorCores que deseas usar. Por ejemplo, para crear una porción de v4 con 32 TensorCores, usarías el --accelerator-type=v4-32.

Usa el comando gcloud compute tpus tpu-vm create para crear una porción de TPU v4 con 512 TensorCores con 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 TPU (v4) especifica la cantidad de TensorCores. Hay dos TensorCores en una TPU v4, por lo que la cantidad de chips TPU sería 512 / 2 = 256.

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

Usa AcceleratorConfig

Usa AcceleratorConfig cuando quieras personalizar la topología física de tu porción de TPU. Por lo general, esto es necesario para ajustar el rendimiento con porciones de más de 256 chips.

Para configurar las TPU v4 con AcceleratorConfig, usa las marcas --type y --topology. Establece --type en la versión de TPU que deseas usar y --topology en la disposición física de los chips TPU en la porción.

Debes especificar una topología de TPU con una tupla de 3 elementos, AxBxC, en la que 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 los recuentos de chips en cada una de las tres dimensiones. Por ejemplo, para crear una porción de v4 con 16 chips, deberías establecer --type=v4 y --topology=2x2x4.

Usa el comando gcloud compute tpus tpu-vm create para crear una porción de TPU v4 con 128 chips TPU dispuestos en un array 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, por ejemplo, 4×4×8, 8×8×16 y 12×12×24. Estas se conocen como topologías de toros retorcidos.

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

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

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

Topologías de toros retorcidos

Algunas formas de porciones de toro 3D de v4 tienen la opción de usar lo que se conoce como topología de toro retorcido. Por ejemplo, dos cubos v4 se pueden organizar como una porción de 4x4x8 o 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 en comparación con una porción 4x4x8 no retorcida. 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 los que más se benefician son los que tienen grandes cargas de trabajo de embedding de TPU.

En el caso de las cargas de trabajo que usan el paralelismo de datos como la única estrategia de paralelismo, las topologías retorcidas pueden tener un rendimiento ligeramente mejor. En el caso de los LLM, el rendimiento con una topología retorcida puede variar según el tipo de paralelismo (DP, MP, etc.). La práctica recomendada es entrenar tu LLM con y sin una topología retorcida para determinar cuál proporciona el mejor rendimiento a tu modelo. En algunos experimentos con el modelo FSDP MaxText, se observaron mejoras de 1 a 2 MFU con una topología retorcida.

El principal beneficio 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 muchos beneficios:

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

En última instancia, estos beneficios se traducen en un mejor rendimiento para muchos patrones de comunicación global.

El software de la TPU admite toros retorcidos en segmentos en los 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.

Como ejemplo, considera esta topología de toro de 4×2 con las TPU etiquetadas con sus coordenadas (X,Y) en la porción:

Para mayor claridad, los bordes de este grafo de topología se muestran como bordes no dirigidos. En la práctica, cada borde es una conexión bidireccional entre las TPU. Nos referimos a los bordes entre un lado de esta cuadrícula y el lado opuesto como bordes de envolvente, como se indica en el diagrama.

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

Lo único que cambió entre este diagrama y el anterior son las conexiones de integración Y que unen los extremos de una red. En lugar de conectarse a otra TPU con la misma coordenada X, se desplazaron para conectarse a la TPU con la coordenada X+2 modelo 4.

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

Consulta Usa AcceleratorConfig para obtener detalles sobre cómo especificar una configuración de toros retorcidos cuando creas una Cloud TPU.

En la siguiente tabla, se muestran las topologías retorcidas admitidas y un aumento teórico en el ancho de banda de bisección con ellas en comparación con las topologías no retorcidas.

Topología retorcida	Aumento teórico en el ancho de banda de la bisección 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, debes especificar la marca topology de la siguiente manera: 4x4x8_twisted.

Variantes de topología de TPU v4

Algunas topologías que contienen la misma cantidad de chips se pueden organizar de diferentes maneras. Por ejemplo, una porción de TPU con 512 chips (1,024 TensorCores) se puede configurar con las siguientes topologías: 4x4x32, 4x8x16 o 8x8x8. Una porción de TPU con 2,048 chips (4,096 TensorCores) ofrece aún más opciones de topología: 4x4x128, 4x8x64, 4x16x32 y 8x16x16.

La topología predeterminada asociada a una cantidad determinada de chips es la más similar a un cubo. Es probable que esta forma sea la mejor opción para el entrenamiento de AA paralelo a los 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, colocar un paralelismo del 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 coincide con una topología de 4x16x16.

Los modelos con múltiples dimensiones de paralelismo funcionan mejor cuando sus dimensiones de paralelismo se asignan a las dimensiones de topología de la TPU. Por lo general, se trata de modelos de lenguaje grandes (LLM) paralelos a los datos y al modelo. Por ejemplo, para una porción de TPU v4 con topología 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 del modelo de 4 vías sería subóptimo con esta topología, ya que no se alinea con ninguna de las dimensiones de la topología de TPU, pero sería óptimo con una topología de 4x16x32 en la misma cantidad de chips.

Los parámetros de configuración de TPU v4 constan de dos grupos: aquellos con topologías de menos de 64 chips (topologías pequeñas) y aquellos con topologías de más de 64 chips (topologías grandes).

Topologías v4 pequeñas

Cloud TPU admite los siguientes segmentos de TPU v4 más pequeños que 64 chips, un cubo de 4x4x4. Puedes crear estas topologías de v4 pequeñas con su nombre basado en TensorCore (por ejemplo, v4-32) o su topología (por ejemplo, 2x2x4):

Nombre (según el recuento de TensorCore)	Cantidad de chips	Topología
v4-8	4	2x2x1
v4-16	8	2x2x2
v4-32	16	2x2x4
v4-64	32	2x4x4

Topologías v4 grandes

Los segmentos 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 "personalizadas" y solo se pueden crear con las marcas --type y --topology porque hay más de una forma de organizar los chips.

Nombre (según el recuento de TensorCore)	Cantidad de chips	Topología
v4-128	64	4x4x4
v4-256	128	4x4x8
v4-512	256	4x8x8
topología personalizada: se deben usar las marcas --type y --topology	256	4x4x16
v4-1024	512	8x8x8
v4-1536	768	8x8x12
v4-2048	1024	8x8x16
topología personalizada: se deben usar las marcas --type y --topology	1024	4x16x16
v4-4096	2,048	8x16x16
…	…	…