FlashAttention-3: Attention más rápido y preciso con asincronía y baja precisión

 (together.ai)
1 puntos por GN⁺ 2024-07-12 | 1 comentarios | Compartir por WhatsApp
  • El cuello de botella de Attention en los Transformer ha limitado la velocidad de entrenamiento e inferencia de los LLM con contexto largo, y FlashAttention-3 es una nueva versión que busca reducirlo aún más aprovechando funciones de las GPU Hopper
  • La clave es usar la asincronía de Tensor Core y TMA para superponer cómputo y movimiento de datos, e intercalar matmul y softmax para reducir el tiempo ocioso de la GPU
  • En H100, FlashAttention-2 solo aprovechaba el 35% de los FLOPS teóricos máximos, pero FlashAttention-3 alcanza hasta 740 TFLOPS en FP16, llegando al 75% del valor teórico de H100
  • En baja precisión FP8, los valores atípicos en las activaciones de los LLM amplifican los errores, por lo que usa incoherent processing basado en transformada de Hadamard para reducir el error en 2.6 veces frente al attention FP8 de referencia
  • FlashAttention-3 ya está publicado en GitHub, y cuanto más se reduzca el costo de Attention, mayor será el margen para procesar contextos más largos y mejorar la eficiencia del entrenamiento e inferencia de LLM

Objetivos y rendimiento de FlashAttention-3

  • Attention es la capa central de los Transformer, pero se ha convertido en un cuello de botella principal en modelos de lenguaje grandes y aplicaciones de contexto largo
  • FlashAttention y FlashAttention-2 abrieron el camino para acelerar Attention reduciendo las lecturas y escrituras en memoria de la GPU, y hoy la mayoría de las librerías los usan para acelerar el entrenamiento y la inferencia de Transformer
  • Este enfoque ha contribuido a que la longitud de contexto de los LLM pase en los últimos 2 años de 2~4K en GPT-3 y OPT a 128K en GPT-4, y hasta 1M en Llama 3
  • FlashAttention-2 se quedaba en un 35% de utilización de los FLOPS teóricos máximos en GPU H100, pero FlashAttention-3 eleva ese nivel aprovechando nuevas funciones de las GPU Hopper
  • El rendimiento publicado de FlashAttention-3 es el siguiente
    • Entre 1.5 y 2.0 veces más rápido que FlashAttention-2 en FP16
    • Hasta 740 TFLOPS en FP16
    • 75% de utilización de los FLOPS teóricos máximos de H100
    • Casi 1.2 PFLOPS en FP8
    • 2.6 veces menos error que el attention FP8 de referencia

Repaso del enfoque de FlashAttention

  • FlashAttention acelera el cálculo de Attention reorganizando el orden del cómputo y usando tiling y recálculo, además de reducir el uso de memoria respecto de la longitud de secuencia de cuadrático a lineal
  • Carga bloques de entrada desde HBM hacia SRAM, ejecuta Attention sobre ese bloque y luego actualiza la salida en HBM
  • Como no escribe la gran matriz intermedia de Attention en HBM, reduce las lecturas y escrituras de memoria, lo que permite una mejora real de 2 a 4 veces en tiempo de ejecución
  • Al combinar tiling y rescaling de softmax, puede procesar por bloques y aun así obtener la salida correcta sin aproximaciones

Funciones de GPU Hopper: WGMMA, TMA, FP8

  • FlashAttention-2 puede llegar hasta el 70% de los FLOPS teóricos máximos en GPU Ampere A100, pero no aprovecha suficientemente las nuevas funciones de las GPU Hopper
  • FlashAttention-3 usa tres funciones de Hopper
    • WGMMA: función de multiplicación-acumulación matricial por warpgroup que aprovecha los nuevos Tensor Core de Hopper y ofrece más rendimiento que mma.sync de Ampere
    • TMA: unidad de hardware dedicada que acelera la transferencia de datos entre memoria global y memoria compartida, manejando el cálculo de índices y la predicación fuera de rango para reducir el uso de registros
    • FP8: puede duplicar el rendimiento de Tensor Core frente a FP16, pero al representar valores de punto flotante con menos bits implica un tradeoff de precisión
  • FlashAttention-3 aprovecha las funciones de Hopper usando las abstracciones de NVIDIA CUTLASS
  • Solo con reescribir FlashAttention para usar estas nuevas funciones, el rendimiento del forward pass en FP16 mejora desde unos 350 TFLOPS en FlashAttention-2 hasta alrededor de 540~570 TFLOPS

Superponer GEMM y softmax con asincronía

  • Las operaciones principales de Attention son GEMM entre Q-K y P-V, además de softmax
  • En aceleradores modernos, las operaciones que no son matmul son mucho más lentas que matmul, y funciones especiales como la exponencial de softmax se procesan en unidades separadas de las de multiplicación y suma en punto flotante o matrix multiply-add
  • H100 SXM5 ofrece 989 TFLOPS en multiplicación matricial FP16, pero el rendimiento para funciones especiales es de 3.9 TFLOPS, 256 veces menor
  • Con dimensión de cabeza 128, aunque los FLOPS de matmul sean 512 veces mayores que los de la exponencial, la exponencial puede ocupar el 50% del tiempo frente a matmul
  • En FP8, los FLOPS de matmul se duplican pero la velocidad de la exponencial no cambia, así que se vuelve más importante ejecutar matmul y softmax en paralelo

  • Programación pingpong entre warpgroups

    • El scheduler de warps de la GPU ejecuta automáticamente algunos overlap al correr otros warps mientras ciertos warps esperan resultados de GEMM
    • FlashAttention-3 usa barreras de sincronización para superponer manualmente mejor el GEMM y el softmax de dos warpgroups
    • El warpgroup 1 ejecuta primero GEMM1 de una iteración y GEMM0 de la siguiente
    • Luego, mientras el warpgroup 2 ejecuta GEMM, el warpgroup 1 procesa softmax
    • Este esquema pingpong oculta el softmax detrás del tiempo de ejecución de GEMM del otro warpgroup
    • La programación real no es tan perfectamente limpia como en el diagrama, pero en el forward pass de attention FP16, con dimensión de cabeza 128 y longitud de secuencia de 8K, eleva el rendimiento de unos 570 TFLOPS a 620 TFLOPS

  • Overlap dentro del warpgroup

    • Incluso dentro de un mismo warpgroup, puede ejecutarse parte del softmax mientras corre el GEMM de ese warpgroup
    • Este pipelining eleva el rendimiento del forward de attention FP16 desde unos 620 TFLOPS hasta 640~660 TFLOPS
    • A cambio, aumenta la presión sobre registros, porque hay que mantener al mismo tiempo el acumulador de GEMM y las entradas y salidas de softmax
    • En conjunto, esta técnica ofrece un tradeoff favorable

Baja precisión FP8 e incoherent processing

  • Las activaciones de los LLM pueden contener outliers con una magnitud mucho mayor que el resto de las features
  • Los outliers dificultan la cuantización y aumentan mucho el error de cuantización
  • FlashAttention-3 aprovecha incoherent processing, usado en literatura de cuantización como QuIP
  • Multiplica query y key por una matriz ortogonal aleatoria para dispersar los outliers y reducir el error de cuantización
  • En la implementación se usa una transformada de Hadamard con signo aleatorio
    • Si la dimensión de cabeza es d, puede aplicarse por cabeza de attention en tiempo O(d log d) en lugar de O(d²)
    • La transformada de Hadamard está limitada por ancho de banda de memoria, así que puede fusionarse con una operación previa también limitada por ancho de banda, como rotary embedding, sin costo adicional
  • En un experimento donde Q, K y V se generan a partir de una distribución normal estándar y se insertan magnitudes grandes en el 0.1% de las entradas para simular outliers, incoherent processing reduce el error de cuantización en 2.6 veces

Benchmarks y estado de publicación

  • FlashAttention-3 se compara no solo con FlashAttention-2, sino también con implementaciones de Triton y cuDNN que ya usan nuevas funciones de hardware de las GPU Hopper
  • En FP16 muestra una mejora de velocidad de alrededor de 1.6~1.8 veces frente a FlashAttention-2
  • En FP8 alcanza casi 1.2 PFLOPS
  • Ya está publicado el repositorio de GitHub de FlashAttention-3
  • El paper también puede consultarse en el mismo repositorio flash-attention

Optimizaciones pendientes e integración futura

  • Además de lo tratado en el blog, el paper incluye optimizaciones como variable length sequence, persistent kernel e in-kernel transpose para FP8
  • Diseñar el algoritmo de acuerdo con el hardware de ejecución puede generar grandes mejoras de eficiencia y habilitar nuevas capacidades del modelo, como contexto más largo
  • El trabajo futuro incluye optimización para inferencia de LLM y generalización de estas técnicas a otras arquitecturas de hardware
  • Se espera que FlashAttention-3 se integre en futuras versiones de PyTorch

Lecturas relacionadas

1 comentarios

 
GN⁺ 2024-07-12
Opiniones de Hacker News

  • Por los comentarios en el código, parece que Tri Dao venía trabajando en FA3 desde abril de 2022, justo después del anuncio de Hopper/H100.
    Es un poco interesante que hayan pasado más de dos años hasta que el código se publicara hoy; quizá sea porque ya se están preparando mejores soluciones.
    El historial reciente de papers de Tri se inclina hacia las arquitecturas de tipo SSM y Mamba. FlashAttention tiene complejidad temporal cuadrática respecto de la longitud de la secuencia, pero los algoritmos más recientes son subcuadráticos, así que no solo hacen el mismo cálculo de forma más eficiente, sino que reducen muchísimo la cantidad de cómputo en sí.
    Dao y Gu mostraron este año, en un paper largo, que Mamba/SSM también puede formularse de manera que sea fácil de acelerar con las mismas primitivas de hardware de las que se beneficia Transformer.

    • Hasta que se demuestre o refute la Hipótesis fuerte del tiempo exponencial (SETH), hace falta pagar un costo cuadrático o renunciar a algo. Al final, es el costo de la búsqueda exhaustiva.
      Si se demuestra o refuta SETH, también se resolvería el problema P versus NP, así que es difícil esperar que ocurra pronto.
      La clave es si un caso de uso específico puede asumir ese costo.

  • Me pregunto qué tan atado al hardware está el algoritmo FlashAttention.
    Por ejemplo, en este anuncio dicen que aprovecha las funciones asíncronas de las GPU H100, así que eso parece significar que en tarjetas que no sean de la serie H no se obtendría esa mejora de velocidad.
    Además, la biblioteca real de FlashAttention requiere CUDA, pero parece que el algoritmo fue portado a Metal[^0]. Si el algoritmo es algo cercano a una función pura, debería poder implementarse en cualquier GPU/framework de machine learning, ¿no?
    [0]: https://github.com/philipturner/metal-flash-attention

    • Hay muchas buenas respuestas, pero en resumen: “en la práctica, bastante” atado al hardware. Lo siguiente es un buen ejemplo:

      https://github.com/karpathy/nanoGPT/blob/master/model.py#L45
      nanoGPT de Karpathy verifica si existe torch.nn.functional.scaled_dot_product_attention y llama a FlashAttention.
      https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
      Si ves la documentación, en la práctica uno termina queriendo que llame sobre todo a FA2, y FA2 optimiza los kernels del dispositivo para dividir la operación Softmax de la matriz triangular y reducir los viajes innecesarios de lotes de punto flotante entre la GPU y la CPU.
      https://arxiv.org/pdf/2307.08691
      El paper de FA2 está descrito casi por completo desde la perspectiva del hardware en el que se ejecuta.

    • La mejora algorítmica de FlashAttention consiste principalmente en dividir y recombinar la parte Softmax de la atención, y eso en sí no es una idea totalmente nueva. La contribución abrumadora está en haber implementado ese método y sus detalles de forma eficiente en hardware de Nvidia.
    • Originalmente, FlashAttention casi no dependía del hardware.
      En las versiones más recientes depende del nivel de abstracción. ThunderKittens[0] ofrece mejoras de velocidad de alrededor de 1.3 a 2 veces frente a FA2, de forma similar a lo mencionado en el artículo, y aun así se aplica de manera relativamente general en distintas GPU.
      Cada nuevo hardware puede traer funciones específicas que permitan extraer rendimiento adicional. Por lo general, los vendors adoptan funciones que los ponen por delante, pero, como ya ocurre con CUDA, las API y bibliotecas se fragmentan.
      [0]: https://hazyresearch.stanford.edu/blog/2024-05-12-tk
    • Conceptualmente, está un poco atado; desde la perspectiva de una implementación práctica, mucho. Incluso la implementación estándar en Python compila internamente kernels adaptados a hardware específico.
    • Sumando desde el punto de vista práctico: el hardware de AMD todavía carece de una implementación sólida de flash-attention-2. ROCm se está volviendo utilizable lentamente, pero aún no está al nivel de CUDA.

  • Quisiera preguntarles a quienes trabajan en compiladores: ¿existe la posibilidad de que un compilador encuentre por sí solo optimizaciones como FlashAttention? TVM y tinygrad parecen ir en esa dirección, pero cuesta creer que sea posible.

    • En teoría, sí es posible. Gracias a las propiedades algebraicas de las matemáticas se puede reordenar mucho, y a eso se le puede sumar un tiling de bucles poliédricos relativamente estructurado.
      Pero, como el costo es alto, habría que cachear los resultados de esa búsqueda.
      La optimización con e-graphs parece encajar bien en esta área. Sin embargo, requiere un gran cambio de paradigma en la forma de procesar los pases de optimización, así que casi no se ha desplegado fuera de algunas herramientas de nicho. Por ejemplo, no encaja bien con los grafos de llamadas tradicionales; para desplegar e-graphs atravesando dentro, fuera y entre bloques básicos y bucles for, habría que cambiar mucho el flujo de control, y tampoco se admiten break ni return.
    • Parece un problema extremadamente difícil, pero no imposible.
      No sé hasta dónde ha llegado el estado del arte en optimización de compiladores en términos de disposición de datos y maximización del uso del procesador.
      Recuerdo haber visto en un video sobre optimización que una pequeña optimización sí aumentaba la velocidad, pero que su impacto era mínimo comparado con la variación de velocidad causada por las diferencias en la disposición de memoria que producían esa optimización, o incluso cambios aleatorios.
      Esa charla se enfocaba más en distinguir la señal dentro del ruido, pero ese ruido en sí es una señal de que los compiladores ni siquiera manejan bien formas mucho más simples del problema del que se habla aquí.
      La arquitectura CPU-memoria, donde las cachés y los patrones de acceso afectan la velocidad, ya es compleja; si encima se suma la arquitectura de GPU, parece un terreno bastante inexplorado.
      Quizá algún día sea posible. Como esto pertenece al campo de la IA, también surge la pregunta de si una IA lo bastante inteligente podría hacerlo, pero depende de qué signifique “lo bastante”.
      Como una prueba de muy alto nivel para un modelo de IA, se podría imaginar darle algo como micrograd y pedirle que cree algo más rápido que torch manteniendo la misma interfaz. Todavía no estamos ni cerca de eso, pero sería interesante si llegara a ser posible.
    • Yo diría que no. Hay que pensarlo como un algoritmo distinto. En vez de considerar solo las matemáticas, se diseña el algoritmo teniendo en cuenta la forma del hardware.
      TVM tiene sentido. Estrictamente hablando hace otra cosa, pero está en un área bastante cercana.
      Pero no sé por qué tinygrad te dio esa impresión.
    • https://github.com/uwplse/tensat
    • Llamar operadores de alto nivel desde un lenguaje envoltorio como Python es bastante complicado.

  • Si alguien quiere portar esto a ROCm / AMD MI300x, que escriba a hello@hotaisle.xyz. No enviaré spam, jamás.
    Puedo donar tiempo de cómputo para este trabajo.

    • ¡Son una empresa de servidores con aceleradores AMD! Gran trabajo; ojalá alguien tome la posta :)
    • No quiero ser grosero, pero me da curiosidad la intención de esta oferta. ¿Alguien haría este port gratis solo por tener acceso al hardware? ¿Qué gana esa persona?

  • FlashAttention-3 is optimized for Hopper GPUs (e.g. H100).
    ¿Qué nivel de rendimiento ofrece FA3 en GPU de consumo como las 3090 y 4090?

  • TMA (Tensor Memory Accelerator). This is a special hardware unit that accelerates the transfer of data between global memory and shared memory, taking care of all index calculation and out-of-bound predication. This frees up registers, which is a valuable resource to increase tile size and efficiency.
    Según entiendo, es cierto que TMA reduce el uso de registros, pero lo más importante es que permite que el hardware se encargue de la generación de direcciones. A medida que las operaciones alrededor se vuelven más rápidas, la generación de direcciones puede convertirse en un cuello de botella.

  • Esta es una de las mejoras más importantes en toda la IA. Permite usar más cosas más rápido con el mismo hardware y ofrece beneficios casi sin concesiones para la mayoría de los usuarios de IA.

    • Para los usuarios que tengan una H100, claro.

  • Me pregunto por qué FlashAttention es unas 5 veces más lento cuando usa enmascaramiento variable que cuando no lo usa. Si no hay buen soporte para masking, el efecto de la optimización casi desaparece.

    • ¿Dónde estás viendo ese benchmark?

  • Me gustaría que algún experto respondiera algunas preguntas :)
    ¿FlashAttention es un reemplazo drop-in de la operación de atención en los LLM? ¿Se puede usar en cualquier lugar donde se use una operación de “atención”, o hay que entrenar el LLM por separado para que use FA?
    ¿Qué relación tiene FA con estrategias como GQA y la atención de ventana deslizante? ¿Son conceptos ortogonales entre sí, o se necesita una implementación de FA distinta para cada estrategia?
    Hace poco llama.cpp agregó soporte para FlashAttention; ¿eso significa que empezó a usar algo como los kernels CUDA que ofrece FlashAttention?
    Por último, en este artículo se compara FlashAttention con Triton. ¿Triton no es algo así como una capa de abstracción? ¿No se puede implementar FA con Triton? No termino de entender la expresión “FlashAttention vs. Triton”.

      1. Es casi correcto. Son matemáticamente equivalentes. En software, los únicos problemas son cosas como la gestión de versiones de dependencias o el formato de los datos en memoria, y FlashAttention 2 ya está integrado en HuggingFace y varias bibliotecas populares. Es muy probable que FlashAttention 3 también se integre pronto, pero para ejecutarlo se necesita una GPU H100.
      2. FlashAttention 2 agregó soporte para GQA en una actualización anterior:
        https://github.com/Dao-AILab/flash-attention
      3. Aquí se compara esta implementación de FlashAttention escrita en CUDA C++ puro con una implementación en Triton de un algoritmo similar escrita en Triton: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html
    • FlashAttention puede usarse como reemplazo drop-in de la operación de atención de un LLM.
      FlashAttention es una forma de calcular la parte Softmax(QK^T)V de la atención, mientras que GQA es una forma de calcular las matrices Q, K y V. No estoy tan seguro sobre la atención de ventana deslizante, pero es una forma de cambiar la máscara de atención que controla qué claves puede ver cada consulta.
      No he usado llama.cpp, pero la explicación de que empezó a usar kernels CUDA en general parece correcta.
      La última pregunta se refiere a una implementación de FlashAttention escrita previamente en Triton.

  • Me da curiosidad que el artículo diga que operaciones como sigmoid son muy lentas.
    Los LLM modernos usan muchas funciones de activación que incluyen sigmoid o Softmax, como SiLU, Swish y SOLU.
    ¿ReLU tiene menos pérdida de rendimiento? Si es así, ¿tal vez convendría volver al buen viejo ReLU?

    • ReLU es literalmente una función lineal que se corta en 0 en cierto punto, así que requiere mucho menos cómputo que las que incluyen funciones exponenciales. Dicho eso, parece difícil obtener resultados competitivos con una función de activación tan simple.