Inicio>Relacionado con la programación>Código Fuente de IA
Descargar

Sumérjase en el aprendizaje profundo, rehecho por la revista Quanta

Atención de similitud de coseno flash

Implementación de atención de similitud de coseno fusionado en el mismo estilo que Flash Attention. La observación es que al adoptar claves y consultas normalizadas l2, ya no es necesario realizar un seguimiento de los máximos de fila para lograr estabilidad numérica. Esto simplifica enormemente el algoritmo de atención flash, asumiendo que la atención por similitud de coseno no tiene costo de generalización.

En otras palabras, atención al contexto estable, rápida, eficiente en memoria y más prolongada sin inconvenientes.

Actualización: Desafortunadamente, los experimentos de Robin mostraron puntuaciones de evaluación FID mucho peores que no se reflejan en la pérdida. Pendiente de más experimentos. Utilice esta biblioteca con precaución.

Actualización 2: La única salvación sería utilizar l2norm agrupado, lo que potencialmente podría permitir una mayor expresividad. Se lo agradecería mucho si alguien pudiera evaluar esta técnica en su trabajo generativo y obtener algunas puntuaciones FID.

Actualización 3: Se ha probado a escala un enfoque similar a la atención de simulación de coseno, con un modelo de visión de parámetros 22B de Brain.

Estado (wip)

Por el momento, las secuencias autorregresivas y de longitud variable deberían ser más rápidas en todas las arquitecturas. Para secuencias de más de 2048, también será eficiente en memoria donde la atención regular no lo sería.

Sin embargo, para los no autorregresivos sin enmascaramiento, la arquitectura es aún más lenta en A100 para F16. El objetivo es lograr que funcione más rápido en A100 hacia adelante y hacia atrás tanto para F32 como para F16, ya que la memoria compartida aún no está explotada por completo.

Con tarjetas gráficas más antiguas y sin suficiente memoria compartida, habrá que evaluar la compensación entre la eficiencia y la velocidad de la memoria dependiendo de la longitud de la secuencia a la que se entrena.

Apreciación

  • Arthur Hennequin por guiarme a través de mi primer kernel CUDA y por codificar una implementación de referencia simple, que me ayudó a iniciar el primer kernel que tiene un rendimiento razonable hasta la línea de base. Este trabajo no habría sido posible sin su experiencia.

  • Boris Dayma y Robin Rombach por realizar experimentos con la atención de simulación de coseno simplificada con escala fija en algunos modelos importantes de texto a imagen y por verificar que realmente funciona tan bien como la atención normal.

  • Markus Rabe por escribir el artículo que demostró que la atención no requiere memoria O(n²), y Tri Dao por ponerlo todo junto en una implementación del kernel CUDA para una atención regular, demostrando superioridad en velocidad utilizando el enfoque en mosaico que minimiza los accesos a HBM (y por calcular out dO * O == dP * P para pase hacia atrás). No habría podido completar mi peregrinaje en busca de la fórmula de atención definitiva sin sus descubrimientos.

  • Stability.ai por el generoso patrocinio para trabajar en investigaciones de vanguardia en inteligencia artificial

Instalar 

$ pip install flash-cosine-sim-attention

Uso

Auto atención 

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention

q = torch . randn ( 1 , 8 , 1024 , 64 ). cuda ()
k = torch . randn ( 1 , 8 , 1024 , 64 ). cuda ()
v = torch . randn ( 1 , 8 , 1024 , 64 ). cuda ()

out = flash_cosine_sim_attention ( q , k , v )  # (1, 8, 1024, 64)

atención cruzada 

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention

q = torch . randn ( 1 , 8 , 1024 , 64 ). cuda ()
k = torch . randn ( 1 , 8 , 2048 , 64 ). cuda ()
v = torch . randn ( 1 , 8 , 2048 , 64 ). cuda ()

out = flash_cosine_sim_attention ( q , k , v ) # (1, 8, 1024, 64)

Con enmascaramiento de clave/valor 

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention

q = torch . randn ( 1 , 8 , 1024 , 64 ). cuda ()
k = torch . randn ( 1 , 8 , 2048 , 64 ). cuda ()
v = torch . randn ( 1 , 8 , 2048 , 64 ). cuda ()

mask = torch . ones ( 1 , 2048 ). bool (). cuda ()

out = flash_cosine_sim_attention ( q , k , v , mask = mask ) # (1, 8, 1024, 64)

autorregresivo 

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention

q = torch . randn ( 4 , 8 , 1024 , 64 ). cuda ()
k = torch . randn ( 4 , 8 , 1024 , 64 ). cuda ()
v = torch . randn ( 4 , 8 , 1024 , 64 ). cuda ()

out = flash_cosine_sim_attention ( q , k , v , causal = True )  # (4, 8, 1024, 64)

Misceláneas

Claves/valores de un solo cabezal (Shazeer et al y utilizados en PaLM) 

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention

q = torch . randn ( 4 , 8 , 1024 , 64 ). cuda ()
k = torch . randn ( 4 , 1024 , 64 ). cuda ()
v = torch . randn ( 4 , 1024 , 64 ). cuda ()

out = flash_cosine_sim_attention ( q , k , v , causal = True )  # (4, 8, 1024, 64)

Si necesita realizar operaciones en las consultas y claves entre l2norm y el paso de atención real, simplemente configure l2norm_qk = False

ex. 

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention , l2norm_tensors

q = torch . randn ( 4 , 8 , 1024 , 64 ). cuda ()
k = torch . randn ( 4 , 1024 , 64 ). cuda ()
v = torch . randn ( 4 , 1024 , 64 ). cuda ()

q , k = l2norm_tensors ( q , k )

# do your rotation of queries and keys
# say with https://github.com/lucidrains/rotary-embedding-torch

out = flash_cosine_sim_attention ( q , k , v , l2norm_qk = False )  # (4, 8, 1024, 64)

Atención cruzada con trabajos causales como se esperaba (almacenamiento en caché de claves y valores en autorregresivo durante la inferencia, o entrenamiento tipo transformador-xl) 

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention

q = torch . randn ( 1 , 8 , 1024 , 64 ). cuda ()
k = torch . randn ( 1 , 8 , 2048 , 64 ). cuda ()
v = torch . randn ( 1 , 8 , 2048 , 64 ). cuda ()

out = flash_cosine_sim_attention ( q , k , v , causal = True ) # (1, 8, 1024, 64)

Si tiene las dimensiones del lote y del cabezal fusionadas, está bien. 

 import torch
from flash_cosine_sim_attention import flash_cosine_sim_attention

q = torch . randn ( 32 , 1024 , 64 ). cuda ()
k = torch . randn ( 32 , 2048 , 64 ). cuda ()
v = torch . randn ( 32 , 2048 , 64 ). cuda ()

out = flash_cosine_sim_attention ( q , k , v , causal = True ) # (32, 1024, 64)

Dimensiones del cabezal soportado

  • 16-f32

  • 32

  • 64

  • 96

  • 128

  • 16-f16

  • 80 - en progreso

Hacer

  • Soporte bfloat16, use sfinae según lo recomendado por Arthur

  • transmita desde qk_mma a la memoria compartida en fragmentos para calcular mma, vea si el smem liberado se puede usar para almacenar en caché más

  • soporte O(n) 1d sesgo posicional dinámico

  • descubrir por qué el almacenamiento en caché de fragmentos smem provocaría una degradación del rendimiento, no tiene sentido

  • Piense en el uso de logsumexp: funciona, pero el registro adicional provoca un rendimiento degradado.

  • preparar un mecanismo de almacenamiento en caché de fragmentos smem, para permitir tanto almacenamiento en caché como esté permitido en A100 (o f16)

  • haga que el procesamiento del tamaño del mosaico de atención sea personalizable para pasar hacia atrás

  • mover agregado atómico a la función sobrecargada dentro de mma

  • flexible qué tipo se utiliza para la acumulación

  • Pruebe mosaicos de 64x96 en f16.

  • incorpore una versión eficiente de la memoria de la CPU (solo para inferencia, ya que el entrenamiento no tiene sentido) usando simplemente código pytorch

  • descubrir cómo distribuir de manera diferente para las arquitecturas (digamos A100), en caso de que al revés se pueda hacer uso del aumento en la memoria compartida de manera diferente

  • desacople los tamaños de filas y columnas para mosaicos de atención

  • dk y dv ahora están en f16 cuando puede ser (kv sin una sola cabeza)

  • admite dimensiones de cabeza más estándar (wip)

  • depurar y corregir gradientes sesgados hacia atrás una vez más para un tamaño de cabeza de 32

  • corregir gradientes de sesgo de atención

  • permitir claves/valores de un solo encabezado, como en PaLM

  • arreglar el complemento atómico para f16

  • El sesgo de atención debería poder aceptar dimensiones de una dimensión de lote adicional, para Alphafold2 como el sesgo de atención.

  • automatizar la eliminación de caché del kernel usando la versión como sufijo del nombre del paquete

  • resolver problemas numéricos causales f16

  • adopte todos los aprendizajes desde el kernel hacia adelante hasta el kernel hacia atrás y asegúrese de que tenga un rendimiento superior al menos en A100

Descripción

Hasta ahora, la atención a la similitud del coseno no se utiliza ampliamente en la industria. El único modelo grande que se ha entrenado con él hasta ahora es SwinV2. Si alguien puede invalidar el enfoque, abra un problema o envíeme un correo electrónico. Puede ejecutar experimentos con atención regular utilizando el repositorio de x-transformers.

Actualización: Boris Dayma gentilmente inició un experimento (azul con rojo como línea de base) para validar la atención de similitud de cosenos con una escala fija de 10 en un entorno de modelo del mundo real.

Actualización 2: La atención por similitud de coseno se ha demostrado en una red de atención de texto a imagen del mundo real, utilizando una escala constante de 10 . No peor que la atención habitual. El crédito es para Boris Dayma por invertir el tiempo para ejecutar el experimento y eliminar dudas sobre la técnica.

Actualización 3: Robin Rombach ha probado el kernel en este repositorio con un tamaño de cabeza de 64 y una escala fija de 10 en un modelo de texto a imagen, sin observar ninguna diferencia con la atención normal. Más evaluaciones pendientes.

Actualización 4: La mejora en el rendimiento observada en los experimentos de Boris probablemente se deba al hecho de que la atención de coseno-sim permite cambiar de la configuración pre-layernorm a post-layernorm en los transformadores (ya que l2norm efectivamente toma el lugar de la pre-layernorm). norma de capa). La atención de simulación de coseno probablemente producirá los mismos resultados que la atención normal, sin ningún otro cambio en el transformador.

Pruebas

Para probar la salida y los gradientes son iguales para escenarios autorregresivos y no autorregresivos 

$ python setup.py test

Evaluación comparativa

Asegúrese de instalar primero el kernel CUDA 

$ python setup . py install

Entonces 

$ python benchmark . py

Para realizar evaluaciones comparativas solo hacia adelante o hacia atrás, agregue el indicador --only-forwards o --only-backwards a lo anterior. Para comparar autorregresivo, agregue --causal

Puntos de referencia - wip

GTX 2080Ti

Adelante 

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 1.05x   kernel: 0.24ms  baseline: 0.23ms
seq_len: 256    slower: 1.27x   kernel: 0.38ms  baseline: 0.30ms
seq_len: 512    slower: 1.28x   kernel: 0.87ms  baseline: 0.68ms
seq_len: 1024   slower: 1.15x   kernel: 2.63ms  baseline: 2.28ms
seq_len: 2048   slower: 0.99x   kernel: 7.99ms  baseline: 8.10ms
seq_len: 4096   slower: 0.88x   kernel: 30.82ms baseline: 34.84ms
seq_len: 8192   slower: 0.00x   kernel: 121.96ms    baseline: oom
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.85x   kernel: 0.20ms  baseline: 0.24ms
seq_len: 256    slower: 0.97x   kernel: 0.24ms  baseline: 0.25ms
seq_len: 512    slower: 1.22x   kernel: 0.43ms  baseline: 0.35ms
seq_len: 1024   slower: 0.95x   kernel: 0.93ms  baseline: 0.98ms
seq_len: 2048   slower: 0.90x   kernel: 3.16ms  baseline: 3.50ms
seq_len: 4096   slower: 0.85x   kernel: 11.06ms baseline: 13.07ms
seq_len: 8192   slower: 0.00x   kernel: 42.61ms baseline: oom

Al revés: todavía necesita trabajo 

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 1.07x   kernel: 0.61ms  baseline: 0.57ms
seq_len: 256    slower: 1.40x   kernel: 0.91ms  baseline: 0.65ms
seq_len: 512    slower: 1.70x   kernel: 2.34ms  baseline: 1.38ms
seq_len: 1024   slower: 1.26x   kernel: 5.67ms  baseline: 4.50ms
seq_len: 2048   slower: 1.29x   kernel: 20.60ms baseline: 15.91ms
seq_len: 4096   slower: 1.30x   kernel: 78.93ms baseline: 60.81ms
seq_len: 8192   slower: 0.00x   kernel: 314.51ms    baseline: oom
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.91x   kernel: 0.50ms  baseline: 0.55ms
seq_len: 256    slower: 1.06x   kernel: 0.58ms  baseline: 0.55ms
seq_len: 512    slower: 1.13x   kernel: 0.81ms  baseline: 0.72ms
seq_len: 1024   slower: 0.97x   kernel: 2.09ms  baseline: 2.16ms
seq_len: 2048   slower: 0.96x   kernel: 7.06ms  baseline: 7.35ms
seq_len: 4096   slower: 0.97x   kernel: 26.08ms baseline: 26.84ms
seq_len: 8192   slower: 0.00x   kernel: 101.02ms    baseline: oom

Adelante y atrás: F32 es definitivamente más lento 

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 1.05x   kernel: 0.83ms  baseline: 0.79ms
seq_len: 256    slower: 1.34x   kernel: 1.26ms  baseline: 0.95ms
seq_len: 512    slower: 1.44x   kernel: 3.14ms  baseline: 2.18ms
seq_len: 1024   slower: 1.15x   kernel: 7.83ms  baseline: 6.81ms
seq_len: 2048   slower: 1.20x   kernel: 28.83ms baseline: 24.03ms
seq_len: 4096   slower: 1.20x   kernel: 111.13ms    baseline: 92.51ms
seq_len: 8192   slower: 0.00x   kernel: 441.70ms    baseline: oom
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 0.89x   kernel: 0.68ms  baseline: 0.77ms
seq_len: 256    slower: 1.03x   kernel: 0.80ms  baseline: 0.77ms
seq_len: 512    slower: 1.06x   kernel: 1.16ms  baseline: 1.10ms
seq_len: 1024   slower: 0.93x   kernel: 2.94ms  baseline: 3.16ms
seq_len: 2048   slower: 0.93x   kernel: 10.06ms baseline: 10.87ms
seq_len: 4096   slower: 0.93x   kernel: 37.09ms baseline: 39.96ms
seq_len: 8192   slower: 0.00x   kernel: 143.13ms    baseline: oom

Para autorregresivo, una victoria clara python benchmark.py --causal 

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 0.97x   kernel: 0.81ms  baseline: 0.84ms
seq_len: 256    slower: 1.07x   kernel: 1.12ms  baseline: 1.05ms
seq_len: 512    slower: 0.83x   kernel: 2.23ms  baseline: 2.68ms
seq_len: 1024   slower: 0.55x   kernel: 4.83ms  baseline: 8.82ms
seq_len: 2048   slower: 0.49x   kernel: 15.89ms baseline: 32.68ms
seq_len: 4096   slower: 0.46x   kernel: 57.50ms baseline: 126.00ms
seq_len: 8192   slower: 0.00x   kernel: 224.76ms    baseline: oom
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 0.82x   kernel: 0.69ms  baseline: 0.84ms
seq_len: 256    slower: 0.95x   kernel: 0.79ms  baseline: 0.83ms
seq_len: 512    slower: 0.78x   kernel: 1.06ms  baseline: 1.37ms
seq_len: 1024   slower: 0.50x   kernel: 2.10ms  baseline: 4.24ms
seq_len: 2048   slower: 0.37x   kernel: 5.85ms  baseline: 15.92ms
seq_len: 4096   slower: 0.31x   kernel: 19.80ms baseline: 64.42ms
seq_len: 8192   slower: 0.00x   kernel: 75.25ms baseline: oom

Para secuencias de longitud variable con enmascaramiento, también una clara victoria. Supongamos que en promedio el 25% de los tokens están enmascarados python benchmark.py --mask-prob 0.25 

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.95x   kernel: 0.84ms  baseline: 0.89ms
seq_len: 256    slower: 1.19x   kernel: 1.28ms  baseline: 1.08ms
seq_len: 512    slower: 1.23x   kernel: 3.19ms  baseline: 2.59ms
seq_len: 1024   slower: 0.92x   kernel: 8.19ms  baseline: 8.88ms
seq_len: 2048   slower: 0.92x   kernel: 30.08ms baseline: 32.57ms
seq_len: 4096   slower: 0.94x   kernel: 123.20ms    baseline: 131.22ms
seq_len: 8192   slower: 0.00x   kernel: 461.77ms    baseline: oom
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.85x   kernel: 0.77ms  baseline: 0.90ms
seq_len: 256    slower: 0.93x   kernel: 0.86ms  baseline: 0.93ms
seq_len: 512    slower: 0.93x   kernel: 1.31ms  baseline: 1.40ms
seq_len: 1024   slower: 0.76x   kernel: 3.31ms  baseline: 4.35ms
seq_len: 2048   slower: 0.71x   kernel: 11.19ms baseline: 15.65ms
seq_len: 4096   slower: 0.70x   kernel: 41.27ms baseline: 59.01ms
seq_len: 8192   slower: 0.00x   kernel: 158.60ms    baseline: oom

A100 40GB (wip)

Gracias a Stability por brindar acceso a los A100 para realizar pruebas. Gracias a Enrico por tomarse el tiempo de ejecutar algunos puntos de referencia cuando todavía no tenía acceso.

A100 es todavía un trabajo en progreso. La memoria compartida aún no está totalmente explotada. Curiosamente, el F32 parece funcionar mejor que el F16.

Hacia adelante 

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.98x   kernel: 0.29ms  baseline: 0.30ms
seq_len: 256    slower: 1.19x   kernel: 0.35ms  baseline: 0.29ms
seq_len: 512    slower: 0.94x   kernel: 0.52ms  baseline: 0.55ms
seq_len: 1024   slower: 0.75x   kernel: 1.23ms  baseline: 1.65ms
seq_len: 2048   slower: 0.88x   kernel: 4.17ms  baseline: 4.73ms
seq_len: 4096   slower: 0.79x   kernel: 14.53ms baseline: 18.36ms
seq_len: 8192   slower: 0.64x   kernel: 55.01ms baseline: 85.93ms
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.84x   kernel: 0.24ms  baseline: 0.29ms
seq_len: 256    slower: 1.02x   kernel: 0.29ms  baseline: 0.29ms
seq_len: 512    slower: 1.24x   kernel: 0.36ms  baseline: 0.29ms
seq_len: 1024   slower: 1.48x   kernel: 0.79ms  baseline: 0.54ms
seq_len: 2048   slower: 1.31x   kernel: 2.08ms  baseline: 1.59ms
seq_len: 4096   slower: 1.21x   kernel: 6.89ms  baseline: 5.70ms
seq_len: 8192   slower: 1.07x   kernel: 24.80ms baseline: 23.15ms

Hacia atrás 

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.94x   kernel: 0.57ms  baseline: 0.60ms
seq_len: 256    slower: 1.29x   kernel: 0.75ms  baseline: 0.58ms
seq_len: 512    slower: 1.16x   kernel: 1.30ms  baseline: 1.12ms
seq_len: 1024   slower: 0.98x   kernel: 3.14ms  baseline: 3.19ms
seq_len: 2048   slower: 1.05x   kernel: 11.13ms baseline: 10.63ms
seq_len: 4096   slower: 0.98x   kernel: 40.11ms baseline: 40.79ms
seq_len: 8192   slower: 0.97x   kernel: 154.96ms    baseline: 159.70ms
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.91x   kernel: 0.55ms  baseline: 0.60ms
seq_len: 256    slower: 1.03x   kernel: 0.62ms  baseline: 0.60ms
seq_len: 512    slower: 1.36x   kernel: 0.82ms  baseline: 0.60ms
seq_len: 1024   slower: 1.52x   kernel: 1.52ms  baseline: 1.01ms
seq_len: 2048   slower: 1.37x   kernel: 4.14ms  baseline: 3.03ms
seq_len: 4096   slower: 1.33x   kernel: 14.23ms baseline: 10.71ms
seq_len: 8192   slower: 1.34x   kernel: 53.90ms baseline: 40.28ms

Hacia adelante y hacia atrás 

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.92x   kernel: 0.80ms  baseline: 0.87ms
seq_len: 256    slower: 1.23x   kernel: 1.07ms  baseline: 0.87ms
seq_len: 512    slower: 1.08x   kernel: 1.80ms  baseline: 1.66ms
seq_len: 1024   slower: 0.94x   kernel: 4.33ms  baseline: 4.62ms
seq_len: 2048   slower: 0.99x   kernel: 15.26ms baseline: 15.44ms
seq_len: 4096   slower: 0.93x   kernel: 54.78ms baseline: 59.21ms
seq_len: 8192   slower: 0.91x   kernel: 210.38ms    baseline: 230.97ms
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64
------------------------------------------------------------
seq_len: 128    slower: 0.90x   kernel: 0.78ms  baseline: 0.86ms
seq_len: 256    slower: 1.00x   kernel: 0.87ms  baseline: 0.87ms
seq_len: 512    slower: 1.36x   kernel: 1.18ms  baseline: 0.86ms
seq_len: 1024   slower: 1.49x   kernel: 2.31ms  baseline: 1.55ms
seq_len: 2048   slower: 1.33x   kernel: 6.17ms  baseline: 4.63ms
seq_len: 4096   slower: 1.28x   kernel: 21.08ms baseline: 16.44ms
seq_len: 8192   slower: 1.24x   kernel: 78.75ms baseline: 63.45ms

autorregresivo 

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 0.82x   kernel: 0.82ms  baseline: 1.01ms
seq_len: 256    slower: 1.02x   kernel: 1.00ms  baseline: 0.98ms
seq_len: 512    slower: 0.82x   kernel: 1.55ms  baseline: 1.89ms
seq_len: 1024   slower: 0.51x   kernel: 2.79ms  baseline: 5.44ms
seq_len: 2048   slower: 0.45x   kernel: 8.37ms  baseline: 18.67ms
seq_len: 4096   slower: 0.40x   kernel: 29.16ms baseline: 72.97ms
seq_len: 8192   slower: 0.38x   kernel: 108.68ms    baseline: 285.47ms
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 0.82x   kernel: 0.81ms  baseline: 0.98ms
seq_len: 256    slower: 0.90x   kernel: 0.88ms  baseline: 0.98ms
seq_len: 512    slower: 1.16x   kernel: 1.13ms  baseline: 0.97ms
seq_len: 1024   slower: 0.80x   kernel: 1.68ms  baseline: 2.10ms
seq_len: 2048   slower: 0.54x   kernel: 3.66ms  baseline: 6.81ms
seq_len: 4096   slower: 0.45x   kernel: 11.43ms baseline: 25.32ms
seq_len: 8192   slower: 0.41x   kernel: 40.58ms baseline: 99.14ms

Secuencias de longitud variable (hasta un 25% de tokens enmascarados) 

------------------------------------------------------------
float32     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 0.80x   kernel: 0.85ms  baseline: 1.07ms
seq_len: 256    slower: 1.07x   kernel: 1.15ms  baseline: 1.08ms
seq_len: 512    slower: 1.00x   kernel: 1.94ms  baseline: 1.94ms
seq_len: 1024   slower: 0.84x   kernel: 4.64ms  baseline: 5.55ms
seq_len: 2048   slower: 0.84x   kernel: 15.86ms baseline: 18.86ms
seq_len: 4096   slower: 0.76x   kernel: 55.19ms baseline: 72.47ms
seq_len: 8192   slower: 0.75x   kernel: 212.48ms    baseline: 282.71ms
------------------------------------------------------------
float16     batch: 4    heads: 8    dim 64  
------------------------------------------------------------
seq_len: 128    slower: 0.80x   kernel: 0.83ms  baseline: 1.04ms
seq_len: 256    slower: 0.90x   kernel: 0.93ms  baseline: 1.03ms
seq_len: 512    slower: 1.18x   kernel: 1.22ms  baseline: 1.04ms
seq_len: 1024   slower: 1.10x   kernel: 2.40ms  baseline: 2.17ms
seq_len: 2048   slower: 0.89x   kernel: 6.27ms  baseline: 7.06ms
seq_len: 4096   slower: 0.82x   kernel: 21.19ms baseline: 25.95ms
seq_len: 8192   slower: 0.78x   kernel: 79.45ms baseline: 101.83ms

Entrenando un pequeño GPT en Enwik8 

$ make train

Pruebe la longitud de secuencia 8192. Será lento pero funcionará (la atención normal se interrumpirá en > 2048, verá esto si elimina el indicador --use-cuda-kernel ) 

$ python train . py - - seq - len 8192 - - use - cuda - kernel

Citas 

 @article { Dao2022FlashAttentionFA ,
    title   = { FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness } ,
    author  = { Tri Dao and Daniel Y. Fu and Stefano Ermon and Atri Rudra and Christopher R'e } ,
    journal = { ArXiv } ,
    year    = { 2022 } ,
    volume  = { abs/2205.14135 }
}
 @misc { rabe2021selfattention ,
    title   = { Self-attention Does Not Need $O(n^2)$ Memory } , 
    author  = { Markus N. Rabe and Charles Staats } ,
    year    = { 2021 } ,
    eprint  = { 2112.05682 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { cs.LG }
}
 @inproceedings { Henry2020QueryKeyNF ,
    title   = { Query-Key Normalization for Transformers } ,
    author  = { Alex Henry and Prudhvi Raj Dachapally and Shubham Vivek Pawar and Yuxuan Chen } ,
    booktitle = { FINDINGS } ,
    year    = { 2020 }
}
 @article { Wang2022DeepNetST ,
    title   = { DeepNet: Scaling Transformers to 1, 000 Layers } ,
    author  = { Hongyu Wang and Shuming Ma and Li Dong and Shaohan Huang and Dongdong Zhang and Furu Wei } ,
    journal = { ArXiv } ,
    year    = { 2022 } ,
    volume  = { abs/2203.00555 }
}
Expandir
Información adicional
Aplicaciones relacionadas
Recomendado para ti
Información relacionada Todo