FlashAttention-3 acelerará los transformadores al doble con 75% de carga de GPU

Together AI, NVIDIA y Meta lanzaron FlashAttention-3, un algoritmo mejorado para acelerar la atención en transformadores. La nueva versión utiliza el 75% de la potencia de GPU moderna en lugar del anterior 35% y funciona de 1.5 a 2 veces más rápido en aceleradores NVIDIA H100. Esto es crítico: mientras los cálculos en la nube se encarecen y los LLM requieren cada vez más recursos, ahora cada porcentaje de eficiencia se cuenta en dinero.

Por qué fue un cuello de botella

La atención (attention) es el corazón de los transformadores. Pero es la parte más costosa de los cálculos. Al procesar texto largo, requiere una cantidad cuadrática de memoria: duplica la longitud del contexto y necesitas cuatro veces más memoria.

FlashAttention resolvió este problema en 2022, reorganizando los cálculos para leer/escribir en la memoria de GPU de manera más eficiente. Esto dio una aceleración de 2-4x y permitió que los LLM extiendan el contexto de 4K tokens a 128K, e recientemente a un millón. Sin FlashAttention, tales contextos largos simplemente eran imposibles.

Pero FlashAttention-2 se estancó en el 35% de eficiencia de H100. Las nuevas GPU de la serie Hopper trajeron núcleos asincronos (WGMMA) y transferencia de datos asincronos (TMA). FlashAttention-3 finalmente los utiliza.

Cómo funciona la aceleración

FlashAttention-3 aplica tres mejoras clave. Primero, asincronía. Los núcleos tensoriales y los sistemas de transferencia de datos funcionan simultáneamente sin esperarse uno al otro. Esta intersección de cálculo y memoria es la aceleración principal. Segundo, entrelazamiento de operaciones. En lugar de calcular primero todas las multiplicaciones de matrices por bloques y luego softmax, el algoritmo los intercala. Esto reduce la latencia y mejora el uso de caché. Tercero, precisión de cálculo baja. FP8 (números de ocho bits en lugar de FP16) requiere la mitad de memoria y permite el doble de operaciones por segundo. FlashAttention-3 con FP8 alcanza aproximadamente 1.2 PFLOPS (petaflops).

¿No pierde calidad con FP8?

El riesgo principal de la baja precisión es que los errores se acumulan y dañan los resultados. Pero los autores muestran que FlashAttention-3 con FP8 tiene un error solo 2.6 veces mayor que la atención FP8 de base. Y en comparación con FP16, el error es apenas perceptible. Esto es importante para contextos largos. Cuando un LLM procesa un millón de tokens, los errores en una capa de atención se acumulan a través de 70+ capas del modelo. FlashAttention-3 mantiene estos errores bajo control, permitiendo usar FP8 sin pérdida de calidad.

Quién se beneficiará

Diferentes escenarios se beneficiarán de FlashAttention-3:

• Entrenamiento de modelos grandes - la aceleración ahorra meses de tiempo computacional. Para Meta y OpenAI, esto significa poder entrenar más variantes con el mismo presupuesto. Las facturas en la nube se reducen en 30-50%.
• Inferencia rápida en producción - en una sola H100 se pueden servir más usuarios simultáneamente, haciendo que las API comerciales sean más económicas
• Contextos largos - sistemas RAG, análisis de documentos grandes, búsqueda en archivos ahora funcionan sin ralentización incluso con un millón de tokens
• Cómputo móvil y edge - FP8 y la reducción de memoria permiten ejecutar modelos en hardware débil

El algoritmo ya está incorporado en las principales bibliotecas de ML PyTorch, JAX y TensorFlow. El código se ha publicado en GitHub y los desarrolladores han comenzado a implementarlo.

Qué significa para la industria

FlashAttention-3 llega en el momento oportuno. Las GPU en la nube se encarecen, la demanda de LLM crece y los contextos se hacen más largos. Los algoritmos que hacen que el hardware funcione al 75% de su potencia en lugar del 35% se vuelven no solo útiles, sino críticos para la economía. Esto reduce el costo de entrenar modelos, acelera el despliegue en producción y abre posibilidades para aplicaciones que antes eran simplemente poco rentables. Ya hay empresas que se apresuraron a integrar FlashAttention-3 en sus sistemas y obtuvieron ahorros notables en sus facturas. Para los investigadores también es una buena noticia: pueden experimentar más rápido, probar más variantes de arquitecturas y tamaños de modelos.

*Meta ha sido reconocida como una organización extremista y está prohibida en Rusia.