今天小编分享的科学经验:清华开源混合精度推理系统MixQ,实现大模型近无损量化并提升推理吞吐,欢迎阅读。
一键部署 LLM 混合精度推理,端到端吞吐比 AWQ 最大提升 6 倍!
清华大学计算机系 PACMAN 实验室发布开源混合精度推理系统——MixQ。
MixQ 支持 8 比特和 4 比特混合精度推理,可实现近无损的量化部署并提升推理的吞吐。
△图 1 MixQ 吞吐与已有开源工作比较
MixQ 同时量化权重和激活,使用低精度张量核心(INT8/INT4 Tensor Core)实现推理加速;同时,MixQ 提取激活中少量的离群值,使用高精度张量核心(FP16 Tensor Core)保持推理准确性,通过系统优化掩盖高精度访存开销。
不仅保持推理的准确性,而且通过使用低精度算力有效提升吞吐,充分发挥硬體计算潜力(图 1)。
同时,研究团队提供了基于 VLLM 和 Tensorrt-LLM 的混合精度推理,用户可以方便地一键部署模型。
△图 2 使用 VLLM 一键部署 4 比特和 8 比特混合精度量化并推理
MixQ 已支持多个主流大模型 LLaMA3,Qwen2,Baichuan2,ChatGLM 等。据了解,目前 MixQ 开源技术已被清程极智等 AI 行业公司应用在实际产品中。
该工作同时于高性能计算领網域顶级国际会议 SC ’ 24 发表,第一作者清华大学博士后陈逸东、通讯作者为翟季冬教授。
研究背景:已有量化技术总结
量化的主要技术路线有两条,第一条是权重量化。
权重量化的理论加速比是 16/ 量化的比特数。例如,将模型压缩成为 4bit,那么理论加速比为 16/4=4 倍。
然而,当服务商面临大量的用户同时访问时,权重量化的系统吞吐会低于 FP16 的吞吐,其主要原因是权重量化计算过程中将低精度权重恢复成 FP16 然后计算,这导致权重量化并不使用低精度算力,当场景表现为 compute bound 的时候,性能较低。
△图 3 用户请求多权重量化吞吐低于 FP16
第二条技术路线是量化权重和激活,使用低精度的张量核心来提升系统的吞吐。
直接将激活量化为低比特可能会出现较大的精度损失。其原因在于激活矩阵中存在离群值(图 4)。
一个有效的方法是 SmoothQuant,主要思想是通过平滑激活矩阵来降低量化激活的误差。
△图 4 激活矩阵中存在离群值
混合精度量化则是一类全新的量化方法,该方案先做了一个矩阵分解,对绝大部分权重和激活用低比特存储,将离群值用 FP16 存储,分别做矩阵乘法。
△图 5 混合精度量化示意图
混合精度量化的一个优势就是可以实现近乎无损精度的量化。使用混合精度量化的 LlaMA 模型在 MMLU 20 个领網域上的数据集进行推理准确率测试表明,采用 8bit 混合精度量化后的准确率下降不到 0.1%:
△图 6 混合精度量化分类准确率
不过,此前已有的混合精度量化的系统的性能普遍不高,主要瓶颈在针对离群点进行查找、访存和计算的开销占比大。
以混合精度库 Bitsandbytes 为例,实测试表明,Bitsandbytes 在用户请求数量为 512 时仅有 1.08 倍的加速。
△图 7 Bitsandbytes 的在 LLaMA70B 上的 Kernel 性能测试
△图 8 Atomic operator 是混合精度推理系统的瓶颈之一
那么,如何优化对离群点的查找、访存和计算的开销呢?
MixQ 的解决方案
MixQ 的核心思想是基于离群点的局部性对混合精度的计算图做等价变换,使得变换后的混合精度的计算图可以避免离群点查找的额外开销;在此基础上,通过图层融合和设计高效的混合精度数据结构降低访存开销;最后通过 CUTLASS 生成高性能的混合精度算子,达到提升系统性能的效果。
MixQ 的设计基于以下的观察:
离群点的局部性。对 LLM 的激活矩阵分析发现,在不同的 decode 阶段的离群点的分布是有规律的。
如图 9,红色的点表示的是第一次出现的离群点,绿色的点表示的是重复出现的离群点,随着 decode 的进行,多数离群点出现在了固定的 channel。
△图 9 decode 阶段离群点的分布规律
因此,研究人员得到一个重要的结论:在大部分的 decode 阶段是不需要重复检测离群点的,也就是说我们可以避免检查离群点的开销。
剩下的问题是,如何知道哪些时候不需要重复检查离群点呢?这个答案就隐藏在量化系数中。
在量化的过程中需要对矩阵进行 amax 的操作。因此,通过 amax 得到的结果可以判断矩阵中是否存在离群点。如 amax 的值大于阈值,那矩阵中存在离群点。反之则不存在。
更重要的是,amax 操作可以和前一个操作融合。这样不仅以极低的代价检测离群点的存在,还通过对图层进行融合来降低量化的开销。
基于以上的分析,MixQ 的设计使用了三个关键技术:
一是对计算图的等价变换。
针对混合精度的计算逻辑进行了等价变换以后,通过计算激活矩阵的 amax 的值,避免了检测离群点的开销。
△图 10 优化混合精度的计算逻辑
二是设计混合精度数据结构。
MixQ 将离群点 " 拼接 " 成了一个新的矩阵。这一方法相较于 ATOM 采用的重排列(reorder)具有更低的开销。
△图 11 MixQ:order-reserved 数据结构
三是使用 CUTLASS 编写高性能的混合精度的算子,这一关键技术的实现依赖于 NVIDIA 提供的高性能矩阵乘法模板 CUTLASS 3.x。
MixQ 在寄存器中反量化低精度的计算结果并与高精度的结果进行相加。
△图 12 融合 dequantize、scale 和 add 操作
下面来看 MixQ 的实验结果,以 LLaMA 70B 为例。
在准确率表现方面,MixQ 的准确率和 Bitsandbytes 一致。
△图 13 MixQ 的推理精度
在性能表现方面,MixQ 8bit kernel 是 Bitsandbytes 的 1.9 倍。
MixQ 4bit Kernel 的性能达 724TFLOPs,是 FP16 的 3.13 倍。
△图 14 MixQ Kernel 性能
端到端测试下,MixQ 在 batch=512 相对 Bitsandbytes 和 AWQ 加速 1.78 和 6 倍。
△图 15 多 batch 测试;上:MIXQ 的推理输出(19.21it/s);下:FP16 的推理输出 (13.56it/s)
项目地址:
[ 1 ] https://github.com/Qcompiler/MixQ_Tensorrt_LLM
[ 2 ] https://github.com/Qcompiler/MIXQ
[ 3 ] https://github.com/Qcompiler/vllm-mixed-precision
— 完 —
投稿请发邮件到:
标题注明【投稿】,告诉我们:
你是谁,从哪来,投稿内容
附上论文 / 项目主页链接,以及联系方式哦
我们会(尽量)及时回复你
点这里关注我,记得标星哦~
一键三连「分享」、「点赞」和「在看」
科技前沿进展日日相见 ~
>