TPU v4

Este documento describe la arquitectura y las configuraciones admitidas 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 de vector 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 o máxima medida	90/170/192 O
Tamaño del pod de TPU	Chips 4,096
Topología de interconexión	Malla 3D
Procesamiento máximo por Pod	1.1 exaflops (bf16 o int8)
Ancho de banda de Allreduce 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 vecinos más cercanos en 3 de red, lo que da como resultado 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 un torus 3D configuración será mejor que una 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. Una CPU y su memoria se denomina 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 multiCPU puede colocar los datos que necesita una CPU su nodo NUMA, lo que aumenta la capacidad de procesamiento de la memoria. Para obtener más información sobre NUMA, consulta Acceso a la 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 tu código de secuencia de comandos con el 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, (cargas de trabajo de recomendación, clasificació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.
• Mejor rendimiento de HBM con las velocidades y los estándares 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 MXU y una frecuencia de reloj más alta con un máximo de 275 TFLOPS
• El doble de ancho de banda de transposición y permutación.
• Modelo de acceso a memoria de almacenamiento de carga para memoria común (Cmem).
• Ancho de banda de carga de MXU más rápido y compatibilidad con el modo de 8 bits para permitir una menor y mejoró la latencia de la inferencia.

Interconexión entre chips

Seis vínculos de interconexión por chip para habilitar topologías de red los diámetros de la red.

Otro

• Interfaz PCIE 3 de x16 para alojar (conexión directa).
• Modelo de seguridad mejorado.
• Eficiencia energética mejorada

Configuraciones

Un Pod de TPU v4 se compone de 4,096 chips interconectados con cables reconfigurables vínculos de alta velocidad. 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, especifica la versión de TPU y la cantidad de recursos 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 después de la versión de TPU (v4) especifica la cantidad de tensores. Hay dos núcleos tensores 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 --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 de pod.

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

El siguiente comando crea una porción de pod de TPU v4 con 128 chips TPU organizados en un array 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.17.0-pod-pjrt

Las topologías en las que 2A=B=C o 2A=2B=C también tienen variantes de topología Optimizada para la comunicación general, 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 4x4x4. de chips.

Topologías tori torcidas

Algunas formas de rebanadas de torus 3D v4 tienen la opción de usar lo que se conoce como torus trenzado. Por ejemplo, dos cubos v4 pueden organizarse como una resolución de 4x4x8 porción o 4x4x8_twisted. Las topologías torcidas ofrecen una bisección significativamente mayor ancho de banda. Un mayor ancho de banda de bisección es útil para cargas de trabajo que usan los patrones de comunicación. Las topologías trenzados pueden mejorar el rendimiento de la mayoría con grandes cargas de trabajo de incorporación de TPU se beneficien más.

Para las cargas de trabajo que usan el paralelismo de datos como la única estrategia de paralelismo, trenzado las topologías podrían funcionar un poco mejor. Para los LLM, el rendimiento con una estructura La topología puede variar según el tipo de paralelismo (DP, MP, etcétera). Recom. es entrenar tu LLM con y sin una topología retorcida para determinar que proporciona el mejor rendimiento para tu modelo. Algunos experimentos El modelo MaxText de FSDP tienen notó mejoras de 1 a 2 MFU mediante 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 tori trenzado en porciones en las que el tamaño de cada sea igual o doble al tamaño de la dimensión más pequeña. Para 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 su (X,Y) en la porción:

Las aristas de este gráfico de topología se muestran como aristas no dirigidas para mayor claridad. En cada perímetro es una conexión bidireccional entre las TPU. Nos referimos al entre un lado de esta cuadrícula y el opuesto como bordes envolventes, como se muestra 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 distintos tamaños de dimensiones y cantidades distintas dimensiones. La red resultante es simétrica, siempre que cada dimensión esté igual o doble al tamaño de la dimensión más pequeña.

Consulta cómo usar AcceleratorConfig para obtener más información para 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 una lista teórica de aumento en el ancho de banda de bisección con ellos en comparación con las topologías no torcidas.

Topología trenzada	Aumento teórico de la bisección ancho de banda frente a un toro
4 × 4 × 8_trenzado	Aprox. un 70%
8x8x16_trenzado
12 × 12 × 24_trenzado
4 × 8 × 8_trenzado	Aprox. un 40%
8 × 16 × 16_trenzado

Nota: La topología trenzada no se usará automáticamente, a menos que se use el sufijo _twisted se especifica con la marca --topology. Por ejemplo, para especificar una cadena topología, 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 en diferentes maneras. Por ejemplo, una porción de pod de TPU con 512 chips (1,024 núcleos de tensor) puede ser configurada con las siguientes topologías: 4x4x32, 4x8x16 u 8x8x8. 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 de datos paralelos. 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 grandes de lenguaje (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. Sí más eficaz si usas un paralelismo de modelos de 8 o 16 vías (asignado a uno de los dimensiones de la topología de TPU física). Un paralelismo de modelo de 4 vías sería menos ó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 4x16x32 en la misma cantidad de chips.

Las configuraciones de TPU v4 constan de dos grupos, aquellos con topologías más pequeñas que 64 chips (topologías pequeñas) y aquellos con topologías de más de 64 chips (topologías grandes).

Topologías de v4 pequeñas

Cloud TPU admite las siguientes porciones de pod de TPU v4 más pequeñas que 64 chips: un cubo de 4 x 4 x 4. Puedes crear estas topologías pequeñas v4 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 creciente. En la siguiente tabla, se muestran varios ejemplos. Algunos de estas topologías son “personalizadas” que solo pueden crearse usando el Marcas --type y --topology porque 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 × 4 × 8
v4-512	256	4 × 8 × 8
N/A: 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 × 8 × 16
N/A (debe usar las marcas --type y --topology)	1024	4 × 16 × 16
v4-4096	2,048	8 × 16 × 16
…	…	…