清华开源MixQ:实现大模型近无损量化并提升推理吞吐
清华开源MixQ:实现大模型近无损量化并提升推理吞吐
清华大学计算机系PACMAN实验室发布开源混合精度推理系统MixQ,支持8比特和4比特混合精度推理,可实现近无损的量化部署并提升推理的吞吐。MixQ已支持多个主流大模型,包括LLaMA3、Qwen2、Baichuan2和ChatGLM等。
MixQ的技术优势
MixQ的核心优势在于其独特的混合精度推理机制。通过同时量化权重和激活,并使用低精度张量核心(INT8/INT4 Tensor Core)实现推理加速,MixQ能够在保持高推理准确性的同时显著提升系统吞吐量。
图1 MixQ吞吐与已有开源工作比较
混合精度量化的技术原理
混合精度量化的主要思想是将权重和激活矩阵分解为低精度部分和高精度部分。低精度部分使用INT8或INT4存储,而高精度部分(即离群值)则使用FP16存储。这种混合存储方式既保持了高精度计算的准确性,又充分利用了低精度算力的加速优势。
图5 混合精度量化示意图
MixQ的创新优化
MixQ在混合精度量化的基础上进行了多项创新优化:
离群点局部性优化:通过分析发现,不同decode阶段的离群点分布具有规律性。MixQ利用这一特性,避免了重复检测离群点的额外开销。
混合精度数据结构设计:MixQ将离群点“拼接”成新的矩阵,相比传统的重排列方法具有更低的开销。
高性能算子优化:使用NVIDIA提供的CUTLASS 3.x模板,MixQ在寄存器中反量化低精度计算结果并与高精度结果相加,进一步提升了计算效率。
性能表现
实验结果显示,MixQ在LLaMA 70B模型上的表现尤为突出:
- 在8bit模式下,MixQ的性能是Bitsandbytes的1.9倍。
- 在4bit模式下,MixQ的性能达到724TFLOPs,是FP16的3.13倍。
- 在端到端测试中,MixQ在batch=512时相对Bitsandbytes和AWQ分别加速1.78倍和6倍。
图11 MixQ:order-reserved数据结构
应用前景
MixQ的出现为大模型的量化部署提供了新的解决方案。其近无损的量化效果和显著的性能提升,使其在实际应用中具有广阔前景。目前,MixQ已支持多个主流大模型,并在清程极智等AI公司得到实际应用。
这项研究成果已在高性能计算领域顶级国际会议SC’24发表,第一作者为清华大学博士后陈逸东,通讯作者为翟季冬教授。