专栏名称: PaperWeekly

PaperWeekly是一个分享知识和交流学问的学术组织，关注的领域是自然语言处理的各个方向。我们热爱知识，分享知识，希望通过我们大家的努力为自然语言处理的发展做出一点点贡献。我们每周会分享一期特定话题的论文笔记和本周值得读的相关论文。

又快又准、即插即用！清华提出8比特量化Attention，两倍加速不掉点

PaperWeekly · 公众号 · 科研 · 2024-10-21 13:40

正文

©作者 | 张金涛

单位 | 清华大学

大模型中，线性层的低比特量化（例如 INT8, INT4）已经逐步落地；对于注意力模块，目前几乎各个模型都还在用高精度（例如 FP16 或 FP32）的注意力运算进行训练和推理。然而，随着大型模型需要处理的序列长度不断增加，Attention（注意力运算）的时间开销逐渐成为网络优化的主要瓶颈。

为了提高注意力运算的效率，清华大学陈键飞团队提出了 8Bit 的 Attention（SageAttention）。实现了 2 倍以及 2.7 倍相比于 FlashAttention2 和 xformers 的即插即用的推理加速，且在视频、图像、文本生成等大模型上均没有端到端的精度损失。

论文标题：

SageAttention: Accurate 8-Bit Attention for Plug-and-play Inference Acceleration

文章链接：

https://arxiv.org/abs/2410.02367

代码链接：

https://github.com/thu-ml/SageAttention

即插即用举例

SageAttention 可以一行代码轻松替换掉 torch 中当前最优的 Attention 接口（scaled_dot_product_attention），实现即插即用的推理加速。

具体来说，SageAttention 的使用非常方便，使用 pip install sageattention 后，只需要在模型的推理脚本前加入以下三行代码即可：

效果上，以开源视频生成模型 CogvideoX 为例，使用 SageAttention 可以端到端加速 35%，且生成的视频无损：

▲ 全精度 Attention

▲ SageAttention

接下来，将从背景与挑战，技术方案，以及实验效果介绍 SageAttention。

背景

随着大模型需要处理的序列长度越来越长（比如 Llama3.1 支持 128K 的序列长度），Attention 的速度优化变得越来越重要。下图展示了一个标准的 Transformer 模型中各运算随着序列长度变化的时间占比：

挑战

为了方便指代注意力元算中包含的矩阵，我们先回顾一下注意力的计算公式：

将神经网络中各运算的数值类型从高比特量化至低比特是一种有效提升计算和访存效率的方法。然而，研究团队发现直接将注意力运算中的 Q, K, P, V 从 FP16 量化为 INT8 或者 FP8 后将会导致在几乎所有模型和任务上都会得到极差的结果，例如，在 Unidiffuser 文生图模型中，会得到一张完全模糊的图像；在 Llama2-7B 进行四选一选择题任务上得到 25.5% 的准确率。

经过仔细分析后，研究团队发现主要是两个原因导致了量化注意力的不准确：

1. 大多视频、图像生成模型中，矩阵 K 表现出了极强的通道维度的异常值分布，直接使用 INT8 或者 FP8 数据类型对其进行量化会导致巨大的误差。

2. 在所有模型中，对矩阵 P, V 进行量化不能保证一个模型中所有层的精度。下表展示了对 P, V 量化后，Llama2-7B 和 Unidiffuser 模型所有层中，最差情况的层对应的量化注意力的准确度（该准确度为量化注意力相比全精度注意力的误差），可以发现不管对 P, V 矩阵进行何种 8Bit（INT8，E4M3，E5M2）量化，总有些层的准确率非常差，导致了端到端效果的下降。