MHA, MQA и GQA: как аттеншн меняет расход памяти KV
Автоматический перевод
Эта статья была автоматически переведена с оригинальной английской версии.
MHA назначает каждой голове запросов собственную голову ключей и собственную голову значений. MQA использует одну общую голову ключей/значений для всех голов запросов. GQA использует общую голову ключей/значений внутри каждой группы голов запросов. Меньшее число голов KV уменьшает обычный кэш аттеншна и объём данных, которые декодирование считывает из него.
При оценке размера кэша проверяйте num_key_value_heads. Если учитывать только число голов запросов, можно существенно завысить размер кэша.
Прямое сравнение совместного использования голов
Предположим, что у нас восемь голов запросов и одинаковые размерность головы, число слоёв, длина последовательности и точность кэша:
| Аттеншн | Головы запросов | Головы KV | Кэш относительно MHA |
|---|---|---|---|
| Многоголовый аттеншн, MHA | 8 | 8 | 100% |
| Аттеншн с несколькими запросами, MQA | 8 | 1 | 12.5% |
| Аттеншн с группировкой запросов, GQA | 8 | 2 | 25% |
Голова KV — это одна голова ключей и одна голова значений. Проекции запросов остаются отдельными. Статья Шазира о MQA и статья Эйнсли и соавторов о GQA определяют эти изменения.
В Llama 3 70B 64 головы запросов совместно используют восемь голов KV: это восемь раз меньше исходных данных KV, чем у модели с 64 головами KV, идентичной во всём остальном. У Llama 3.1 405B 128 голов запросов и восемь голов KV, что при таком же сравнении даёт уменьшение в шестнадцать раз. Таблица архитектуры Meta содержит эти числа.
Чего уменьшение не гарантирует
KV cache в восемь раз меньшего размера не означает в восемь раз меньший общий расход памяти GPU или в восемь раз более быструю генерацию. Веса, активации, выполнение ядер и обмен данными остаются. Распределённые реализации также могут реплицировать головы KV между разделами.
В статье о GQA авторы нашли полезный компромисс между качеством и скоростью на протестированных моделях семейства T5. Перевод обученного чекпоинта MHA на GQA — это изменение архитектуры, которое требует адаптации и оценки; это не настройка выделения памяти.
Проверьте конфигурацию аттеншна в чекпоинте, используйте число голов KV в формуле размера кэша и измерьте производительность сервера при нужной длине контекста и числе одновременных запросов. При выборе между чекпоинтами учитывайте в сравнении качество решения ваших задач.
Инженерное руководство: GQA и MQA показывает схемы совместного использования голов.