DuoAttention: эффективный метод для работы с длинными контекстами в LLM
DuoAttention – метод, который решает проблему LLM при работе с длинными текстами. Эта неэффективность проявляется в большом потреблении памяти и замедлении работы модели. DuoAttention позволяет сократить использование памяти и ускорить обработку длинных текстов без ущерба для точности модели.
Он основан на разделении attention heads в LLM на два типа: «Retrieval Heads» и «Streaming Heads»:
🟢 Retrieval Heads – это небольшая часть attention heads, критически важная для обработки длинных контекстов, которым требуется полное внимание ко всем токенам.
🟢 Streaming Heads, напротив, составляют большинство attention heads и фокусируются преимущественно на последних токенах и так называемых “точках притяжения” внимания (attention sinks).
DuoAttention использует эту дихотомию, применяя полный KV-кэш только к Retrieval Heads, а для Streaming Heads использует облегченный KV-кэш постоянной длины, который хранит только “точки притяжения” и последние токены.
Проведенные эксперименты показывают, что DuoAttention может снижать использование памяти до 2,55 раз для моделей MHA и 1,67 для моделей GQA, а также ускоряет декодирование в 2,18 раз для моделей MHA и 1,50 для моделей GQA.
В сочетании с квантованием DuoAttention позволяет модели Llama-3-8B декодировать текст с длиной контекста 3,3 млн. токенов на одном GPU A100, это примерно 6,4-кратное увеличение емкости по сравнению со стандартным развертыванием FP16 с полным вниманием.
▶ Модели, которые поддерживают инференс с патчем DuoAttention:
🟢Llama-2-7B-32K-Instruct;
🟢Llama-3-8B-Instruct-Gradient-1048k;
🟢Llama-3-8B-Instruct-Gradient-4194k;
🟠Mistral-7B-Instruct-v0.2;
🟠Mistral-7B-Instruct-v0.3;
🟠Meta-Llama-3.1-8B-Instruct.
▶ Установка окружения для запуска инференса с DuoAttention :
📌Лицензирование: MIT License.