FlashAttention-3: Eşzamansızlık ve Düşük Hassasiyetle Daha Hızlı ve Daha Doğru Attention

 (together.ai)
1 puan yazan GN⁺ 2024-07-12 | 1 yorum | WhatsApp'ta paylaş
  • Transformer’larda Attention darboğazı, uzun bağlamlı LLM’lerin eğitim ve çıkarım hızını sınırlıyordu; FlashAttention-3 ise Hopper GPU özelliklerinden yararlanarak bunu daha da azaltmayı hedefleyen yeni sürüm.
  • Temel fikir, Tensor Core ve TMA’nin eşzamansızlığını kullanarak hesaplama ile veri taşımayı örtüştürmek ve matmul ile softmax işlemlerini dönüşümlü yerleştirerek GPU’nun boşta kalma süresini azaltmak.
  • H100’de FlashAttention-2, teorik azami FLOPS’un yalnızca %35’ini kullanabilirken FlashAttention-3, FP16’da en fazla 740 TFLOPS ile H100’ün teorik değerinin %75’ine ulaşıyor.
  • FP8 düşük hassasiyette, LLM activation’lardaki aykırı değerler hatayı büyüttüğü için Hadamard transform tabanlı incoherent processing, temel FP8 attention’a kıyasla hatayı 2,6 kat azaltıyor.
  • FlashAttention-3 GitHub’da yayımlandı ve Attention maliyeti düştükçe daha uzun bağlam işleme ile LLM eğitim/çıkarım verimliliğini artırma alanı büyüyor.

FlashAttention-3’ün hedefi ve performansı

  • Attention, Transformer’ın temel katmanı olsa da büyük dil modelleri ve uzun bağlamlı uygulamalarda başlıca darboğaz haline geliyor.
  • FlashAttention ve FlashAttention-2, GPU bellek okuma/yazmasını azaltarak Attention’ı hızlandıran yaklaşımın öncüsü oldu; bugün çoğu kütüphane bunu Transformer eğitimi ve çıkarımını hızlandırmak için kullanıyor.
  • Bu yaklaşım, LLM bağlam uzunluğunun son 2 yılda GPT-3·OPT’deki 2~4K seviyesinden GPT-4’te 128K’ya ve Llama 3’te 1M’ye çıkmasına katkı sağladı.
  • FlashAttention-2, H100 GPU’da teorik azami FLOPS’un yalnızca %35’ini kullanabilirken, FlashAttention-3 Hopper GPU’nun yeni özelliklerinden yararlanarak bunu yükseltiyor.
  • Açıklanan FlashAttention-3 performansı şöyle:
    • FP16’da FlashAttention-2’den 1,5~2,0 kat daha hızlı
    • FP16’da en fazla 740 TFLOPS
    • H100’ün teorik azami FLOPS’unun %75’i
    • FP8’de yaklaşık 1,2 PFLOPS
    • Temel FP8 attention’a kıyasla 2,6 kat daha düşük hata

FlashAttention yaklaşımına kısa bir bakış

  • FlashAttention, Attention hesaplama sırasını yeniden düzenleyip tiling ve yeniden hesaplama kullanarak hızı artırıyor; ayrıca sequence length’e göre bellek kullanımını quadratic’ten linear’e indiriyor.
  • Girdi blokları HBM’den SRAM’e yükleniyor, ilgili blok üzerinde Attention çalıştırılıyor ve ardından çıktı HBM’de güncelleniyor.
  • Büyük ara Attention matrisleri HBM’e yazılmadığı için bellek okuma/yazması azalıyor ve gerçek çalışma süresi bazında 2~4 kat hız artışı mümkün oluyor.
  • tiling ile softmax rescaling birlikte kullanıldığında, blok bazlı işleme rağmen yaklaşık hesap olmadan doğru çıktı elde edilebiliyor.

Hopper GPU özellikleri: WGMMA, TMA, FP8

  • FlashAttention-2, Ampere A100 GPU’da teorik azami FLOPS’un %70’ine kadar çıkabilse de Hopper GPU’nun yeni özelliklerini yeterince kullanamıyor.
  • FlashAttention-3, Hopper’ın üç özelliğini kullanıyor:
    • WGMMA: Hopper’ın yeni Tensor Core’larını kullanan warpgroup matrix multiply-accumulate özelliği; Ampere’deki mma.sync’ten daha yüksek throughput sağlıyor.
    • TMA: global memory ile shared memory arasındaki veri aktarımını hızlandıran özel donanım birimi; index hesaplama ve out-of-bound predication işlemlerini üstlenerek register kullanımını azaltıyor.
    • FP8: FP16’ya göre Tensor Core throughput’unu iki katına çıkarabiliyor, ancak kayan noktalı değerleri daha az bitle temsil ettiği için doğrulukla bir tradeoff getiriyor.
  • FlashAttention-3, Hopper özelliklerinden yararlanmak için NVIDIA CUTLASS soyutlamalarını kullanıyor.
  • FlashAttention’ın yalnızca bu yeni özellikleri kullanacak şekilde yeniden yazılması bile FP16 forward pass performansını FlashAttention-2’deki yaklaşık 350 TFLOPS’tan 540~570 TFLOPS seviyesine çıkarıyor.

Eşzamansızlıkla GEMM ve softmax’ı örtüştürmek

  • Attention’daki başlıca işlemler, Q-K ve P-V arasındaki GEMM ile softmax.
  • Modern hızlandırıcılarda matmul dışı işlemler, matmul’a göre çok daha yavaş; softmax içindeki exponential gibi özel fonksiyonlar da floating point multiply-add veya matrix multiply-add’den farklı birimde işleniyor.
  • H100 SXM5, FP16 matrix multiply için 989 TFLOPS sunarken özel fonksiyon throughput’u yalnızca 3,9 TFLOPS; yani 256 kat daha düşük.
  • Head dimension 128’de, matmul FLOPS’u exponential’dan 512 kat fazla olsa da exponential, matmul süresinin %50’sini kaplayabiliyor.
  • FP8’de matmul FLOPS’u iki kat hızlanırken exponential hızı aynı kaldığı için matmul ile softmax’ı paralel yürütmek daha da önemli hale geliyor.

  • warpgrouplar arasında pingpong scheduling

    • GPU warp scheduler, bazı warplar GEMM sonucunu beklerken diğer warpları çalıştırarak belli ölçüde overlap’ı otomatik sağlıyor.
    • FlashAttention-3 ise synchronization barrier kullanarak iki warpgroup’un GEMM ve softmax işlemlerini elle daha iyi örtüştürüyor.
    • warpgroup 1, bir iterasyondaki GEMM1 ile sonraki iterasyondaki GEMM0’ı önce yürütüyor.
    • Sonrasında warpgroup 2 GEMM çalıştırırken warpgroup 1 softmax işliyor.
    • Bu pingpong schedule, softmax süresini diğer warpgroup’un GEMM çalışma süresinin arkasına gizliyor.
    • Gerçek zamanlama diyagramdaki kadar kusursuz olmasa da, FP16 attention forward pass’te head dimension 128 ve sequence length 8K için performansı yaklaşık 570 TFLOPS’tan 620 TFLOPS’a çıkarıyor.

  • warpgroup içinde overlap

    • Tek bir warpgroup içinde de o warpgroup’un GEMM’i çalışırken softmax’ın bir kısmı yürütülebiliyor.
    • Bu pipelining, FP16 attention forward throughput’unu yaklaşık 620 TFLOPS’tan 640~660 TFLOPS’a yükseltiyor.
    • Ancak GEMM accumulator’ları ile softmax giriş/çıkışlarını birlikte tutmak gerektiğinden register pressure artıyor.
    • Genel olarak bu teknik avantajlı bir tradeoff sunuyor.

FP8 düşük hassasiyet ve incoherent processing

  • LLM activation’larında, diğer feature’lara göre magnitude’u çok daha büyük olan outlier değerler bulunabiliyor.
  • Bu outlier’lar quantization’ı zorlaştırıyor ve quantization error’u ciddi biçimde büyütüyor.
  • FlashAttention-3, QuIP gibi quantization literatüründe kullanılan incoherent processing yaklaşımından yararlanıyor.
  • Query ve key üzerine rastgele orthogonal matrix çarpılarak outlier’lar dağıtılıyor ve quantization error azaltılıyor.
  • Uygulamada random sign içeren Hadamard transform kullanılıyor.
    • Head dimension d olduğunda, attention head başına işlem O(d²) değil O(d log d) zamanda yapılabiliyor.
    • Hadamard transform, memory-bandwidth bound olduğu için rotary embedding gibi yine memory-bandwidth bound olan önceki işlemlerle fuse edilirse ek maliyet olmadan uygulanabiliyor.
  • Q, K, V’nin standart normal dağılımdan üretildiği ve girdilerin %0,1’ine büyük magnitude verilerek outlier simülasyonu yapılan deneyde, incoherent processing quantization error’u 2,6 kat azaltıyor.

Benchmark’lar ve yayın durumu

  • FlashAttention-3, yalnızca FlashAttention-2 ile değil, Hopper GPU’nun yeni donanım özelliklerini zaten kullanan Triton ve cuDNN uygulamalarıyla da karşılaştırıldı.
  • FP16’da FlashAttention-2’ye göre yaklaşık 1,6~1,8 kat hız artışı gösteriyor.
  • FP8’de yaklaşık 1,2 PFLOPS seviyesine ulaşıyor.
  • FlashAttention-3 GitHub deposu yayımlandı.
  • Makaleye de aynı flash-attention deposundan ulaşılabiliyor.

Kalan optimizasyonlar ve gelecekteki entegrasyon

  • Makale, blog yazısında ele alınanların yanı sıra variable length sequence, persistent kernel ve FP8 için in-kernel transpose gibi optimizasyonları da içeriyor.
  • Algoritma, çalıştığı donanıma göre tasarlandığında büyük verimlilik artışları ve uzun bağlam gibi yeni model yeteneklerinin önü açılabiliyor.
  • Gelecekteki çalışmalar arasında LLM inference optimizasyonu ve tekniğin diğer donanım mimarilerine genellenmesi bulunuyor.
  • FlashAttention-3’ün gelecekteki PyTorch sürümlerine entegre edilmesi bekleniyor.

İlgili okumalar

1 yorum

 
GN⁺ 2024-07-12
Hacker News yorumları

  • Kod yorumlarına bakılırsa Tri Dao, FA3 üzerinde Hopper/H100 duyurusunun hemen ardından, Nisan 2022’den beri çalışıyormuş gibi görünüyor.
    Kodun bugün yayımlanmasına kadar iki yıldan fazla zaman geçmiş olması biraz ilginç; belki de daha iyi çözümler hazırlık aşamasında olduğu içindir.
    Tri’nin son dönem makale geçmişi SSM ve Mamba ailesi mimarilere doğru kayıyor. FlashAttention, dizi uzunluğuna göre ikinci dereceden zaman karmaşıklığına sahip; ancak güncel algoritmalar alt-ikinci dereceden olduğu için aynı hesaplamayı daha verimli yapmaktan öte, hesap miktarının kendisini çok daha fazla azaltıyor.
    Dao ve Gu, bu yılki uzun makalelerinde Mamba/SSM’nin de Transformer’ın faydalandığı aynı donanım ilkel işlemleriyle hızlandırılmaya elverişli olacak şekilde formüle edilebileceğini gösterdi.

    • Güçlü Üstel Zaman Hipotezi (SETH) kanıtlanana veya çürütülene kadar, ya ikinci dereceden maliyet gerekecek ya da bir şeylerden vazgeçmek gerekecek. Sonuçta bu, tam aramanın maliyeti.
      SETH’yi kanıtlamak veya çürütmek P-NP problemini de çözeceğinden, bunun yakında olmasını beklemek zor.
      Asıl mesele, belirli bir kullanım senaryosunun bu maliyeti kaldırıp kaldıramadığı.

  • FlashAttention algoritmasının donanıma ne kadar bağlı olduğunu merak ediyorum.
    Örneğin bu duyuruda H100 GPU’nun asenkron özelliklerinden yararlandığı söyleniyor; bu da H serisi olmayan kartlarda bu hız artışının elde edilemeyeceği anlamına geliyor gibi görünüyor.
    Ayrıca gerçek FlashAttention kütüphanesi CUDA gerektiriyor, ancak algoritma Metal’e port edilmiş gibi görünüyor[^0]. Algoritma saf fonksiyona yakın bir şeyse, herhangi bir GPU/makine öğrenimi framework’ünde uygulanabilir olması gerekmez mi diye düşünüyorum.
    [0]: https://github.com/philipturner/metal-flash-attention

    • İyi yanıt çok, ama kısaca söylemek gerekirse “pratikte epey fazla” donanıma bağlı. Aşağıdaki iyi bir örnek:

      https://github.com/karpathy/nanoGPT/blob/master/model.py#L45
      Karpathy’nin nanoGPT’si, torch.nn.functional.scaled_dot_product_attention var mı diye kontrol edip FlashAttention’ı çağırıyor.
      https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
      Dokümanlara bakınca, gerçekte çoğunlukla FA2’nin çağrılmasını isteyeceğiniz görülüyor; FA2, cihazın kernel’ını optimize ederek üçgensel matrisin Softmax işlemini parçalara ayırıyor ve gereksiz kayan nokta gruplarını GPU ile CPU arasında gidip getirmeyi azaltıyor.
      https://arxiv.org/pdf/2307.08691
      FA2 makalesi neredeyse tamamen üzerinde çalıştığı donanım perspektifinden anlatılıyor.

    • FlashAttention’daki algoritmik iyileştirme esas olarak attention’ın Softmax kısmını bölüp birleştirmekten ibaret; bu da başlı başına tamamen yeni bir fikir değil. Asıl ezici katkı, bu yöntemi ve ayrıntılarını Nvidia donanımında verimli biçimde uygulamasında yatıyor.
    • Orijinal FlashAttention’ın donanım bağımlılığı neredeyse yoktu.
      En yeni sürüm, soyutlama düzeyine bağlı. ThunderKittens[0], yazıda anlatılana benzer şekilde FA2’ye kıyasla 1,3–2 kat hız artışı sağlarken GPU’lar genelinde de görece yaygın uygulanabilirlik sunuyor.
      Her yeni donanımda ek performans çıkarabilecek donanıma özgü özellikler bulunabilir. Genelde tedarikçiler kendilerini öne çıkaran özellikleri benimser, ancak CUDA’da da zaten olduğu gibi API’ler ve kütüphaneler parçalanır.
      [0]: https://hazyresearch.stanford.edu/blog/2024-05-12-tk
    • Kavramsal olarak biraz, pratik uygulama açısından ise çok bağlı. Standart Python uygulaması bile içeride belirli donanıma göre kernel derliyor.
    • Pratik açıdan eklemek gerekirse, AMD donanımında hâlâ flash-attention-2’nin düzgün bir uygulaması eksik. ROCm yavaş yavaş kullanılabilir hâle geliyor ama CUDA ile karşılaştırılacak seviyede değil.

  • Derleyici tarafındaki kişilere sormak istiyorum. Bir derleyicinin FlashAttention gibi optimizasyonları kendi kendine bulma ihtimali var mı? TVM ve tinygrad bu yöne gidiyor gibi görünüyor ama bunun mümkün olduğuna inanmak zor

    • Teorik olarak mümkün. Matematiğin cebirsel özellikleri sayesinde büyük çaplı yeniden sıralama yapılabilir; buna nispeten düzenli çokyüzlü döngü döşemesi eklemek yeterli olur
      Ancak maliyeti yüksek olduğu için bu aramanın sonuçlarını önbelleğe almak gerekir
      e-graph optimizasyonu bu alana iyi uyuyor gibi görünüyor. Ancak optimizasyon pass’lerinin işlenme biçiminde büyük bir paradigma değişimi gerektiğinden, bazı niş araçlar dışında pek dağıtıma girmedi. Örneğin geleneksel çağrı grafiğiyle iyi uyuşmuyor; temel blokların ve for döngülerinin dışına/arasına taşan şekilde e-graph dağıtmak için kontrol akışını ciddi biçimde değiştirmek gerekiyor, ayrıca break ve return de desteklenmiyor
    • Aşırı zor ama imkânsız olmayan bir problem gibi görünüyor
      Derleyici optimizasyonunun güncel seviyesinin veri yerleşimi ve işlemci kullanımını en üst düzeye çıkarma açısından nereye geldiğini pek bilmiyorum
      Eskiden optimizasyonla ilgili bir videoda, küçük bir optimizasyonun hızı artırdığı ama bu optimizasyonun, hatta rastgele değişikliklerin yarattığı bellek yerleşimi farklarından kaynaklanan hız dalgalanmalarına kıyasla etkisinin önemsiz kaldığı anlatılıyordu
      O sunum daha çok gürültü içindeki sinyali ayırt etmeye odaklanıyordu; fakat gürültünün kendisi bile derleyicinin burada sözü edilen problemden çok daha basit biçimleri bile iyi ele alamadığının bir izi
      Yalnızca önbelleklerin ve erişim desenlerinin hızı etkilediği CPU-bellek mimarisi bile karmaşıkken, buna GPU mimarisini de ekleyince oldukça bakir bir alan gibi duruyor
      Bir gün mümkün olabilir. Konu AI alanı olduğuna göre, yeterince akıllı bir AI’ın bunu yapıp yapamayacağı sorusu da doğuyor; ama mesele “yeterince”in ölçütünün ne olduğuna bağlı
      AI modelleri için çok üst seviye bir test olarak, micrograd gibi bir şey verip aynı arayüzü koruyarak torch’tan daha hızlı bir şey yapmasını istemek düşünülebilir. Henüz bunun yakınında bile değiliz ama mümkün olursa ilginç olur
    • Bence hayır. Bunu farklı bir algoritma gibi düşünmek gerekir. Yalnızca matematiği dikkate almak yerine, algoritmayı donanımın biçimini hesaba katarak tasarlıyorsunuz
      TVM anlaşılır. Kesin konuşursak farklı bir iş yapıyor ama oldukça yakın bir alan
      Ama tinygrad için neden öyle hissettirdiğini bilmiyorum
    • https://github.com/uwplse/tensat
    • Python gibi sarmalayıcı bir dilden yüksek seviyeli operatörleri çağırmak oldukça zahmetli

  • Bunu ROCm / AMD MI300x’e port etmek isteyen varsa hello@hotaisle.xyz adresinden iletişime geçsin. Asla spam göndermem
    Bu iş için hesaplama zamanı bağışlayabilirim

    • AMD hızlandırıcı sunucuları şirketisiniz demek! Harika bir iş; umarım biri üstlenir :)
    • Kaba olmak istemem ama bu teklifin niyetini merak ediyorum. Yalnızca donanıma erişim hakkı verildi diye biri bu portu ücretsiz yapar mı? O kişiye ne faydası var?

  • FlashAttention-3 is optimized for Hopper GPUs (e.g. H100).
    FA3, 3090 ve 4090 gibi tüketici GPU’larında ne düzeyde performans veriyor?

  • 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.
    Anladığım kadarıyla TMA’nın register kullanımını azaltması doğru, ama daha önemlisi adres üretimini donanıma yaptırması. Etraftaki işlemler hızlandıkça adres üretimi darboğaz hâline gelebilir

  • Bu, tüm AI alanındaki en önemli iyileştirmelerden biri. Aynı donanımla daha fazlasını daha hızlı kullanmayı sağlıyor ve çoğu AI kullanıcısına neredeyse tavizsiz kazanç sunuyor

    • H100 sahibi kullanıcılar için öyledir tabii

  • FlashAttention’ın değişken maskeleme kullandığında kullanmadığı duruma göre neden yaklaşık 5 kat yavaş olduğunu merak ediyorum. İyi maskeleme desteği yoksa optimizasyon etkisi neredeyse ortadan kalkıyor

    • O benchmark’ı nerede görüyorsunuz?

  • Bir uzmanın birkaç soruyu yanıtlaması güzel olurdu :)
    FlashAttention, LLM'in attention işleminin yerine doğrudan konabilen bir drop-in mi? “Attention” işleminin kullanıldığı her yerde kullanılabilir mi, yoksa FA kullanmak için LLM'i ayrıca eğitmek mi gerekir?
    FA'nın GQA, sliding window attention gibi stratejilerle ilişkisi nedir? Birbirinden bağımsız kavramlar mı, yoksa her strateji için ayrı bir FA implementasyonu mu gerekir?
    Kısa süre önce llama.cpp FlashAttention desteği ekledi; bu, FlashAttention'ın sağladığı CUDA kernel'ları gibi şeyleri kullanmaya başladığı anlamına mı geliyor?
    Son olarak, bu yazıda FlashAttention ile Triton karşılaştırılıyor. Triton bir soyutlama katmanı gibi bir şey değil mi? FA Triton ile implemente edilemez mi? “FlashAttention vs Triton” ifadesini pek anlayamadım

      1. Neredeyse doğru. Matematiksel olarak eşdeğer. Yazılım tarafındaki sorunlar yalnızca bağımlılık sürüm yönetimi ya da bellekteki veri biçimi gibi şeyler; FlashAttention 2 zaten HuggingFace'e ve birçok popüler kütüphaneye girmiş durumda. FlashAttention 3'ün de yakında girmesi muhtemel, ancak çalıştırmak için H100 GPU gerekiyor
      2. FlashAttention 2, önceki sürüm güncellemelerinde GQA desteği eklemişti:
        https://github.com/Dao-AILab/flash-attention
      3. Burada saf CUDA C++ ile yazılmış bu FlashAttention implementasyonu ile Triton'da yazılmış benzer bir algoritmanın Triton implementasyonu karşılaştırılıyor: https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html
    • FlashAttention, LLM'in attention işleminin yerine drop-in olarak geçebilir
      FlashAttention, attention'ın Softmax(QK^T)V bölümünü hesaplama yöntemidir; GQA ise Q, K, V matrislerini hesaplama yöntemidir. Sliding window attention konusunda o kadar emin değilim, ama hangi query'nin hangi key'leri görebileceğini kontrol eden attention maskesini değiştirme yöntemi
      llama.cpp'yi kullanmadım, ancak CUDA kernel'larını alıp kullanmaya başladığı açıklaması genel olarak doğru görünüyor
      Son soru, daha önce Triton ile yazılmış FlashAttention implementasyonundan bahsediyor

  • Bu yazıda sigmoid gibi işlemlerin çok yavaş olduğu söylendiği için merak ettim
    Modern LLM'ler SiLU, Swish, SOLU gibi sigmoid veya Softmax içeren aktivasyon fonksiyonlarını çok kullanıyor
    ReLU'da performans kaybı daha mı az? Öyleyse eski iyi ReLU'ya dönmek daha iyi olabilir mi?

    • ReLU kelimenin tam anlamıyla bir noktada 0'a kırpılan doğrusal bir fonksiyon olduğu için, üstel fonksiyon içerenlere göre hesaplama yükü çok daha az. Ancak bu kadar basit bir aktivasyon fonksiyonuyla rekabetçi sonuçlar elde etmek zor olur gibi görünüyor