🌟 SageAttention: метод квантования механизма внимания в архитектурах трансформеров.

Внимание - ключевой компонент трансформеров, но его квадратичная сложность вычислений становится проблемой при обработке длинных последовательностей. Квантование успешно применяется для ускорения линейных слоев, но оно мало изучено применительно к механизму внимания.

SageAttention - экспериментальный метод, который использует 8-битное квантование механизма внимания для ускорения вычислений и сохранения точности модели.

Метод не требует специального обучения и конвертации моделей в какой-либо формат, он применяется к существующим трансформеным моделям в режиме "plug-and-play".

Ключевые особенности метода:

🟢Для уменьшения ошибки квантования используется сглаживания матртицы К (среднее значение K вычитается по всем токенам);

🟢Квантование Q и K в INT8;
INT8 в четыре раза быстрее, чем в FP16, и в два раза быстрее, чем в FP8.

🟢Matmul PV выполняется с FP16-накопителем;
Умножение матриц в высокой разрядности позволяет ускорить вычисления без потери точности.

🟢Адаптивное квантование;
Для каждого слоя внимания выбирается наиболее быстрый вариант квантования.

SageAttention реализован с использованием Triton и оптимизирован для GPU RTX4090 и 3090. Метод превосходит FlashAttention2 и xformers по скорости примерно в 2,1 и 2,7 раза соответственно.

Тестирование на Llama2, CogvideoX, Unidiffuser и TIMM подтвердило сохранение метрик точности при использовании SageAttention.

⚠️ Использование SageAttention рекомендуется с версиями:

🟠python>=3.11;
🟠torch>=2.4.0;
🟠triton-nightly.

⚠️ SageAttention оптимизирован для RTX4090 и RTX3090. На других архитектурах GPU прирост производительности может быть незначительным.

▶️Пример использования:

# Install sageattention
pip install sageattention

# How to use
from sageattention import sageattn
attn_output = sageattn(q, k, v, is_causal=False, smooth_k=True)

# Plug-and-play example with Cogvideo
# add the following codes and run
from sageattention import sageattn
import torch.nn.functional as F

F.scaled_dot_product_attention = sageattn

# Specifically
cd example
python sageattn_cogvideo.py

📌Лицензирование: BSD-3-Clause license.


🟡Arxiv
🖥GitHub


@ai_machinelearning_big_data

#AI #ML #SageAttention #Transformers