Twój cache KV nie ma problemu z bitami. Ma problem z geometrią.

Przy identycznym 2-bitowym dokładności, jeden decyzja o tym, który axis jest używany do kwantyzacji, zmienia wynik testu benchmarkowego z 2,88 do 63,53. Klucze i wartości wymagają przeciwstawnych podejść — a powodem jest równanie uwagi, a nie sprzęt.

Weźmy Llama-2-13B. Grupujmy jej cache key-value według rozmiaru grupy kwantyzacji 32 w dwa bity, pozostawiając wszystko inne w miejscu — ten sam model, ten sam budżet bitowy, te same rozmiary grup, te same testy.

W zależności od jednego elementu decyzji implementacyjnej, wyniki dokładności CoQA wynoszą albo 2,88, albo 63,53. Wynik przy pełnej precyzji wynosi 66,37.

Decyzja nie dotyczy ilości użytych bitów. Pytanie brzmi: jaki axis wybieramy do grupowania przy obliczaniu każdego czynnika skali? Kiedy wybieramy kanał jako wymiar grupowania (klucze) i token jako wymiar grupowania (wartości), kończymy w granicach czterech punktów od wyniku pełnej precyzji. Jeśli odwrócimy którąkolwiek z tych decyzji, doświadczymy straty jakości. Jeśli odwrócimy obie te decyzje, model przestaje działać.

Cztery sposoby wydatkowania tych samych 2 bitów na ten sam cache. Wyniki z ablacją KIVI na Llama-2-13B przy rozmiarze grupy 32.

Kwantyzacja jest zwykle postrzegana jako jeden przełącznik: 8 bitów, 4 bity, 2 bity, z gładkim kosztem dokładności. Wewnątrz cache KV nie jest to takie proste. Chodzi o wybór układu współrzędnych, a różne układy stosuje się do kluczy i wartości. Ten artykuł wyjaśnia, dlaczego. Krótko mówiąc: błąd kwantyzacji zależy od zakresu wartości w grupach; klucze i wartości mają bardzo różną strukturę; ludzie często popełniają błąd, ponieważ nie można wywnioskować prawidłowej osi z dystrybucji wartości w ogóle. Trzeba spojrzeć, jak błąd zmienia się po tym, jak uwaga go spożytkuje. To daje ogólny principle kompresji aktywacji pośrednich i dobry powód, aby podważyć błąd rekonstrukcji jako proxy dla jakości.

Dlaczego cache KV jest miejscem, gdzie to ugryzie

Podczas fazy generacji, transformator przechowuje wszystkie dane projekcji klucza i wartości (KV) tokenów, które wcześniej przetworzył, w cache, aby nie musieć ponownie obliczać tych danych. Ten cache rośnie liniowo z długością kontekstu i rozmiarem partii. Ostatecznie spowoduje to, że cache będzie większy niż sam model.

Ten wzrost można łatwo zidentyfikować, patrząc na zużycie pamięci różnych części modelu. W analizie KVQuant LLaMA-7B, wagi stanowią około 98 procent pamięci przy długości sekwencji 512, a aktywacje — 2 procent. Przy kontekście 128K stosunek odwraca się do około 16 procent wag i 84 procent cache KV. Kiedy spojrzymy na analizę OPT-175B, cytowaną przez autorów KIVI, znaleźli oni podobne wyniki. Konkretnie, przy rozmiarze partii 512 i promocie 512-tokenowej, cache KV osiąga 1,2 TB — kilka razy większy niż rozmiar wag modelu.

Jednak pojemność to tylko połowa problemu. GPU musi odczytać cały cache KV z pamięci urządzenia dla każdego wygenerowanego tokenu. Oznacza to, że podczas odczytywania cache KV przez GPU, rdzenie obliczeniowe pozostają bezczynne. Zmniejszając rozmiar cache, zwiększamy dostępną przestrzeń obliczeniową i zmniejszamy czas spędzany na oczekiwaniu na transfery danych.

Czym tak naprawdę jest błąd kwantyzacji

Jednorodna kwantyzacja liczb całkowitych jest prosta matematycznie. Dla grupy liczb, rejestrujemy najmniejszą liczbę jako punkt zero, a następnie dzielimy zakres grupy przez liczbę poziomów, które można reprezentować, aby uzyskać rozmiar kroku. Następnie zaokrąglamy każdy element do najbliższego kroku. Dwa natychmiastowe wyniki następują. Po pierwsze, błąd na element jest ograniczony do połowy kroku. Po drugie, rozmiar kroku to zakres grupy podzielony przez 2^B – 1. Przy 2 bitach mamy tylko 4 poziomy, aby pokryć jakiś zakres, jaki istnieje w grupie. Więc element, który jest stu razy większy w porównaniu z sąsiadami, nie tylko działa źle. Nadmienia rozmiar kroku dla wszystkich innych elementów, które dzielą tę samą grupę, i wszystkie stają się bardziej szorstkie razem. Grupa jest jednostką uszkodzenia. Wybór osi oznacza decyzję, które elementy cierpią razem. Innymi słowy, nie jest to już „ile bitów mogę sobie pozwolić?” ale „gdzie są ekstremalne wartości i czy mogę je odizolować?”

Klucze: Outliery żyją w stałych kanałach

Duże modele językowe zawierają aktywacje, które są niezwykle duże w porównaniu z większością aktywacji. Sun i współpracownicy skatalogowali te bardzo duże aktywacje w różnych rodzinach modeli: w Mixtral 8x7B, największa wielkość jest w pobliżu 7000, a medianowa wielkość cechy wynosi około 0,3 — czyli około czterech rzędów wielkości różnicy. Są one bardzo rzadkie; pozostają w stałych wymiarach, które rzadko zmieniają się z wejściem, i nie są przypadkowe. Działają jako niejawne uprzedzenia, i są tym, co skupia uwagę na kilku tokenach: zachowanie uwagi. W cache kluczy, struktura ta jest bardzo wyraźna: określone kanały przenoszą bardzo duże wielkości w sposób ciągły w całej sekwencji tokenów. Grupujmy wzdłuż tokenów, a każda grupa zawiera te outliery kanałów, więc każda grupa ma rozmiar kroku ustawiony przez outliery, a wszystkie zwykłe kanały płacą za to. Grupujmy wzdłuż kanałów, a outliery kanałów tworzą własne grupy. Ich wewnętrzny zakres jest duży, ale samodzielny; zwykłe kanały pozostają same. Wyniki są zgodne. Średnio przez warstwy i głowy na Llama-2-13B, KIVI raportuje błąd rekonstrukcji klucza 13,67 przy grupowaniu per-token wobec 4,55 przy grupowaniu per-kanał. A co więcej — błąd wyniku uwagi 47,00 wobec 9,60. Kwantyzacja kluczy per-token produkuje około pięć razy większy błąd wyniku. Wyniki zgadzają się z istotnymi metrykami dla kluczy; kwantyzacja kanałowa wyróżnia się w obu aspektach.

Wartości: Tam, gdzie intuicja się łamie

Cache wartości nie wykazuje kanałowo-outlierowego wzorca. Wydaje się, że jest dość płaski. Samodzielnie, według argumentu zakresu, można by oczekiwać, że albo któraś z tych osi wyprodukuje podobną jakość kompresji.

One nie są. Niezależnie od tego, jak zarządzane są klucze (wyniki 2,80 i 2,88), kompresja per-kanałowa wartości powoduje załamanie modelu.

I tutaj jest haczyk: jeśli zmierzymy tę stratę za pomocą surowego błędu rekonstrukcji na oryginalnym tensory, dla którego każda wartość była skompresowana, kwantyzacja per-kanałowa wartości wygląda trochę lepiej, bo 3,73 wobec 4,57. Jeśli walidujemy kompresję w oczywisty sposób, wybralibyśmy konfigurację, która niszczy model.

Błąd kwantyzacji cache wartości na Llama-2-13B, zmierzony na dwa sposoby. Metryka tensora przechowywanego i metryka wyjścia zużytego różnią się o więcej niż rząd wielkości.

Rozwiązanie polega na tym, że cache wartości nie jest czytany bezpośrednio. Jest zużywany przez produkt macierzy: wynik uwagi jest ważoną sumą wektorów wartości wzdłuż tokenów, z wagami uwagi softmax jako wagami. Ze względu na to, istotny błąd to ten, który jest wprowadzony podczas tego procesu, a nie wewnątrz samych tensorów. Zmierzony w kategoriach wyniku uwagi, kolejność była całkowicie odwrócona. Względny błąd raportowany przez KIVI dla wyniku uwagi z powodu kwantyzacji wektorów wartości per-token wyniósł 3,55 wobec 49,89 dla kwantyzacji per-kanał — ponad czternaście razy większy dla tego, co wydawało się lepszym wyborem na podstawie tego, jak dobrze było to skompresowane.

Wyjaśnienie to rzadkość uwagi, którą oni zmierzyli na 84,3 procent. Większość informacji zawartych w wyjściu może być przypisana do niewielkiej liczby bardzo ważnych tokenów. Kwantyzacja per-token ogranicza błąd każdego tokenu do tego tokenu, więc błędy na nieistotnych tokenach mnożone są przez wagę uwagi bliską zeru i skutecznie znikają. Kwantyzacja per-kanałowa rozprzestrzenia błąd każdego tokenu na wspólną skalę kanałową, więc źle reprezentowane tokeny zanieczyszczają reprezentację tych, które są istotne. Rzadkość, która sprawia, że uwaga jest wydajna, jest tą samą właściwością, która sprawia, że kwantyzacja per-token jest bezpieczna.

Przenoszalna lekcja jest szersza niż cache KV: mierzyć błąd kompresji tam, gdzie tensor jest zużywany, a nie tam, gdzie jest przechowywany. Niejawne założenie, które jest robione przez błąd rekonstrukcji, to to, że każdy komponent tensora ma równy ciężar, przyczyniając się do końcowego wyjścia. Uwaga nie robi tego jawnie. Każda operacja downstream, która waży, bramuje lub rozprzestrzenia swój wejście, łamie to założenie. Czytelnicy, którzy znają mój poprzedni artykuł dotyczący ślepych punktów w metrykach oceny w systemach pobierania, rozpoznają, że te wyniki są podobne do wcześniej opisanych niepowodzeń: łatwo obliczalne metryki, które raportują coś innego niż to, co było zamierzone.

Rotary Embeddings komplikują klucze

Istnieją pewne problemy z użyciem Rotary Position Embeddings (RoPE). RoPE obraca pary kanałów w zależności od względnej pozycji każdego tokenu. To mieszanie częściowo rozpuszcza strukturę kanałową, która sprawiła, że kwantyzacja per-kanałowa kluczy działała. Kanał outlier zostaje obrócony w swoich sąsiadach, a sąsiedzi odziedziczą zakres. Odpowiedź KVQuant to kolejność: kwantyzuj klucze przed zastosowaniem obrotu, a następnie zastosuj RoPE po dekwantyzacji. Wraz z kwantyzacją per-kanałową kluczy, niejednolitymi typami danych i izolacją małej frakcji outlierów, to pozwala im uzyskać mniej niż 0,1 degradację perplexity przy 3 bitach i umożliwia obsługę LLaMA-7B do 1 miliona tokenów kontekstu na jednym A100-80GB.

Ważne jest również zrozumienie poziomu wpływu RoPE. Autorzy artykułu „RotateKV” zgłosili wzrost o 145 procent błędów kwantyzacji, gdy RoPE zostało dodane, i zauważyli, że kanały outlier różnią się między głowami uwagi — dlatego też zastosowanie jednej wspólnej macierzy obrotu wszędzie jest niewystarczające, a obroty dostosowane do głowy robią lepiej.

Podatek systemowy i dlaczego nie jest to detail

Kwantyzacja per-token nadaje się dobrze do dekodowania. Każdy token przychodzi; kwantyzujemy go, dodajemy go do sekwencji (wzdłuż wymiaru tokenu), nic innego się nie porusza.

Jednak kwantyzacja per-kanałowa nie pasuje. Ponieważ statystyki kanału rozciągają się na tokeny, które jeszcze nie zostały wygenerowane, nie można obliczyć czynnika skali, gdy token przychodzi. Rozwiązanie KIVI polega na przechowywaniu najnowszych tokenów — do 128 — w pełnej precyzji w buforze resztkowym i kwantyzowaniu w grupach, gdy tylko wystarczająco dużo z nich się nagromadzi.

Okazuje się, że bufor resztkowy staje się nośny, a nie tylko przypadkowy. Na GSM8K z Llama-2-7B, wynik w pełnej precyzji wynosi 13,50. W pełni skwantyzowany do 2 bitów z odpowiednimi osiami, wynosi 5,76. Te same osie i te same bity, plus bufor resztkowy ostatnio wygenerowanych tokenów w pełnej precyzji, wynoszą 12,74. Przesuwający się okno ostatnio wygenerowanych tokenów w pełnej precyzji odzyska większość tego, co zostało utracone z powodu agresywnej kwantyzacji na trudnych wieloetapowych problemach — co miałoby sens, gdybyśmy rozważyli, które tokeny były uwzględniane przez łańcuch operacji arytmetycznych.

Istnieje znacząca korzyść z robienia wszystkich tych rzeczy poprawnie — jak donosi KIVI, 2,6 razy mniej szczytowego zużycia pamięci dla Llama-2-7B, umożliwiając partie o rozmiarze do 4 razy większym, a także 2,35 do 3,47 razy lepszy przepływ na rzeczywistym zadaniu usługi.

Co z tym zrobić

1. Nigdy nie używaj jednego kwantyzatora dla obu. Używaj różnych kwantyzatorów dla kluczy (per-kanał) i dla wartości (per-token). Potok, który stosuje jeden kwantyzator do „cache KV”, najprawdopodobniej już poświęcił większość możliwej jakości przy użyciu małej liczby bitów do reprezentowania każdej wartości.
2. Kwantyzuj klucze przed RoPE. To jest kwestia poprawności, a nie preferencji.
3. Przechowuj okno w pełnej precyzji ostatnio wygenerowanych tokenów. Chociaż przechowywanie takiego okna zajmuje bardzo mało pamięci w porównaniu z tym, jak duży może być cache, to właśnie ta okolica generuje większość dokładności dla trudnych zadań.
4. Nie waliduj na błędzie rekonstrukcji. Zawsze waliduj na podstawie wyniku uwagi lub na podstawie końcowej wydajności zadania. Metryka przechowywania nie jest po prostu szumem — dla wartości wskazuje w złym kierunku.
5. Nie waliduj na krótkich benchmarkach wielokrotnego wyboru. Autorzy KIVI celowo unikają zadań zamkniętych, takich jak MMLU, w tej ocenie, ponieważ jeden krok dekodowania czytający wyjście logitów ledwo ćwiczy cache w ogóle. Każda ocena, która nie buduje cache w czasie i nie wykonuje generacji z niego, nie będzie w stanie zaobserwować niepowodzeń tkwiących w projekcie systemu.

Gdzie zmierza praca

Chociaż pozostaje jeszcze coś do zrobienia dotyczącego geometrycznej natury problemu, wielu badaczy nadal studiuje sposoby, w jakie outliery kanałów są rozłożone wśród różnych głów transformatora, i w jaki sposób ograniczenia sprzętowe wpływają na to, które grupowania są najtańsze: InnerQ zamyka normalizację kanałową klucza w wagach klucza i zapytania podczas wstępnego wypełnienia. Nie jest więc generowany żaden dodatkowy nakład na runtime. Ponadto, InnerQ przechowuje okna o wysokiej precyzji zarówno dla ostatnio wygenerowanych tokenów, jak i tokenów sink uwagi. W ten sposób InnerQ eliminuje możliwość, że outliery w kanale sink zanieczyszczą sąsiednie kanały.

Inni proponują, że zamiast przechowywać cały cache, powinniśmy przechowywać tylko tyle informacji, aby móc ponownie utworzyć klucz i/lub wartości na żądanie z mniejszej reprezentacji cache.

Wreszcie, należy pamiętać, że dokładność nie jest jedynym parametrem, którym kwantyzacja się zajmuje. Niedawno opublikowane badania wykazały degradację wyrównania wynikającą z kwantyzacji cache KV. Co więcej, to badanie udokumentowało degradację wyrównania nawet w środowiskach produkcyjnych vLLM z cache FP8 oraz protokołem odzyskiwania bez treningu, który przywrócił do 97 procent tego, co zostało utracone pod względem wyrównania. Jak zatem można się spodziewać, że nawet jeśli konfiguracja utrzyma wyniki benchmarkowe, niekoniecznie oznacza to, że zachowuje wszystkie inne istotne parametry, o które dbasz.

Ogólny principle

Pomysł kwantyzacji został ukształtowany jako „budżet precyzji”: ile bitów mogę poświęcić? Cache KV pokazuje, że bardziej użyteczne pytanie jest strukturalne. Precyzja jest przydzielana w grupach; grupa jest jednostką uszkodzenia, a oś, wzdłuż której grupujemy, decyduje, które elementy dzielą swoje losy. Prawidłowa oś to ta, wzdłuż której tensor jest zużywany, czyli sposób, w jaki używasz tensora, a nie to, jak wygląda on, gdy jest przechowywany w pamięci. Klucze są używane za pomocą obliczenia iloczynu skalarnego z zapytaniem. Jeden uszkodzony kanał zatruwa wszystkie wyniki. Wartości są zużywane przez obliczenie średniej ważonej wzdłuż tokenów, z wagami uwagi softmax jako wagami.

Dwa tensory o identycznych wymiarach i wygenerowane przez dwie kolejne warstwy są traktowane inaczej. Warto zapytać się o każdej aktywacji, którą planujesz skompresować: jaka operacja to usuwa, i czy moje grupowanie szanuje to?