[Блогпост]
Чтобы не пересчитывать каждый раз все заново при авторегрессивной генерации, в трансформерах используются kv-кэши, ключи и значения, посчитанные для прошлых токенов, сохраняются в памяти и используются при подсчете attention для каждого нового токена.
Однако есть проблема - если последовательность длинная, кэш слишком тяжелый становится.
Для 🦙 (1 или 2) 7b для 10к токенов потребуется 5Gb памяти (2 * 2 * 32 * 32 * 128 * 10000 ≈ 5GB). А для более крупных моделей и , если контекст переваливает за миллион, страшно представить сколько.
Существуют разные подходы к сжатию KV-кэшей, и один из наиболее очевидных - квантизация кэша, преставление ключей и значений в более низкой битности.
Метод
В работе KIVI предложили 2-битное ассиметричное (с zero point) квантование кэшей. Ключи квантуются поканально, значения - по токенно. Поканальное квантование помогает бороться с outliers.
В HF использует по-токенное квантование для ключей и значений из соображений скорости, несколько жертвуя качеством. Кроме того, чтобы немного улучшить качество, самые свежие 128 токенов держат в исходной точности. Они обычно самые важные для текущего токена, и не добавляют много памяти для длинного контекста.
HF поддерживает два метода квантизации:
- quanto (простой rount-to-nearest)
- hqq
Оба метода data-free.
Валидация
На замерах перплексии на Llama-2-7b на PG-19, 4-битный HQQ не просаживает качество, а
quanto_int4 умеренно просаживает. 2-битные квантизации сильно роняют качество. На LongBench в 4 битах почти нет просадки по качеству.На длинных последовательностях удается достичь до 2.5x экономии потребляемой памяти.
Скорость, генерации, однако замедляется так как на текущий момент нет оптимизированных ядер. Просадка по скорости 30-50% судя по графикам (крайняя правая точка на графике number_of_tokens/sec явно забагованная).
Кроме того, данная стратегия помогает именно на стадии генерации, а заполнение промпта (pre-fill) приходится делать стандартным способом.
Метод совместим с Flash Attention.
Использование
Достаточно установить quanto, и прописать аргументы в методе
.generate.