FlashAttention-2: Szybsze Transformery i Dłuższy Kontekst

Poniższy tekst jest tłumaczeniem mojego artykułu w języku angielskim.
Polecam również, oprócz oczywiście oryginalnej publikacji naukowej, wpis jej autorów na blogu.

Naukowcy z Uniwersytetu Stanforda opracowali nową technikę o nazwie FlashAttention-2, która może znacząco przyspieszyć trenowanie dużych modeli Transformer na dłuższych sekwencjach. W publikacji umieszczonej na ArXiv, autorzy demonstrują nawet 2-krotnie szybsze szkolenie w porównaniu do najnowszych dotychczasowych metod, takich jak FlashAttention.

Transformery stały się dominującą architekturą dla zadań związanych z przetwarzaniem języka naturalnego, ale kluczowym ograniczeniem jest ich kwadratowe skalowanie wymagań dotyczących pamięci i obliczeń wraz z długością sekwencji. Ograniczyło to większość modeli Transformer, takich jak GPT-3, do sekwencji 2048 tokenów lub mniej. FlashAttention-2 ma na celu przesunięcie tej granicy, optymalizując wzorce dostępu do pamięci i równoległość przy obliczaniu kluczowego mechanizmu atencji (ang. attention) w Transformerach.

Główne innowacje w FlashAttention-2 to:

• • Modyfikacja algorytmu w celu zmniejszenia operacji innych niż mnożenie macierzy, które są znacznie wolniejsze na GPU. Zwiększa to proporcję czasu obliczeń spędzanego na szybkim mnożeniu macierzy.
  • Równoległe obliczanie atencji na długości sekwencji, oprócz rozmiaru partii (ang. batch size) i liczby głów. Poprawia to wykorzystanie GPU dla długich sekwencji.
  • Zoptymalizowany podział pracy w obrębie każdego bloku wątków GPU, aby zmniejszyć niepotrzebny dostęp do pamięci.

W testach wydajności na GPU A100, FlashAttention-2 osiąga do 2-krotnie większą prędkość w porównaniu z oryginalnym FlashAttention i do 10-krotnie większą w porównaniu ze standardowymi implementacjami PyTorch. Gdy jest używany do trenowania modeli w stylu GPT, umożliwia przepustowość powyżej 225 TFLOPs na GPU dla długości sekwencji do 8k tokenów.

Badacze podkreślają, że FlashAttention-2 sprawia, że jest ekonomicznie opłacalne trenować modele na znacznie dłuższych sekwencjach niż dotychczas. Na przykład, można szkolić model na 16k tokenów za tę samą cenę, która wcześniej była wymagana dla 8k tokenów. Mogłoby to umożliwić modelom Transformer rozumienie całych książek, artykułów, obrazów czy filmów, a nie tylko krótkich fragmentów.

Poza kosztem, korzyści wynikają również z dłuższego kontekstu. Najnowsze modele, takie jak GPT-3 i Claude, wykazały znaczące zyski w zdolnościach do wnioskowania i zdrowego rozsądku, gdy były trenowane na sekwencjach 32-100k tokenów. FlashAttention-2 czyni te korzyści bardziej dostępnymi dla badaczy.

Patrząc w przyszłość, zespół planuje dalej optymalizować FlashAttention-2 pod kątem nowego sprzętu, takiego jak GPU H100 od Nvidia. Chcą również połączyć te niskopoziomowe optymalizacje z proponowanymi ulepszeniami do mechanizmu atencji. Końcowym celem jest całkowite usunięcie problemu z długością kontekstu dla Transformerów i umożliwienie tworzenia modeli, które mogą wnioskować na podstawie dowolnie długich dokumentów lub danych.