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