TPU v4

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

Arquitectura del sistema

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

Especificaciones clave	Valores de Pods de 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/máxima medida	90/170/192 O
Tamaño del pod de TPU	4,096 chips
Topología de interconexión	Malla 3D
Procesamiento máximo por Pod	exaflops 1.1 (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 ilustra un chip TPU v4.

Malla 3D y torus 3D

Las TPU v4 tienen una conexión directa a los chips cercanos 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 torus 3D en porciones en las que la topología, AxBxC, es 2A=B=C o 2A=2B=C, en la que cada dimensión es un múltiplo de 4. Por ejemplo, 4 × 4 × 8, 4 × 8 × 8 o 12 × 12 × 24. En general, el rendimiento de una configuración de torus 3D será mejor que una configuración de malla 3D. Para obtener más información, consulta Topologías de tori trenzado.

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

En esta sección, se describen los beneficios de rendimiento de TPU v4

Sistema de memoria

El Acceso a la memoria no uniforme (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 con otros nodos de NUMA que son directamente 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 acceder a la memoria dentro de un nodo de NUMA.

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

Puedes aprovechar los beneficios de la localidad de NUMA si vinculas la secuencia de comandos de entrenamiento al nodo 0 de NUMA.

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

1. Instalar la herramienta de línea de comandos numactl

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

2. Vincula el código de la secuencia de comandos al nodo 0 de NUMA.

    $ 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 CPU (por ejemplo, clasificación de imágenes o cargas de trabajo de recomendación), sin importar el 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 HBM unificado de 32 GiB en todo el chip, lo que permite una mejor coordinación entre los dos TensorCores en el chip.
• Se mejoró el rendimiento de HBM gracias a los estándares de memoria y las velocidades más recientes.
• Se mejoró el perfil de rendimiento de DMA con compatibilidad integrada para alto rendimiento con niveles de detalle de 512 B.

TensorCores

• El doble de MXU y una mayor frecuencia de reloj, lo que genera un máximo de 275 TFLOPS.
• El doble de ancho de banda de transposición y permutación.
• Modelo de acceso a la memoria de 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 lotes menores y una mejor latencia de inferencia

Interconexión entre chips

Seis vínculos de interconexión por chip para habilitar topologías de red que tienen diámetros de red más pequeños.

Otro

• Interfaz x16 PCIE gen3 para el host (conexión directa).
• Se mejoró el modelo de seguridad.
• Mayor eficiencia energética.

Parámetros de configuración

Un pod de TPU v4 consta de 4,096 chips interconectados con vínculos de alta velocidad reconfigurables. Las redes flexibles de TPU v4 te permiten conectar los chips en una porción de Pod del mismo tamaño de varias maneras. Cuando creas una porción de pod de TPU, debes especificar la versión de TPU y la cantidad de recursos de TPU que necesitas. Cuando creas una porción de pod de TPU v4, puedes especificar su tipo y tamaño de una de estas dos maneras: AcceleratorType y AccleratorConfig.

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

Usa AcceleratorType

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

Con el siguiente comando, se crea una porción de pod de TPU v4 con 512 TensorCores mediante la marca --accelerator-type:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --accelerator-type=v4-512 \
    --version=tpu-vm-tf-2.16.1-pod-pjrt

El número 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.

Usa AcceleratorConfig

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

Para configurar las TPU v4 con AcceleratorConfig, usa las marcas --version y --topology. Establece --version 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 de Pod.

Debes especificar una topología de TPU con una tupla de 3, AxBxC, en la que A<=B<=C y A, B, C son todos <= 4 o todos son 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 Pod v4 con 16 chips, debes configurar --version=v4 y --topology=2x2x4.

El siguiente comando crea una porción de pod de TPU v4 con 128 chips TPU organizados en un arreglo de 4 x 4 x 8:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --type=v4 \
    --topology=4x4x8 \
    --version=tpu-vm-tf-2.16.1-pod-pjrt

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 todo a todos, por ejemplo, 4 × 4 × 8, 8 × 8 × 16 y 12 × 12 × 24. Estas topologías se conocen como topologías torilas retorcidas.

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

Las porciones de Pod más grandes se pueden compilar a partir de uno o más “cubos” de 4×4×4 de chips.

Topologías de tori trenzado

Algunas formas de porciones de torus 3D v4 tienen la opción de usar lo que se conoce como una topología de torus trenzado. Por ejemplo, se pueden organizar dos cubos v4 como una porción de 4x4x8 o 4x4x8_twisted. Las topologías trenzadas ofrecen un ancho de banda de bisección significativamente mayor. Un mayor ancho de banda de bisección es útil para 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 grandes cargas de trabajo de incorporación de TPU son las que más se benefician.

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 un poco mejor. Para los LLM, el rendimiento con una topología retorcida puede variar según el tipo de paralelismo (DP, MP, etcétera). La práctica recomendada es entrenar tu LLM con y sin una topología retorcida para determinar cuál le proporciona el mejor rendimiento a tu modelo. Algunos experimentos en el modelo MaxText para FSDP experimentaron 1 o 2 mejoras de MFU mediante una topología retorcida.

El beneficio principal de las topologías trenzadas es que transforman una topología de torus 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 globales.

El software de la TPU admite tori trenzado en porciones en las que el tamaño de cada dimensión es igual o dos veces el tamaño de la dimensión más pequeña. Por ejemplo, 4 × 4 × 8, 4 × 8 × 8 o 12 × 12 × 24.

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

Para mayor claridad, los bordes de este gráfico de topología se muestran como bordes no dirigidos. En la práctica, cada perímetro 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 envolventes, como se indica en el diagrama.

Cuando se torce esta topología, obtenemos una topología de torus trenzado de 4 × 2 completamente simétrica:

Todo lo que cambió entre este diagrama y el anterior son los bordes envolventes Y. En lugar de conectarse a otra TPU con la misma coordenada X, se cambiaron 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 diferentes cantidades de dimensiones. La red resultante será simétrica, siempre que cada dimensión sea igual o dos veces el tamaño de la dimensión más pequeña.

Consulta Cómo usar AcceleratorConfig para obtener detalles sobre cómo especificar una configuración de tori trenzado cuando se crea una Cloud TPU.

En la siguiente tabla, se muestran las topologías trenzadas 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 trenzada	Aumento teórico del ancho de banda de bisección en comparación con un torus no trenzado
4×4×8_twisted	Aprox. un 70%
8 x 8 x 16_trezado
12×12×24_twisted
4×8×8_twisted	Aprox. un 40%
8 × 16 × 16_twist

Nota: La topología trenzada 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 4 x 4 x 8, 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 pod de TPU con 512 chips (1, 024 TensorCores) se puede configurar mediante las siguientes topologías: 4 × 4 × 32, 4 × 8 × 16 y 8 × 8 × 8. Una porción de pod de TPU con 2,048 chips (4,096 TensorCores) ofrece aún más opciones de topología: 4 × 4 × 128, 4 × 8 × 64, 4 × 16 × 32 y 8 × 16 × 16.

La topología predeterminada asociada con un recuento de chips determinado es la más similar a un cubo. Es probable que esta forma sea la mejor opción para el entrenamiento del AA con datos paralelos. Otras topologías pueden ser útiles para cargas de trabajo con varios tipos de paralelismo (por ejemplo, paralelismo de modelos y datos o partición espacial de una simulación). Estas cargas de trabajo tienen un mejor rendimiento si la topología coincide con el paralelismo utilizado. Por ejemplo, ubicar el paralelismo de modelos de 4 vías en la dimensión X y el paralelismo de datos de 256 vías en las dimensiones Y y Z coincide con una topología de 4 × 16 × 16.

Los modelos con varias dimensiones de paralelismo tienen un mejor rendimiento cuando las dimensiones de paralelismo se asignan a las dimensiones de topología de TPU. Por lo general, estos son datos y modelan modelos grandes de lenguaje (LLM) paralelos. Por ejemplo, para una porción de pod de TPU v4 con una topología de 8 × 16 × 16, las dimensiones de topología de TPU son 8, 16 y 16. Es más eficaz usar paralelismo de modelos de 8 o 16 vías (asignado a una de las dimensiones de la topología de TPU física). Un paralelismo de modelos de 4 vías sería subóptimo con esta topología, ya que no está alineada con ninguna de las dimensiones de topología de TPU, pero sería óptimo con una topología de 4 × 16 × 32 en la misma cantidad de chips.

Las configuraciones de TPU v4 constan de dos grupos: los que tienen topologías menores a 64 chips (topologías pequeñas) y las que tienen topologías de más de 64 chips (topologías grandes).

Topologías v4 pequeñas

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

Nombre (según la cantidad 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

Las porciones de pod 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 creciente. 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 porque hay más de una manera de organizar los chips.

Nombre (según la cantidad de TensorCore)	Cantidad de chips	Topología
v4-128	64	4 × 4 × 4
v4-256	128	4 × 4 × 8
v4-512	256	4 × 8 × 8
N/A: Debes usar las marcas --type y --topology.	256	4 × 4 × 16
v4-1024	512	8 × 8 × 8
v4-1536	768	8 × 8 × 12
v4-2048	1024	8 × 8 × 16
N/A: Debes usar las marcas --type y --topology	1024	4x16x16
v4-4096	2,048	8 × 16 × 16
…	…	…