Ваш KV-кеш не має проблеми з бітами. У нього проблема з геометрією.

При однаковій точності 2 біти одне рішення про вибір осі, яку потрібно квантувати, може змінити результати бенчмарку з 2,88 до 63,53. Ключі та значення потребують протилежного лікування — і причиною цього є рівняння уваги, а не апаратне забезпечення.

Розглянемо Llama-2-13B. Групуємо ключ-значення кеш по розміру групи квантування 32 у два біти, залишаючи все інше на місці — ту саму модель, той самий бюджет біт, ті самі розміри групи, ті самі бенчмарки.

Залежно від одного елементу реалізації, результати точності CoQA можуть бути як 2,88, так і 63,53. Результат з повною точністю становить 66,37.

Рішення не полягає в тому, скільки бітів використовується в цілому. Питання полягає просто в тому, яку вісь обрати для групування при розрахунку кожного коефіцієнта масштабування? Коли ви вирішуєте використовувати канал як розмір групи (ключі) і токен як розмір групи (значення), ви потрапляєте в діапазон чотирьох пунктів від результату повної точності. Якщо ви зміните будь-який з цих виборів, ви відчуєте втрату якості. Якщо ви зміните обидва цих виборів, модель більше не працює.

Чотири способи витратити ті самі 2 біти на одному й тому ж кеші. Результати з KIVI на Llama-2-13B при розмірі групи 32.

Квантування зазвичай вважається одним з регуляторів: 8 біт, 4 біти, 2 біти, з гладкою вартістю точності. Усередині KV-кеша все інакше. Це вибір систем координат, і різні системи застосовуються до ключів і значень. Ця стаття пояснює, чому. Коротко: похибка квантування залежить від діапазону значень у групах; ключі та значення мають дуже різну структуру; і люди часто спотикаються, оскільки ви не можете вивести правильну вісь з розподілу значень зовсім. Ви повинні дивитися, як помилка змінюється після того, як увага її споживає. Це дає загальне правило для стиснення проміжних активацій і хорошу причину сумніватися у похибці реконструкції як заміні якості.

Чому KV-кеш так важливий

Під час фази генерації трансформер зберігає всі дані проєкції ключ-значення (KV) токенів, які він раніше обробив, у кеші, щоб не довелося повторно обчислювати ці дані. Цей кеш зростає лінійно з довжиною контексту та розміром партії. Зрештою, це призведе до того, що кеш зросте більшим за саму модель.

Цей зростання можна легко ідентифікувати, розглянувши витрати пам’яті різних частин моделі. У аналізі KVQuant LLaMA-7B ваги становлять близько 98 відсотків пам’яті при довжині послідовності 512, а активації — 2 відсотка. При контексті 128K співвідношення змінюється до близько 16 відсотків ваг і 84 відсотків кешу KV. Коли ми розглядаємо аналіз OPT-175B, цитований авторами KIVI, вони знайшли подібні результати. Зокрема, при розмірі партії 512 з промптом 512 токенів кеш KV досягає 1,2 ТБ — кілька разів більше розміру ваг моделі.

Однак потенціал є тільки половinou проблеми. GPU повинен прочитати весь кеш KV з пристрою пам’яті для кожного генерованого токена. Це означає, що під час читання кешу KV ядра обчислення простоюють. Таким чином, зменшення загального розміру кешу збільшує доступну обчислювальну потужність і зменшує час, витрачений на очікування передачі даних.

Що таке похибка квантування

Уніформне цілісне квантування математично просте. Для групи чисел ви реєструєте найменше число як нульову точку, а потім ділите діапазон цієї групи на кількість рівнів, які можна представити, щоб отримати розмір кроку. Ви потім округлюєте кожен елемент до найближчого кроку. Дві негайні результати слідують. По-перше, похибка на елемент обмежена половinou кроку. По-друге, розмір кроку дорівнює діапазону групи, розділеному на 2^Б − 1. При 2 бітах у вас є тільки 4 рівні для покриття будь-якого розмаїття, яке існує в цій групі. Отже, елемент, який у сотню разів більший порівняно зі своїми сусідами, не просто погано виконує свою функцію. Він надуває розмір кроку для всіх інших елементів, які ділять одну й ту ж групу, і всі вони стають грубішими разом. Група — це одиниця пошкодження. Вибір осі означає вирішення питання, які елементи будуть страждати разом. Інакше кажучи, питання вже не “скільки бітів я можу собі дозволити?”, а “де екстремальні значення і чи можу я їх ізолювати?”

Ключі: аутсайдери живуть у фіксованих каналах

Великі мовні моделі містять активації, які незвичайно великі порівняно з більшістью активацій. Сан і колеги каталогізували ці дуже великі активації у різних сім’ях моделей: у Mixtral 8x7B найбільша величина становить близько 7000, а медіана величини функції становить близько 0,3 — приблизно чотири порядки величини різниця. Ці аутсайдери дуже рідкісні; вони залишаються фіксованими у розмірах, які рідко змінюються з входом, і вони не є випадковими. Вони діють як неявні упередження, і вони те, що фокусують увагу на декількох токенах: поведінка аутсайдерів. У кеші ключів ця структура дуже очевидна: конкретні канали несуть дуже великі величини послідовно через кожен токен у послідовності. Групування по токенам, і кожна група містить ці аутсайдерські канали, тому крок кожного групи встановлюється аутсайдерами, і всі звичайні канали платять за це. Групування по каналам, і аутсайдерські канали утворюють свої власні групи. Їх внутрішній діапазон великий, але самодостатній; звичайні канали залишаються наодинці. Результати збігаються. У середньому по шарах і головам на Llama-2-13B KIVI повідомляє про похибку реконструкції ключа 13,67 при групуванні по токенам проти 4,55 при групуванні по каналу, і — що ще важливіше — похибку результату уваги 47,00 проти 9,60. Квантування ключів по токенам дає у п’ять разів більшу похибку результату.

Значення: де інтуїція ламає

Кеш значень не показує канал-аутсайдерський шаблон. Він здається досить плоским. Сам по собі, за аргументом діапазону, ми могли б очікувати, що будь-яка з цих осей дасть подібну якість стиснення.

Але вони не дають. Незалежно від того, як реалізовано керування ключами (результати 2,80 і 2,88), стиснення по каналу значень призводить до колапсу моделі.

І ось підхід: якщо ви виміряєте цю втрату за допомогою похибки реконструкції на початковому тензорі, для якого кожне значення було стиснено, квантування значень по каналу насправді виглядає трохи краще, з 3,73 проти 4,57. Якщо ви валідували своє стиснення очевидним способом, ви б обрали конфігурацію, яка знищує модель.

Похибка квантування кешу значень на Llama-2-13B, виміряна двома способами. Метрика збереженого тензора і метрика споживаного виходу розрізняються більш ніж на порядок.

Рішення полягає в тому, що кеш значень ніколи не читається напряму. Вони споживаються матричним добутком: результат уваги є зваженою сумою векторів значень по токенам, з результатами уваги softmax як вагами. Через це відповідна похибка — це та, яка вводиться під час цього процесу, а не всередині самих тензорів. Виміряна за результатом уваги, порядок був повністю змінений. Відносна похибка, зазначена KIVI для результату уваги через квантування векторів значень по токенам, становила 3,55 порівняно з 49,89 для квантування по каналу — більш ніж у чотирнадцять разів вище для того, що здавалося кращим вибором на основі того, як добре воно було стиснено.

Пояснення полягає в спарсності уваги, яку вони виміряли як 84,3 відсотка. Більшість інформації, яку містить вивід, можна віднести до кількох дуже важливих токенів. Квантування по токенам обмежує похибку кожного токена цим токеном, тому похибки на неважливих токенах множаться на майже нульові ваги уваги і фактично зникають. Квантування по каналу розмазує похибку кожного токена по спільному каналу, тому погано представлені токени забруднюють представлення тих, які мають значення. Спарсність, яка робить увагу ефективною, є тією ж властивістю, яка робить квантування по токенам безпечним.

Переносимий урок ширший за кеш KV: виміряйте похибку стиснення там, де тензор споживається, а не там, де він зберігається. Неявне припущення, зроблене похибкою реконструкції, полягає в тому, що кожна складова тензора має рівну вагу при внесенні до кінцевого виходу. Увага явно не робить цього. Будь-яка операція вниз по потоку, яка вагою, брамою або спарсенізацією вводить свій вхід, порушує це припущення. Читачі, знайомі з моєю попередньою статтею щодо сліпих плям у метриках оцінки в системах пошуку, визнають, що ці результати подібні до раніше описаних невдач: легко обчислювані метрики, які повідомляють про щось інше, ніж те, що було задумано.

Ротарі-ембеддінги ускладнюють ключі

Є деякі питання щодо використання ротарі-позиційних ембеддінгів (RoPE). RoPE обертає пари каналів на основі відносної позиції кожного токена. Це змішування частково розчиняє фіксовану структуру каналу, яка зробила квантування ключів по каналу працювати з самого початку — аутсайдерський канал обертається в своїх сусідів, і сусіди успадковують діапазон. Відповідь KVQuant — порядок: квантуйте ключі до застосування обертання, і застосуйте RoPE після де-квантування. Разом з квантуванням ключів по каналу, неуніформними типами даних і ізоляцією малих аутсайдерів це отримує менше 0,1 погіршення перплексії при 3 бітах і дозволяє служити LLaMA-7B до 1 мільйона токенів контексту на одному A100-80GB.

Також важливо зрозуміти рівень впливу від RoPE. Автори статті “RotateKV” повідомили про збільшення на 145 відсотків похибки квантування після додавання RoPE і зазначили, що аутсайдерські канали відрізняються між головами уваги — саме тому застосування однієї спільної матриці обертання всюди є недостатнім, а адаптивні до голови обертання працюють краще.

Системний податок і чому це не деталь

Квантування по токену підходить для декодування. Кожен токен прибуває; ви квантуєте його, додаєте до послідовності (вздовж розміру токену), нічого іншого не рухається.

Однак квантування по каналу не підходить. Оскільки статистика каналу охоплює токени, які ще не були сгенеровані, ви не можете обчислити коефіцієнт масштабування, коли токен прибуває. Робота KIVI полягає в тому, щоб зберігати найбільш недавні токени — до 128 — у повній точності в буфері залишків, і квантувати групи, коли накопичується достатньо.

Як виявилося, буфер залишків стає важливим, а не просто випадковим. На GSM8K з Llama-2-7B повна точність становить 13,50. Повністю квантізовано до 2 біт з правильними осями, воно оцінюється в 5,76. Те ж саме осі і ті ж біти, плюс буфер залишків недавно згенерованих токенів у повній точності, оцінюються в 12,74. Слизький вікно недавно згенерованих токенів у повній точності може відновити більшу частину того, що було втрачено через агресивне квантування на складних багаторівневих задачах — що мав би сенс, якщо ми розглянемо, які токени були звернені ланцюжком арифметичних операцій.

Є значна вигода від правильного виконання всіх цих дій — як повідомляє KIVI, 2,6 рази менше пікової витрати пам’яті для Llama-2-7B, що дозволяє збільшити розмір партії до 4 разів, а також на 2,35-3,47 рази краще пропускна здатність на реальному сервісному завданні.

Що робити з цим

1. Ніколи не використовувати один квантізатор для обох. Використовуйте різні квантізатори для ключів (по каналу) і для значень (по токену). Конвейєр, який застосовує один квантізатор до “кешу KV”, вже, ймовірно, пожертвував більшістю можливої якості при використанні малої кількості біт для представлення кожного значення.
2. Квантуйте ключі до застосування RoPE. Це питання коректності, а не питання уподобань.
3. Зберігайте вікно повної точності недавно згенерованих токенів. Хоча зберігання такого вікна займає дуже мало пам’яті порівняно з тим, як великий може бути кеш, саме ця область генерує більшу частину точності для складних завдань.
4. Не валідуйте на похибці реконструкції. Завжди валідуйте на основі результату уваги або кінцевої продуктивності завдання. Метрика зберігання не просто шумна — для значень вона вказує у неправильному напрямку.
5. Не валідуйте на коротких контекстних багаторівневих бенчмарках. Автори KIVI навмисно уникають закритих завдань, таких як MMLU, для цієї оцінки, оскільки один крок декодування, який читає вивід логітів, майже не використовує кеш зовсім. Будь-яка оцінка, яка не будує кеш з часом і не виконує генерацію з нього, ніколи не зможе спостерігати невдачі, закладені в вашому системному дизайні.

Куди рухається робота

Хоча залишається ще робота щодо геометричної природи проблеми, багато дослідників продовжують вивчати, як аутсайдерські канали розподілені між різними головами трансформера, і як апаратні обмеження впливають на те, які групування є найдешевшими: InnerQ згортає канально-орієнтовану нормалізацію ключа у ваги ключа і запиту під час попереднього заповнення. Тому жодних додаткових витрат на виконання не виникає. Крім того, InnerQ зберігає вікна високої точності як для недавно згенерованих токенів, так і для токенів-мішків уваги. Роблячи це, InnerQ усуває можливість того, що аутсайдери в каналі-синкі будуть забруднювати сусідні канали.

Інші пропонують, що замість зберігання всього кешу ми повинні зберігати тільки достатньо інформації, щоб перематеріалізувати ключ і/або значення(я) на вимогу з меншого кешованого представлення.

Нарешті, важливо пам’ятати, що точність не єдина параметр, який квантування впливає. Нещодавно опубліковані дослідження продемонстрували погіршення вирівнювання, що результатувало з квантуванням кешу KV. Крім того, це дослідження задокументувало погіршення вирівнювання навіть у виробничих середовищах обслуговування vLLM, які використовують кеші FP8 разом з протоколом відновлення без навчання, який відновлює до 97 відсотків того, що було втрачено щодо вирівнювання. Отже, хоча конфігурація може зберігати свої результати бенчмарку, це не означає, що вона зберігає всі інші важливі параметри, про які ви піклуєтеся.

Загальне правило

Ідея квантування була сформульована як “бюджет точності”: скільки бітів я можу собі дозволити пожертвувати? Кеш KV показує, що більш корисне питання — структура. Точність розподіляється в групах; група — це одиниця пошкодження, і вісь, по якій ви групуєте, визначає, які елементи доляться спільною долею. Правильна вісь — це та, по якій ваш тензор споживається, тобто те, як ви використовуєте свій тензор, а не те, як ваш тензор виглядає при збереженні в пам’яті. Ключі використовуються через обчислення скалярного добутку проти запиту. Одночасно пошкоджений канал отруює всі результати. Значення споживаються через розріджену зважену суму по токенам. Отже, одночасно пошкоджений токен просто зважується. Два тензори однакових розмірів і згенерованих двома послідовними шарами обробляються по-різному. Варто запитати про будь-яку активацію, яку ви плануєте стиснути: яка операція усуває це, і чи відповідає моє групування цьому?

Два тензори однакових розмірів і згенерованих двома послідовними шарами обробляються по-різному. Варто запитати про будь-яку активацію, яку ви плануєте стиснути: яка операція усуває це, і чи відповідає моє групування цьому? Отже, одне пошкоджене значення просто зважується.