DuoAttention: эффективный метод для работы с длинными контекстами в LLM.

DuoAttention – метод, который решает проблему LLM при работе с длинными текстами. Эта неэффективность проявляется в большом потреблении памяти и замедлении работы модели. DuoAttention позволяет сократить использование памяти и ускорить обработку длинных текстов без ущерба для точности модели. 

DuoAttention: эффективный метод для работы с длинными контекстами в LLM.

 Он основан на разделении 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: эффективный метод для работы с длинными контекстами в LLM.

Проведенные эксперименты показывают, что DuoAttention может снижать использование памяти до 2,55 раз для моделей MHA и 1,67 для моделей GQA, а также ускоряет декодирование в 2,18 раз для моделей MHA и 1,50 для моделей GQA. 

В сочетании с квантованием DuoAttention позволяет модели Llama-3-8B декодировать текст с длиной контекста 3,3 млн. токенов на одном GPU A100, это примерно 6,4-кратное увеличение емкости по сравнению со стандартным развертыванием FP16 с полным вниманием.

В Телеграм канале @ai_machinelearning_big_data подробно описан этот метод: https://t.me/ai_machinelearning_big_data/

DuoAttention: эффективный метод для работы с длинными контекстами в LLM.

▶️ Модели, которые поддерживают инференс с патчем 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 :

conda create -yn duo_demo python=3.10
conda activate duo_demo

# Install DuoAttention
pip install -e .

conda install -y git
conda install -y nvidia/label/cuda-12.4.0::cuda-toolkit
conda install -y nvidia::cuda-cudart-dev

# Install QServe
git clone git@github.com:mit-han-lab/qserve.git
cd qserve
pip install -e .
pip install ninja packaging
pip install flash-attn==2.4.1 --no-build-isolation
cd kernels
python setup.py install

# Install FlashInfer
pip install flashinfer -i https://flashinfer.ai/whl/cu121/torch2.3/
pip install tensor_parallel

📌Лицензирование: MIT License.

🟡Arxiv

🖥GitHub

#AI #ML #LLM #DuoAttention #LongContext

+1
0
+1
1
+1
0
+1
0
+1
0

Ответить

Ваш адрес email не будет опубликован. Обязательные поля помечены *