TPU v4

En este documento, se describe 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 matrices (MXU), una unidad vectorial y una unidad escalar. En la siguiente tabla, se muestran las especificaciones clave para un 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 O
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 de reducción total 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 con los chips vecinos más cercanos en 3 dimensiones, lo que genera una malla 3D de conexiones de red. Las conexiones se puede configurar como un torus 3D en porciones en las que la topología, AxBxC, 2A=B=C o 2A=2B=C, donde 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 toro 3D será mejor que el de una configuración de malla 3D. Para obtener más información, consulta Topologías tori torcidas.

Beneficios de rendimiento de TPU v4 sobre 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 NUMA. Los nodos de NUMA están conectados a otro Nodos de NUMA que son directamente adyacentes entre sí. Una CPU de un nodo de NUMA puede acceder a la memoria en otro nodo de NUMA, pero este acceso es más lento que el memoria dentro de un nodo 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 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 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 NUMA 0.

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

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

• Si tu carga de trabajo depende en gran medida de las cargas de trabajo de la CPU (por ejemplo, la clasificación de imágenes, las cargas de trabajo de recomendación), independientemente del framework.
• Si usas una versión de 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.
• Rendimiento mejorado de HBM con los estándares y las velocidades de memoria más recientes
• Perfil de rendimiento de DMA mejorado con compatibilidad integrada de alto rendimiento con un nivel de detalle de 512 mil millones.

TensorCores

• El doble de la cantidad de MXUs y una velocidad de reloj más alta que ofrece 275 TFLOPS máximos
• Ancho de banda de transposición y permutación 2x.
• Modelo de acceso a memoria de almacenamiento de carga 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 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 que tengan diámetros de red más pequeños.

Otro

• Interfaz PCIE gen3 x16 al host (conexión directa)
• Modelo de seguridad mejorado.
• Mayor eficiencia energética.

Configuraciones

Un pod de TPU v4 se compone de 4,096 chips interconectados con vínculos de alta velocidad reconfigurables. La red flexible de TPU v4 te permite conectar los chips en una porción de Pod del mismo tamaño de varias maneras. Cuando creas una porción de TPU Pod, 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 dos maneras: AcceleratorType y AccleratorConfig.

Nota: AcceleratorType y AccleratorConfig son mutuamente excluyentes. puedes solo usa una cuando crees 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 pod de TPU. Establece --accelerator-type en una cadena que contenga la TPU. y la cantidad de TensorCores que quieres usar. Por ejemplo, para crear un Porción de Pod v4 con 32 tensores, deberías usar --accelerator-type=v4-32.

El siguiente comando crea una porción de pod de TPU v4 con 512 núcleos tensores mediante la marca --accelerator-type:

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

El número que aparece después de la versión de la 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 quieras personalizar la topología física de tu porción de TPU. Por lo general, se requiere para ajustar el rendimiento con el Pod porciones de más de 256 chips.

Para configurar 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 del pod.

Especificas una topología de TPU con un 3-tuplo, AxBxC, donde A<=B<=C y A, B y C son todos <= 4 o son todos múltiplos enteros de 4. Los valores A, B y C son cuenta el chip 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 un pod de TPU v4 con 128 chips de TPU organizados en un array de 4 × 4 × 8:

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

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 tori trenzado.

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 chips de 4 × 4 × 4.

Topologías de toros trenzados

Algunas formas de rebanada de toro 3D de la v4 tienen la opción de usar lo que se conoce como una topología de toro torcido. Por ejemplo, dos cubos v4 se pueden organizar como una porción 4 × 4 × 8 o 4 × 4 × 8_trenzada. Las topologías torcidas ofrecen una bisección significativamente mayor ancho de banda. 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 trenzadas pueden mejorar el rendimiento de la mayoría de los modelos, y las cargas de trabajo de incorporación de TPU grandes se benefician más.

En el caso de las cargas de trabajo que usan el paralelismo de datos como la única estrategia de paralelismo, las topologías torcidas podrían tener un mejor rendimiento. En el caso de los LLM, el rendimiento con una topología trenzada 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 trenzada para determinar cuál proporciona el mejor rendimiento para tu modelo. Algunos experimentos en el modelo MaxText de FSDP mostraron mejoras de 1 a 2 MFU con una topología trenzada.

El principal beneficio de las topologías tornidas es que transforman una topología de Torus (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 los patrones de comunicación.

El software de TPU admite toros torcidos en rebanadas 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. Para por ejemplo, 4×4×8, 4×8×8 o 12×12×24.

A modo de ejemplo, considera esta topología de toro 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 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 unión, como se indica en el diagrama.

Con la retorsión de esta topología, obtenemos un toro retorcido de 4 × 2 completamente simétrico. Topología:

Todo lo que cambió entre este diagrama y el anterior es la Y con bordes envolventes. En lugar de conectarse a otra TPU con el mismo X se han desplazado para conectarse a la TPU con X+2 mod 4.

La misma idea se generaliza a diferentes tamaños de dimensiones y diferentes cantidades de estas. La red resultante es simétrica, siempre que cada dimensión sea igual o el doble del 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 toro torcido cuando creas una Cloud TPU.

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

Topología trenzada	Aumento teórico en el ancho de banda de la bisección en comparación con un toro no torcido
4×4×8_trenzado	Aprox. un 70%
8x8x16_trenzado
12 × 12 × 24_trenzado
4 × 8 × 8_trenzado	~40%
8 × 16 × 16_trenzado

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 trenzada de 4 × 4 × 8, debes especificar la marca topology como 4x4x8_twisted.

Variantes de topología de TPU v4

Algunas topologías que contienen la misma cantidad de chips se pueden organizar en diferentes maneras. Por ejemplo, una porción de pod de TPU con 512 chips (1,024 TensorCores) se puede configurar con las siguientes topologías: 4 × 4 × 32, 4 × 8 × 16 o 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: 4x4x128, 4 × 8 × 64, 4 × 16 × 32 y 8 × 16 × 16.

La topología predeterminada asociada con un recuento de chips determinado es la que más similar a un cubo. Es probable que esta forma sea la mejor opción para el entrenamiento de AA en paralelo de datos. Otra opción Las topologías pueden ser útiles para cargas de trabajo con varios tipos de paralelismo (para ejemplo, el paralelismo de modelos y datos o la 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 paralelismo de modelos de 4 vías en la dimensión X y datos de 256 vías el paralelismo en las dimensiones Y y Z coincide con una topología de 4 × 16 × 16.

Los modelos con múltiples dimensiones de paralelismo funcionan mejor con sus dimensiones de paralelismo asignadas a las dimensiones de la topología de TPU. Por lo general, son de datos y modelos de lenguaje extenso (LLM) paralelos. Por ejemplo, para una TPU v4 Porción de pod con una topología de 8x16 x 16. 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 TPU). Un paralelismo de 4 vías sería subóptimo con esta topología, ya que no está alineado con ninguna de las dimensiones de la 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 consisten 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 de v4 pequeñas

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

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 grandes de v4

Las porciones de pod de TPU v4 están disponibles en incrementos de 64 chips, con formas que 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, ya que hay más de una forma de organizar los chips.

Nombre (según la cantidad de TensorCore)	Cantidad de chips	Topología
v4-128	64	4x4x4
v4-256	128	4 x 4 x 8
v4-512	256	4 x 8 x 8
topología personalizada: Se deben 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 x 8 x 16
topología personalizada: Se deben usar las marcas --type y --topology.	1024	4 × 16 × 16
v4-4096	2,048	8 × 16 × 16
…	…	…