Arquitecturas de Hardware para Deep Learning: De GPUs a TPUs

Introducción a las Redes Neuronales y el Deep Learning

La necesidad de heterogeneidad

Debido al fin de la ley de Dennard y la desaceleración de la Ley de Moore, la especialización de hardware es la única alternativa viable para seguir mejorando la relación energía-prestaciones-coste. Esto ha impulsado el auge de las Arquitecturas de Dominio Específico (DSAs).

Modelo Roofline

Es un método intuitivo para identificar cuellos de botella en el rendimiento. Relaciona el rendimiento en coma flotante (FLOP/s) con el ancho de banda de memoria mediante la "Intensidad Aritmética" (operaciones por byte accedido). Ayuda a determinar si una aplicación está limitada por la memoria (memory-bound) o por la capacidad de cálculo (compute-bound).

Implicaciones de Hardware de las DNNs

• MLPs y RNNs: Tienen baja intensidad aritmética y suelen estar limitadas por la memoria (memory-bound).
• CNNs y Transformers: Tienen alta reutilización de datos y paralelismo masivo, lo que resulta en alta intensidad aritmética, estando limitadas por el cómputo (compute-bound).

Arquitecturas GPU

¿Por qué las GPUs son perfectas para Deep Learning?

Ofrecen paralelismo masivo de datos, dominan en álgebra lineal (como la multiplicación de matrices) y tienen un ancho de banda de memoria extremadamente elevado (como la memoria HBM).

Arquitectura SIMT (Single Instruction, Multiple Thread)

El hardware organiza y ejecuta hilos en grupos de 32 llamados WARPs. Todos los hilos del warp ejecutan la misma instrucción de forma sincronizada pero con diferentes datos, lo que requiere evitar divergencias de control (como los if-else dispares dentro de un warp) para no perder rendimiento.

Extensiones Hardware para Deep Learning

• Tensor Cores: Unidades especializadas que aceleran dramáticamente la multiplicación de matrices (GEMM) al ejecutar docenas de operaciones simultáneas de multiplicar y sumar (FMA) en un solo ciclo.
• Precisión Reducida (Mixed Precision): Las redes neuronales son tolerantes al ruido, lo que permite usar formatos de 16 bits (FP16, BF16) o incluso de 8 bits para duplicar o cuadruplicar el rendimiento y reducir a la mitad el consumo de memoria sin perder calidad en los modelos.
• NVLink: Interconexión de muy alta velocidad entre GPUs para superar los cuellos de botella del bus PCIe, esencial para escalar el entrenamiento de grandes modelos en múltiples GPUs.

Arquitecturas de Dominio Específico (Google TPUs)

Arrays Sistólicos

Son el núcleo de las TPUs. Consisten en una malla bidimensional de elementos de proceso (PEs) donde los datos fluyen rítmicamente en cada ciclo de reloj. Esto maximiza enormemente la reutilización de datos traídos desde la memoria y mejora la eficiencia energética.

Evolución de las TPUs de Google

• TPU v1 (2015): Primer ASIC para producción, diseñado exclusivamente para inferencia. Usaba enteros de 8 bits y fue hasta 29 veces más rápida que las CPUs contemporáneas.
• TPU v2 y v3: Añadieron soporte para entrenamiento, uso de memoria HBM, el formato numérico bfloat16, e introdujeron una red de interconexión propia (ICI) para unir chips. La v3 implementó refrigeración líquida.
• TPU v4 y v4i: Introdujeron conmutadores ópticos (OCS) para reconfigurar la red en nanosegundos, y el "SparseCore", un bloque de hardware dedicado para acelerar operaciones dispersas (como tablas de embeddings en sistemas de recomendación).
• TPU v5 (v5e y v5p): Marcó una bifurcación: "v5e" orientada a la eficiencia y bajo coste para inferencia, y "v5p" orientada al máximo rendimiento para el entrenamiento de grandes modelos de lenguaje (LLMs) como Gemini.
• TPU v6e y v7 (Ironwood): La v7, diseñada para la era de los LLMs, implementa soporte nativo para el formato FP8 e incorpora cantidades masivas de memoria (192 GB HBM3e) para soportar las inmensas KV Caches necesarias al procesar grandes contextos de texto.