Czym jest FlashAttention i jakie zyski pamięciowe zapewnia?

FlashAttention to zoptymalizowany algorytm obliczania mechanizmu uwagi (attention) w transformatorach, który znacznie zmniejsza zużycie pamięci i przyspiesza obliczenia w porównaniu z klasycznym podejściem.

Główną ideą FlashAttention jest obliczanie uwagi bez jawnego przechowywania dużych macierzy wyników pośrednich (np. macierzy wag uwagi QK^T), korzystając z przetwarzania blokowego i efektywnego wykorzystania pamięci podręcznej GPU. Pozwala to zmniejszyć szczytowe zużycie pamięci o około 3-4 razy, co jest szczególnie ważne podczas trenowania dużych modeli językowych (LLM) z długimi sekwencjami.

Przybliżenie, jeśli klasyczna uwaga wymaga pamięci O(N^2) dla sekwencji długości N, FlashAttention zmniejsza to obciążenie dzięki obliczeniom na bieżąco i optymalizacjom, co pozwala uruchamiać modele z większymi kontekstami na tych samych GPU.

W ten sposób, FlashAttention to sposób na uczynienie uwagi bardziej efektywną pod względem pamięci i szybkości, co jest krytyczne dla dużych modeli.