Flash Attention: Rewolucjonizacja wydajności modeli Transformer

</pre>
import torch

<p>def standardowa_uwaga(Q, K, V):
# Q, K, V kształt: (rozmiar_partii, długość_sekwencji, d_model)
d_k = K.size(-1)
wyniki = torch.matmul(Q, K.transpose(-2, -1)) / torch.sqrt(torch.tensor(d_k))
wagi_uwagi = torch.softmax(wyniki, dim=-1)
return torch.matmul(wagi_uwagi, V)</p>

Ta implementacja, chociaż prosta, cierpi na nieefektywności wymienione powyżej. Tensor wyniki, który ma kształt (rozmiar_partii, długość_sekwencji, długość_sekwencji), może stać się niezwykle duży dla długich sekwencji.

Wejdź Flash Attention

Flash Attention, wprowadzony przez Tri Dao i współpracowników w ich pracy z 2022 roku, jest podejściem do obliczania uwagi, które dramatycznie redukuje użycie pamięci i poprawia wydajność obliczeniową. Kluczowe pomysły za Flash Attention to:

1. Tiling: Rozbicie dużej macierzy uwagi na mniejsze kafelki, które mieszczą się w szybkiej pamięci SRAM.
2. Ponowne obliczanie: Zamiast przechowywać całą macierz uwagi, ponowne obliczanie części jej podczas etapu wstecznego.
3. Implementacja zorientowana na wejście/wyjście: Optymalizacja algorytmu w celu minimalizacji ruchu danych między różnymi poziomami hierarchii pamięci GPU.

Algorytm Flash Attention

W swojej istocie Flash Attention ponownie wyobraża, jak obliczamy mechanizm uwagi. Zamiast obliczania całej macierzy uwagi na raz, przetwarza ją w blokach, wykorzystując hierarchię pamięci nowoczesnych GPU.

Oto ogólny przegląd algorytmu:

1. Wejście: Macierze Q, K, V w pamięci HBM i w pamięci SRAM o rozmiarze M.
2. Rozmiary bloków są obliczane na podstawie dostępnej pamięci SRAM.
3. Inicjacja macierzy wyjściowej O i wektorów pomocniczych l i m.
4. Algorytm dzieli macierze wejściowe na bloki, które mieszczą się w pamięci SRAM.
5. Dwa zagnieżdżone pętle przetwarzają te bloki:
   • Zewnętrzna pętla ładuje bloki K i V
   • Wewnętrzna pętla ładuje bloki Q i wykonuje obliczenia
6. Obliczenia w pamięci SRAM obejmują mnożenie macierzy, softmax i obliczanie wyjścia.
7. Wyniki są zapisywane z powrotem do pamięci HBM po przetworzeniu każdego bloku.

To przetwarzanie blokowe pozwala Flash Attention na utrzymanie znacznie mniejszego śladu pamięci, jednocześnie obliczając dokładną uwagę.

Matematyka za Flash Attention

Kluczem do tego, aby Flash Attention działał, jest sztuczka matematyczna, która pozwala obliczyć softmax w sposób blokowy. Praca wprowadza dwa kluczowe wzory:

1. Rozkład softmax:
   softmax(x) = exp(x - m) / Σexp(x - m)

   gdzie m jest wartością maksymalną w x.

2. Scalanie softmax:
   softmax(x ∪ y) = softmax(softmax(x) * e^(m_x - m), softmax(y) * e^(m_y - m))

   gdzie m = max(m_x, m_y)

Te wzory pozwalają Flash Attention na obliczanie częściowych wyników softmax dla każdego bloku i łączenie ich poprawnie, aby uzyskać wynik końcowy.

Szczegóły implementacji

Zagłębmy się w uproszczoną implementację Flash Attention, aby zilustrować jego podstawowe pojęcia:

import torch

<p>def flash_attention(Q, K, V, block_size=256):
batch_size, seq_len, d_model = Q.shape</p>

<p># Inicjacja macierzy wyjściowej i statystyk bieżących
O = torch.zeros_like(Q)
L = torch.zeros((batch_size, seq_len, 1))
M = torch.full((batch_size, seq_len, 1), float(&#039;-inf&#039;))</p>

<p>for i in range(0, seq_len, block_size):
Q_block = Q[:, i:i+block_size, :]</p>

<p>for j in range(0, seq_len, block_size):
K_block = K[:, j:j+block_size, :]
V_block = V[:, j:j+block_size, :]</p>

<p># Obliczanie wyników uwagi dla tego bloku
S_block = torch.matmul(Q_block, K_block.transpose(-2, -1)) / (d_model ** 0.5)</p>

<p># Aktualizacja maksymalnej wartości
M_new = torch.maximum(M[:, i:i+block_size], S_block.max(dim=-1, keepdim=True).values)</p>

<p># Obliczanie wykładnic
exp_S = torch.exp(S_block - M_new)
exp_M_diff = torch.exp(M[:, i:i+block_size] - M_new)</p>

<p># Aktualizacja sumy
L_new = exp_M_diff * L[:, i:i+block_size] + exp_S.sum(dim=-1, keepdim=True)</p>

<p># Obliczanie wyjścia dla tego bloku
O[:, i:i+block_size] = (
exp_M_diff * O[:, i:i+block_size] +
torch.matmul(exp_S, V_block)
) / L_new</p>

<p># Aktualizacja statystyk bieżących
L[:, i:i+block_size] = L_new
M[:, i:i+block_size] = M_new</p>

return O

Ta implementacja, chociaż uproszczona, ujmuje istotę Flash Attention. Przetwarza dane w blokach, utrzymując statystyki bieżące (M i L), aby poprawnie obliczyć softmax na wszystkich blokach.

Wpływ Flash Attention

Wprowadzenie Flash Attention miało głęboki wpływ na dziedzinę sztucznej inteligencji, szczególnie na duże modele językowe i aplikacje z długimi kontekstami. Niektóre z kluczowych korzyści to:

1. Zmniejszone użycie pamięci: Flash Attention zmniejsza złożoność pamięci z O(N^2) do O(N), gdzie N jest długością sekwencji. Pozwala to na przetwarzanie znacznie dłuższych sekwencji przy użyciu tego samego sprzętu.
2. Poprawiona szybkość: Poprzez minimalizowanie ruchu danych i lepsze wykorzystanie możliwości obliczeniowych GPU, Flash Attention osiąga znaczne przyspieszenia. Autorzy donoszą o przyspieszeniach nawet do 3-krotnych w przypadku treningu modelu GPT-2 w porównaniu z implementacjami standardowymi.
3. Dokładne obliczanie: W przeciwieństwie do niektórych innych technik optymalizacji uwagi, Flash Attention oblicza dokładną uwagę, a nie jej przybliżenie.
4. Skalowalność: Zmniejszony ślad pamięci pozwala na skalowanie do znacznie dłuższych sekwencji, potencjalnie do milionów tokenów.

Wpływ w świecie rzeczywistym

Wpływ Flash Attention sięga poza badania akademickie. Został szybko przyjęty w wielu popularnych bibliotekach i modelach sztucznej inteligencji:

• Hugging Face Transformers: Popularna biblioteka Transformers zawiera już implementację Flash Attention, umożliwiając użytkownikom łatwe wykorzystanie jego zalet.
• GPT-4 i dalej: Chociaż niepotwierdzone, istnieją spekulacje, że zaawansowane modele językowe, takie jak GPT-4, mogą wykorzystywać techniki podobne do Flash Attention, aby radzić sobie z długimi kontekstami.
• Modele z długimi kontekstami: Flash Attention umożliwił powstanie nowego pokolenia modeli, które mogą radzić sobie z niezwykle długimi kontekstami, takimi jak modele, które mogą przetwarzać całe książki lub długie filmy.

FlashAttention: Ostatnie rozwoje

Standardowa uwaga vs Flash Attention

FlashAttention-2

Budując na sukcesie oryginalnego Flash Attention, ta sama grupa wprowadziła FlashAttention-2 w 2023 roku. Ta zaktualizowana wersja przynosi kilka ulepszeń:

1. Dalsza optymalizacja: FlashAttention-2 osiąga jeszcze lepsze wykorzystanie GPU, sięgając do 70% teoretycznego szczytu FLOPS na GPU A100.
2. Poprawiony etap wsteczny: Etap wsteczny został zoptymalizowany, aby był prawie tak szybki, jak etap do przodu, co prowadzi do znacznych przyspieszeń w treningu.
3. Wsparcie dla różnych wariantów uwagi: FlashAttention-2 rozszerza wsparcie na różne warianty uwagi, w tym uwagę z grupowanymi zapytaniami i uwagę wielozapytaniową.

FlashAttention-3

Wydany w 2024 roku, FlashAttention-3 reprezentuje najnowszy krok w tej linii badań. Wprowadza kilka nowych technik, aby dalej poprawić wydajność:

1. Obliczenia asynchroniczne: Wykorzystanie asynchronicznej natury nowych instrukcji GPU do nakładania się różnych obliczeń.
2. Wsparcie dla FP8: Używanie obliczeń o niskiej precyzji FP8 dla jeszcze szybszego przetwarzania.
3. Przetwarzanie niezgodne: Technika redukująca błąd kwantyzacji przy użyciu formatów o niskiej precyzji.

Oto uproszczony przykład, jak FlashAttention-3 może wykorzystywać obliczenia asynchroniczne:

import torch
from torch.cuda.amp import autocast

<p>def flash_attention_3(Q, K, V, block_size=256):
with autocast(dtype=torch.float8): # Używanie FP8 do obliczeń
# ... (podobnie jak w poprzedniej implementacji)</p>

<p># Przykład obliczeń asynchronicznych
with torch.cuda.stream(torch.cuda.Stream()):
# Obliczanie GEMM asynchronicznie
S_block = torch.matmul(Q_block, K_block.transpose(-2, -1)) / (d_model ** 0.5)</p>

<p># Tymczasem, na strumieniu domyślnym:
# Przygotuj do obliczeń softmax</p>

<p># Synchronizuj strumienie
torch.cuda.synchronize()</p>

<p># Kontynuuj z obliczeniami softmax i wyjścia
# ...</p>

return O

Ten fragment kodu ilustruje, jak FlashAttention-3 może wykorzystywać obliczenia asynchroniczne i precyzję FP8. Należy zauważyć, że jest to uproszczony przykład, a rzeczywista implementacja byłaby znacznie bardziej złożona i zależna od sprzętu.

Wdrażanie Flash Attention w Twoich projektach

Jeśli jesteś zainteresowany wykorzystaniem Flash Attention w swoich projektach, masz kilka opcji:

1. Użyj istniejących bibliotek: Wiele popularnych bibliotek, takich jak Hugging Face Transformers, zawiera już implementację Flash Attention. Aktualizacja do najnowszej wersji i włączenie odpowiednich flag może być wystarczająca.
2. Własna implementacja: Dla większej kontroli lub specjalnych przypadków użycia, możesz zdecydować się na własną implementację Flash Attention. Biblioteka xformers zapewnia dobrą implementację referencyjną.
3. Optymalizacje specyficzne dla sprzętu: Jeśli pracujesz z konkretnym sprzętem (np. GPU NVIDIA H100), możesz chcieć wykorzystać funkcje specyficzne dla tego sprzętu, aby osiągnąć maksymalną wydajność.

Oto przykład, jak możesz użyć Flash Attention z biblioteką Hugging Face Transformers:

from transformers import AutoModel, AutoConfig

<p># Włącz Flash Attention
config = AutoConfig.from_pretrained(&quot;bert-base-uncased&quot;)
config.use_flash_attention = True</p>

<p># Załaduj model z Flash Attention
model = AutoModel.from_pretrained(&quot;bert-base-uncased&quot;, config=config)</p>

<p># Użyj modelu jak zwykle
# ...

Wyzwania i kierunki przyszłego rozwoju

Chociaż Flash Attention zrobił znaczne postępy w poprawie wydajności mechanizmów uwagi, nadal istnieją wyzwania i obszary wymagające dalszych badań:

1. Specyfika sprzętu: Obecne implementacje są często zoptymalizowane dla konkretnych architektur GPU. Uogólnienie tych optymalizacji na różny sprzęt pozostaje wyzwaniem.
2. Integracja z innymi technikami: Kombinowanie Flash Attention z innymi technikami optymalizacji, takimi jak pruning, kwantyzacja i kompresja modelu, jest aktywnym obszarem badań.
3. Rozszerzenie na inne dziedziny: Chociaż Flash Attention odniósł duży sukces w NLP, rozszerzenie jego korzyści na inne dziedziny, takie jak widzenie komputerowe i modele multimodalne, jest kontynuowanym wysiłkiem.
4. Teoretyczne zrozumienie: Głębsze zrozumienie, dlaczego Flash Attention działa tak dobrze, mogłoby prowadzić do jeszcze potężniejszych optymalizacji.

Podsumowanie

Poprzez inteligentne wykorzystanie hierarchii pamięci GPU i zastosowanie sztuczek matematycznych, Flash Attention osiąga znaczne poprawy zarówno w szybkości, jak i wykorzystaniu pamięci, bez poświęcania dokładności.

Jak zbadaliśmy w tym artykule, wpływ Flash Attention sięga daleko poza prostą optymalizację. Umożliwił rozwój bardziej potężnych i wydajnych modeli.