Los Transformers de secuencias largas eran lentos y consumían mucha memoria porque la autoatención estándar materializaba grandes matrices intermedias y las movía repetidamente entre la memoria de alto ancho de banda de la GPU y las unidades de cómputo.
La solución tentadora era la atención aproximada: cambiar las matemáticas del modelo para reducir la complejidad. FlashAttention tomó el camino opuesto. Mantuvo la atención exacta, pero procesó bloques en SRAM y recalculó las estadísticas acumuladas del softmax para que nunca hubiera que almacenar la matriz de atención completa.
El artículo reportó una aceleración de 15% en tiempo de reloj de extremo a extremo en BERT-large frente al récord de entrenamiento de MLPerf 1.1, una aceleración de 3x en GPT-2 con longitud de secuencia de 1K, y una aceleración de 2.4x en tareas de Long Range Arena.
Optimizar el tráfico de memoria, no las matemáticas
La restricción central era la E/S entre HBM y SRAM, no solo las operaciones de punto flotante.
La división en bloques permitió que cada bloque hiciera trabajo útil mientras los datos ya estaban en el chip.
Mantener el algoritmo exacto evitó las regresiones de calidad que hacían poco atractivos muchos métodos aproximados.
La misma primitiva habilitó contextos más largos y nuevas tareas de secuencia sin cambiar la abstracción del Transformer.
Cuando el movimiento de hardware es el cuello de botella, un mejor algoritmo puede ser el que calcula la misma respuesta cruzando la jerarquía de memoria menos veces.
FlashAttention se convirtió en un componente estándar para el entrenamiento e inferencia eficientes de Transformers. Su lección más amplia se extendió más allá de la atención: los algoritmos modernos para aceleradores a menudo ganan diseñando en torno a la jerarquía de memoria antes de inventar nuevas matemáticas del modelo.
CONSULTA LAS PRUEBAS
Las fuentes
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness arXiv · arxiv.org ·