Análisis profundo de las TPU

• Las TPU son chips personalizados desarrollados por Google para entrenamiento e inferencia de IA a gran escala, con una filosofía de diseño distinta a la de las GPU
• Enfatizan la escalabilidad y la eficiencia energética, diseñando en conjunto hardware (por ejemplo, configuración de sistema en chip, memoria on-chip de gran tamaño) y software (compilador XLA)
• Su estructura central se basa en arreglos sistólicos, pipeline y compilación anticipada, optimizada para la mayoría de las operaciones de deep learning, especialmente la multiplicación de matrices
• Gracias a las tecnologías OCI y OCS, es posible contar con configuración flexible de nodos, procesamiento paralelo de alto rendimiento y distintas opciones de topología
• Los sistemas TPU escalan jerárquicamente desde un solo chip hasta multipods gigantes, lo que permite entrenar modelos de IA a gran escala y usar los recursos de forma eficiente

---

Panorama general y contexto

• Las TPU son chips basados en ASIC desarrollados por Google para responder a la expansión masiva de sus servicios de IA
• Al inicio se evaluó la adopción de GPU, FPGA y ASIC, pero a partir de 2013, con el aumento del uso de deep learning en búsqueda por voz y otros servicios, creció la necesidad de hardware personalizado
• Hoy cumplen un papel esencial en servicios de IA clave como Gemini y Veo, y también se usan ampliamente en modelos de recomendación como DLRM

Estructura de un chip TPU individual

Configuración básica

• En TPUv4, un chip individual incluye 2 TensorCore (las TPU especializadas en inferencia incluyen 1)
• Cada TensorCore está conectado a memoria CMEM (128MiB) y HBM (32GiB)

Componentes internos de TensorCore

• Matrix Multiply Unit (MXU): basada en un arreglo sistólico de 128x128, dedicada a la multiplicación de matrices
• Vector Unit (VPU): realiza operaciones generales elemento por elemento
• Vector Memory (VMEM; 32MiB): copia datos desde HBM y se usa como memoria de preparación para el trabajo
• Scalar Unit + Scalar Memory (SMEM; 10MiB): se encarga del flujo de control, operaciones escalares y gestión de direcciones de memoria

Diferencias estructurales frente a la GPU

• Las TPU tienen mucha más memoria on-chip (CMEM, VMEM, SMEM) que las GPU
• La capacidad de HBM es mayor en las GPU, y la cantidad de núcleos de cómputo también es mucho más alta en GPU
• En TPUv5p, se logra un alto rendimiento de 500 TFLOPs/sec por chip y 4.45 ExaFLOPs/sec por pod completo (8960 chips)

Filosofía de diseño de las TPU

1. Arreglos sistólicos y pipeline

• En un arreglo sistólico, los elementos de procesamiento (PE) están dispuestos en una matriz y transfieren los resultados a elementos vecinos
• Una vez que entran los datos, se pueden realizar operaciones continuas sin control adicional, y las lecturas/escrituras de memoria solo ocurren en la entrada y salida
• Está optimizado principalmente para multiplicación de matrices y convolución
• Mediante procesamiento en pipeline, el cómputo y el movimiento de datos avanzan al mismo tiempo para optimizar el throughput

Desventaja de los arreglos sistólicos: sparsity

• Como todos los elementos de cómputo están siempre activos, no son adecuados para matrices dispersas
• Si en el futuro los modelos de DL avanzan hacia sparsity irregular, esto podría convertirse en una limitación

2. Compilación anticipada (AoT) y minimización de caché

• El codiseño TPU-XLA minimiza la dependencia de caché, que requiere accesos de memoria irregulares, y reduce el consumo de energía
• El compilador XLA analiza el grafo de operaciones para calcular de antemano los patrones de acceso a memoria, operando con enfoque en memoria scratchpad en lugar de caché
• En JAX, @jit es una forma intermedia entre JIT y AoT: en la primera ejecución genera un grafo estático y luego realiza compilación AoT con XLA
• Si cambia la forma de entrada, se requiere recompilación, y existe ineficiencia en padding dinámico o iteraciones
• Se logra una alta eficiencia energética, aunque con la desventaja de menor flexibilidad

Eficiencia energética de TPUv4

• Los chips modernos usan memoria HBM3 para reducir el consumo energético
• Como las operaciones de memoria consumen decenas o cientos de veces más energía que el cómputo, minimizar los accesos a memoria puede mejorar mucho la eficiencia

Estructura multinodo de TPU

Nivel tray (Tray/Board; 4 chips)

• Un tray está compuesto por 4 chips TPU (8 TensorCore) y un CPU Host
• La conexión Host↔Chip usa PCIe, y la conexión Chip↔Chip usa Inter-Core Interconnect (ICI), que ofrece mayor ancho de banda

Nivel rack (Rack; 4x4x4=64 chips)

• Un rack contiene 64 chips TPU, conectados mediante toro 3D (4x4x4) con ICI y OCS (Optical Circuit Switching)
• En Google se distinguen rack, pod y slice
  • Rack: unidad física de 64 chips (=cubo)
  • Pod: unidad máxima conectable con ICI y OCS (por ejemplo, TPUv4 = 4096 chips = 64 racks)
  • Slice: unidad abstracta de configuración arbitraria entre 4 chips y un Superpod

Ventajas de OCS

1. Wraparound: cada eje se convierte en un anillo (toro 1D), reduciendo la cantidad máxima de saltos entre nodos
2. Slices no contiguos flexibles: con la estructura de switching de OCS, nodos físicamente separados pueden formar un mismo slice, lo que facilita el uso de recursos y el mantenimiento
3. Topologías torcidas: manteniendo el mismo número de chips (x, y, z fijos), se puede cambiar la estructura de conexión para optimizar velocidad según ciertos patrones de operación (por ejemplo, twisted torus)

Casos de uso de topologías

• Forma cúbica: adecuada para paralelismo de datos/tensores (máximo ancho de banda)
• Forma lineal (tipo cigarro): favorable para paralelismo en pipeline
• Twisted torus: mejora la velocidad cuando se necesita comunicación all-to-all (por ejemplo, tensor parallel)

Superpod (Full Pod/Superpod; TPUv4: 4096 chips/64 racks)

• Se construyen sistemas gigantes conectando varios racks con ICI y OCS
• El ancho de banda de comunicación y el rendimiento del procesamiento paralelo varían según el tipo de topología del slice
• Gracias a OCS, se obtiene flexibilidad para slices no contiguos y topologías torcidas

Nivel multipod (Multi-pod/Multislice; TPUv4: más de 4096)

• Varios pods se conectan mediante la red de centro de datos (DCN), proporcionando infraestructura de entrenamiento a gran escala, aunque con menor ancho de banda que ICI
• El entrenamiento de PaLM utilizó 2 pods (6144 TPUv4), gestionando en total recursos de 6 pods
• En el desarrollo de modelos grandes, el compilador XLA optimiza los patrones de comunicación según los parámetros definidos por los investigadores, como las dimensiones de paralelismo
• XLA inserta operaciones de comunicación entre slices y pods, permitiendo entrenamiento distribuido a gran escala con cambios mínimos de código

Ejemplos de hardware real y diagramas

• Rack de TPU: un toro 3D 4x4x4 forma una unidad, con 2 trays (8 chips) por cada fila
• Tray de TPUv4: en la práctica tiene 4 puertos PCIe (uno por cada TPU)
• Chip TPUv4: ASIC en el centro, rodeado por 4 stacks de HBM (con base en 2 TensorCore)
• Floorplan del chip TPUv4i (para inferencia): 1 TensorCore, con gran área ocupada por CMEM

Cierre

• Se expresa agradecimiento al apoyo de investigación de Google TPU Research Cloud(TRC)

Referencias

• TPU Multi-Slice Trainng
• Xu et al., paper de GSPMD
• Jouppi et al., paper de TPUv4i
• How to Scale Your Model - TPUs
• Además, unas 10 publicaciones académicas y materiales oficiales adicionales