flash cosine sim attention

_{Mergulhe no Deep Learning, refeito pela Quanta Magazine}

Atenção à similaridade do cosseno flash

Implementação de atenção de similaridade de cosseno fundido no mesmo estilo que Flash Attention. A observação é que, ao adotar consultas e chaves normalizadas l2, você não precisa mais acompanhar os máximos das linhas para estabilidade numérica. Isso simplifica muito o algoritmo de atenção flash, assumindo que a atenção à similaridade de cosseno não tem custo de generalização.

Em outras palavras, atenção de contexto estável, rápida, eficiente em termos de memória e mais longa, sem desvantagens.

Atualização: Infelizmente, os experimentos de Robin mostraram pontuações de avaliação FID muito piores, não refletidas na perda. Pendente de mais experimentos. Use esta biblioteca com cuidado.

Atualização 2: A única graça salvadora seria usar l2norm agrupado, o que poderia permitir mais expressividade. Se alguém puder avaliar esta técnica em seu trabalho generativo e obter algumas pontuações no FID, ficaria muito grato.

Atualização 3: Uma abordagem semelhante à atenção do cosseno sim foi comprovada em escala, com um modelo de visão de parâmetros de 22B do Brain.

Status (wip)

No momento, sequências autorregressivas e de comprimento variável devem ser mais rápidas em todas as arquiteturas. Para sequências superiores a 2048, também será eficiente em termos de memória onde a atenção regular não o faria.

Porém, para não autorregressivo sem mascaramento, a arquitetura ainda é mais lenta no A100 para F16. O objetivo é fazer com que ele funcione mais rápido no A100 para frente e para trás tanto para F32 quanto para F16, já que a memória compartilhada ainda não está totalmente explorada.

Placas gráficas mais antigas sem memória compartilhada suficiente, será necessário avaliar a compensação entre eficiência e velocidade da memória, dependendo do comprimento da sequência que está sendo treinada.

Apreciação

• Arthur Hennequin por me orientar em meu primeiro kernel CUDA e por codificar uma implementação de referência simples, o que me ajudou a inicializar o primeiro kernel que apresenta desempenho razoável até a linha de base. Este trabalho não teria sido possível sem a sua expertise.

• Boris Dayma e Robin Rombach por realizarem experimentos com a atenção simplificada do cosseno sim com escala fixa em alguns modelos significativos de texto para imagem e verificarem se ele realmente funciona tão bem quanto a atenção regular.

• Markus Rabe por escrever o artigo que mostrou que atenção não requer memória O(n²), e Tri Dao por juntar tudo em uma implementação de kernel CUDA para atenção regular, demonstrando superioridade em velocidade usando a abordagem lado a lado minimizando acessos HBM (e por descobrir out dO * O == dP * P para passe para trás). Não teria sido capaz de completar minha peregrinação em busca da formulação definitiva da atenção sem suas descobertas.

• Stability.ai pelo generoso patrocínio para trabalhar em pesquisas de ponta em inteligência artificial

Instalar

$ pip install flash-cosine-sim-attention

Uso

Autoatenção

 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)

Atenção 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)

Com máscara de chave/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)

Autoregressivo

 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) 

Variado

Chaves/valores de uma só cabeça (Shazeer et al e usados ​​no 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)

Se você precisar fazer operações nas consultas e chaves entre l2norm e a etapa de atenção real, basta definir 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)

A atenção cruzada com o causal funciona conforme o esperado - (armazenamento em cache de chaves e valores em autorregressivo durante a inferência ou treinamento semelhante ao transformer-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)

Se você tiver dimensões de lote e cabeçote mescladas, tudo bem

 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) 

Dimensões da cabeça suportadas

• 16-f32

• 32

• 64

• 96

• 128

• 16-f16

• 80 - em andamento

Pendência

• suporte bfloat16, use sfinae conforme recomendado por Arthur

• transmitir de qk_mma para a memória compartilhada em pedaços para calcular mma, ver se o smem liberado pode ser usado para armazenar mais cache

• suporta O (n) 1d viés posicional dinâmico

• descobrir por que o cache de fragmentos smem levaria à degradação do desempenho, não faz sentido

• pense no uso de logsumexp - funciona, mas log extra leva a desempenho degradado

• preparar um mecanismo de cache de fragmentos smem, para permitir tanto cache quanto permitido no A100 (ou f16)

• tornar o processamento do tamanho do bloco de atenção personalizável para passagem para trás

• mover atômico adicionar à função sobrecarregada dentro do mma

• flexível qual tipo é usado para acumulação

• teste blocos 64x96 em f16

• trazer uma versão eficiente de memória da CPU (apenas para inferência, pois o treinamento não faz sentido) usando apenas código pytorch simples

• descobrir como despachar de maneira diferente para arquiteturas (digamos A100), caso o retrocesso possa fazer uso do aumento na memória compartilhada de maneira diferente

• dissociar tamanhos de linha e coluna para blocos de atenção

• dk e dv estão agora em f16 quando podem estar (kv sem cabeça única)

• suporta dimensões de cabeçote mais padrão (wip)

• depurar e corrigir gradientes de polarização para trás mais uma vez para tamanho de cabeça de 32

• corrigir gradientes de viés de atenção

• permitir chaves/valores de uma só cabeça, como em PaLM

• corrigir adição atômica para f16

• o viés de atenção deve ser capaz de aceitar dimensões de uma dimensão de lote extra, para Alphafold2 como o viés de atenção

• automatizar o bloqueio de cache do kernel usando a versão como sufixo para o nome do pacote

• resolver problemas numéricos causais f16

• adote todos os aprendizados do kernel avançado para o kernel anterior e certifique-se de que ele tenha desempenho superior pelo menos no A100

Descrição

Até agora, a atenção à similaridade de cossenos não é amplamente utilizada na indústria. O único modelo grande que foi treinado até agora é o SwinV2. Se alguém puder invalidar a abordagem, abra um problema ou envie-me um e-mail. Você pode realizar experimentos com atenção regular usando o repositório x-transformers.

Atualização: Boris Dayma gentilmente iniciou um experimento (azul com vermelho como linha de base) para validar a atenção à similaridade de cosseno com uma escala fixa de 10 em um modelo do mundo real.

Atualização 2: A atenção de similaridade de cosseno foi comprovada em uma rede de atenção de texto para imagem do mundo real, usando uma escala constante de 10 . Não é pior do que atenção regular. O crédito vai para Boris Dayma por investir tempo na execução do experimento e tirar dúvidas em torno da técnica.

Atualização 3: Robin Rombach testou o kernel neste repositório com tamanho de cabeçalho de 64 e escala fixa de 10 em um modelo de texto para imagem, não observando nenhuma diferença da atenção regular. Mais avaliações pendentes.

Atualização 4: A melhoria no desempenho observada nos experimentos de Boris provavelmente se deve ao fato de que a atenção do cosseno-sim permite que se mude da configuração pré-camada para a configuração pós-camada nos transformadores (já que o l2norm efetivamente toma o lugar do pré- norma de camada). A atenção do cosseno sim provavelmente produzirá resultados iguais aos da atenção regular, sem quaisquer outras alterações no transformador.

Teste

Para testar, a saída e os gradientes são iguais para cenários não autorregressivos e autorregressivos

$ python setup.py test 

Comparativo de mercado

Certifique-se de instalar primeiro o kernel CUDA

$ python setup . py install

Então

$ python benchmark . py

Para fazer benchmarking apenas para frente ou para trás, anexe o sinalizador --only-forwards ou --only-backwards ao acima. Para avaliar o autorregressivo, acrescente --causal

Referências - wip

GTX 2080Ti

Avançar

------------------------------------------------------------
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

Para trás - ainda precisa de trabalho

------------------------------------------------------------
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

Avançar e retroceder - F32 é definitivamente mais 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 autorregressivo, uma vitória 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 sequências de comprimento variável com mascaramento, também é uma vitória clara. Suponha que, em média, 25% dos tokens sejam mascarados 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 (wi-fi)

Agradecemos à Stability por fornecer acesso aos A100s para testes. Obrigado ao Enrico por dedicar seu tempo para executar alguns benchmarks quando eu ainda não tinha acesso.

A100 ainda é um trabalho em andamento. A memória compartilhada ainda não está totalmente explorada. Estranhamente, o F32 parece estar melhor que o F16

Avançados

------------------------------------------------------------
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

Para trá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

Para frente e para trá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

Autoregressivo

------------------------------------------------------------
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

Sequências de comprimento variável (até 25% de tokens mascarados)

------------------------------------------------------------
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

Treinando um pequeno GPT no Enwik8

$ make train

Experimente o comprimento da sequência 8192. Será lento, mas funcionará (a atenção normal será interrompida em> 2048, você verá isso se remover o sinalizador --use-cuda-kernel )

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

Citações

 @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 }
}