flash cosine sim attention

_{Selami Pembelajaran Mendalam, yang dibuat ulang oleh Majalah Quanta}

Perhatian Kesamaan Flash Cosinus

Implementasi perhatian kesamaan kosinus yang menyatu dalam gaya yang sama dengan Flash Attention. Pengamatannya adalah dengan mengadopsi l2 kueri dan kunci yang dinormalisasi, Anda tidak perlu lagi melacak baris maksimum untuk stabilitas numerik. Hal ini sangat menyederhanakan algoritme perhatian kilat, dengan asumsi perhatian kesamaan kosinus tidak memerlukan biaya generalisasi.

Dengan kata lain, perhatian konteks yang stabil, cepat, hemat memori, dan lebih lama tanpa kerugian.

Pembaruan: Sayangnya, percobaan Robin menunjukkan evaluasi skor FID yang jauh lebih buruk tidak tercermin dalam kerugian. Menunggu eksperimen lebih lanjut. Gunakan perpustakaan ini dengan hati-hati.

Pembaruan 2: Satu-satunya anugrah adalah dengan menggunakan l2norm yang dikelompokkan, yang berpotensi memungkinkan lebih banyak ekspresi. Jika ada yang bisa mengevaluasi teknik ini pada pekerjaan generatif mereka dan memperoleh beberapa skor FID, akan sangat dihargai.

Pembaruan 3: Pendekatan yang mirip dengan perhatian sim kosinus telah terbukti dalam skala besar, dengan model visi parameter 22B dari Brain.

Status (hapus)

Saat ini, urutan autoregresif dan panjang variabel harus lebih cepat di semua arsitektur. Untuk urutan yang lebih lama dari 2048, ini juga akan hemat memori jika tidak ada perhatian rutin.

Namun, untuk non-autoregresif tanpa masking, arsitekturnya masih lebih lambat di A100 untuk F16. Tujuannya adalah agar kinerjanya lebih cepat pada A100 maju dan mundur untuk F32 dan F16, karena memori bersama belum sepenuhnya dieksploitasi.

Kartu grafis lama tanpa memori bersama yang cukup, kita harus mengukur tradeoff efisiensi dan kecepatan memori tergantung pada panjang urutan yang dilatih.

Apresiasi

• Arthur Hennequin yang telah melatih saya melalui kernel CUDA pertama saya, dan untuk membuat kode implementasi referensi sederhana, yang membantu saya melakukan bootstrap pada kernel pertama yang memiliki kinerja wajar hingga dasar. Pekerjaan ini tidak akan mungkin terwujud tanpa keahliannya.

• Boris Dayma dan Robin Rombach yang menjalankan eksperimen perhatian sim kosinus yang disederhanakan dengan penskalaan tetap pada beberapa model teks-ke-gambar yang signifikan dan memverifikasi bahwa kinerjanya memang sama baiknya dengan perhatian biasa.

• Markus Rabe yang menulis makalah yang menunjukkan bahwa perhatian tidak memerlukan memori O(n²), dan Tri Dao yang menyatukan semuanya dalam implementasi kernel CUDA untuk perhatian rutin, menunjukkan keunggulan dalam kecepatan menggunakan pendekatan ubin yang meminimalkan akses HBM (dan untuk menghitung keluar dO * O == dP * P untuk operan mundur). Tidak akan mampu menyelesaikan ziarah saya mencari formulasi perhatian tertinggi tanpa penemuan mereka.

• Stability.ai atas sponsor yang murah hati untuk mengerjakan penelitian kecerdasan buatan yang mutakhir

Memasang

$ pip install flash-cosine-sim-attention

Penggunaan

Perhatian Diri

 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)

Perhatian silang

 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)

Dengan penyembunyian kunci/nilai

 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)

Autoregresif

 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) 

Aneka ragam

Kunci/nilai berkepala tunggal (Shazeer dkk & digunakan di 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)

Jika Anda perlu melakukan operasi pada kueri dan kunci di antara l2norm dan langkah perhatian sebenarnya, cukup setel l2norm_qk = False

mantan.

 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)

Perhatian silang dengan pekerjaan kausal seperti yang diharapkan - (cache kunci dan nilai dalam autoregresif selama inferensi, atau pelatihan seperti transformator-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)

Jika Anda menggabungkan dimensi batch dan head, tidak apa-apa

 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) 

Dimensi kepala yang didukung

• 16 - f32

• 32

• 64

• 96

• 128

• 16 -f16

• 80 - sedang berlangsung

Semua yang harus dilakukan

• dukungan bfloat16, gunakan sfinae seperti yang direkomendasikan oleh Arthur

• streaming dari qk_mma ke memori bersama dalam beberapa bagian untuk menghitung mma, lihat apakah smem yang dibebaskan dapat digunakan untuk menyimpan cache lebih banyak

• mendukung O(n) bias posisi dinamis 1d

• mencari tahu mengapa cache fragmen kecil akan menyebabkan penurunan kinerja, itu tidak masuk akal

• pikirkan tentang penggunaan logsumexp - berfungsi tetapi log tambahan menyebabkan penurunan kinerja

• menyiapkan mekanisme caching fragmen kecil, untuk memungkinkan caching sebanyak yang diizinkan pada A100 (atau f16)

• membuat pemrosesan ukuran ubin perhatian dapat disesuaikan untuk passing mundur

• pindahkan atom tambahkan ke fungsi kelebihan beban di dalam mma

• fleksibel jenis mana yang digunakan untuk akumulasi

• uji ubin 64x96 di f16

• menghadirkan versi hemat memori CPU (hanya untuk inferensi, karena pelatihan tidak masuk akal) hanya dengan menggunakan kode pytorch biasa

• cari tahu cara mengirimkan secara berbeda untuk arsitektur (katakanlah A100), seandainya mundur dapat memanfaatkan peningkatan memori bersama secara berbeda

• memisahkan ukuran baris dan kolom untuk ubin perhatian

• dk dan dv sekarang berada di f16 jika memungkinkan (bukan kv berkepala tunggal)

• mendukung dimensi kepala yang lebih standar (wip)

• debug dan perbaiki bias mundur gradien lagi untuk ukuran kepala 32

• memperbaiki gradien bias perhatian

• mengizinkan kunci/nilai berkepala tunggal, seperti di PaLM

• perbaiki penambahan atom untuk f16

• bias perhatian harus dapat menerima dimensi dimensi batch tambahan, karena Alphafold2 seperti bias perhatian

• mengotomatiskan penghilangan cache kernel menggunakan versi sebagai akhiran pada nama paket

• menyelesaikan masalah numerik kausal f16

• adopsi semua pembelajaran dari kernel maju ke kernel mundur dan pastikan kinerjanya mengungguli setidaknya pada A100

Keterangan

Sejauh ini perhatian cosine-similarity belum banyak digunakan dalam industri. Satu-satunya model besar yang telah dilatih sejauh ini adalah SwinV2. Jika ada yang bisa membatalkan pendekatan ini, silakan buka terbitan atau kirimi saya email. Anda dapat menjalankan eksperimen dengan perhatian rutin menggunakan repositori x-transformers.

Pembaruan: Boris Dayma dengan baik hati memulai eksperimen (biru dengan merah sebagai garis dasar) untuk memvalidasi perhatian kesamaan kosinus dengan skala tetap 10 dalam pengaturan model dunia nyata.

Pembaruan 2: Perhatian kesamaan kosinus telah dibuktikan dalam jaringan perhatian teks-ke-gambar dunia nyata, menggunakan skala konstan 10 . Tidak lebih buruk dari perhatian biasa. Penghargaan diberikan kepada Boris Dayma karena telah menginvestasikan waktu untuk menjalankan eksperimen dan menghilangkan keraguan seputar teknik ini.

Pembaruan 3: Robin Rombach telah menguji kernel di repositori ini dengan ukuran kepala 64 dan skala tetap 10 dalam model teks-ke-gambar, dan tidak menemukan perbedaan dari perhatian biasa. Evaluasi lebih lanjut menunggu keputusan.

Pembaruan 4: Peningkatan kinerja yang terlihat dalam eksperimen Boris kemungkinan besar disebabkan oleh fakta bahwa perhatian cosinus-sim memungkinkan seseorang untuk beralih dari konfigurasi pra-layernorm ke pasca-layernorm di transformator (karena l2norm secara efektif menggantikan pra- norma lapisan). Perhatian cosinus sim kemungkinan besar akan memberikan hasil yang sama dengan perhatian biasa, tanpa adanya perubahan lain pada trafo.

Pengujian

Untuk pengujian, keluaran dan gradiennya sama untuk skenario non-autoregresif dan autoregresif

$ python setup.py test 

Pembandingan

Pastikan untuk menginstal kernel CUDA terlebih dahulu

$ python setup . py install

Kemudian

$ python benchmark . py

Untuk pembandingan maju atau mundur saja, tambahkan tanda --only-forwards atau --only-backwards ke tanda di atas. Untuk mengukur autoregresif, tambahkan --causal

Tolok ukur - hapus

GTX 2080 Ti

Maju

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

Mundur - masih perlu perbaikan

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

Maju & Mundur - F32 jelas lebih lambat

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

Untuk autoregresif, win 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

Untuk urutan panjang variabel dengan masking, juga merupakan kemenangan yang jelas. Asumsikan rata-rata 25% token ditutupi 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 (hapus)

Terima kasih kepada Stabilitas karena telah menyediakan akses ke A100 untuk pengujian. Terima kasih kepada Enrico karena telah meluangkan waktu untuk menjalankan beberapa benchmark ketika saya belum memiliki akses.

A100 masih dalam proses. Memori bersama belum sepenuhnya dieksploitasi. Anehnya, F32 nampaknya lebih baik dari F16

Ke depan

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

Ke belakang

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

Maju & Mundur

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

Autoregresif

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

Urutan panjang variabel (hingga 25% token terselubung)

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

Melatih GPT kecil di Enwik8

$ make train

Coba panjang urutan 8192. Ini akan lambat tetapi akan berhasil (perhatian normal akan berhenti pada > 2048, Anda akan melihat ini jika Anda menghapus tanda --use-cuda-kernel )

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

Kutipan

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