计算所王颖研究员团队联合华北电力大学与上海科技大学团队,在计算机体系结构领域顶级会议ASPLOS 2025上发布并开源了名为COMET的创新量化推理框架。该框架通过系统与算法的紧密协同,在权重(W4)、激活(A4)以及KV缓存(KV4)三个关键维度实现全4比特量化,打破了传统量化性能与精度的瓶颈。在LLaMA-70B等大型模型上,仅带来了0.32的困惑度细微增加,同时实现了端到端推理速度提升2.02倍,树立了大模型量化推理领域的新标杆。
技术亮点包括:
1. 全面量化压缩:将权重、激活和KV缓存全部压缩至4比特,极大程度降低存储与计算资源占用,且精度损失控制在极低水平。
2. 充分挖掘GPU潜能:通过创新的通道重排技术和异步流水线设计,显著提高硬件利用率,GPU算力利用率达76%。
3. 简易集成开源方案:用户仅需5行代码即可将COMET框架集成至TensorRT-LLM,实现对LLaMA-3及Qwen等大模型全系列的支持。
4. 卓越实际效果:在70亿参数大型模型的长序列推理场景中,最高加速可达3.27倍,显著降低端侧设备推理运行成本约70%。
当前,全球大型模型的能耗问题日益严峻,例如GPT-4单次训练耗电接近数千户家庭的年用电量,推理过程的能效瓶颈更是不容忽视。传统量化技术面临明显不足:
• 权重量化(W4A16)虽压缩权重,但激活仍采用高精度,造成GPU算力严重浪费;
• 权重和激活双量化(W8A8)虽然保障了精度,但加速效果受限且在长序列推理中内存需求仍居高不下。
COMET框架通过硬件与算法的深度融合,从GPU架构底层重新设计量化策略和计算核心,彻底解决了权重、激活及KV缓存量化难题,实现了全维度4比特超低精度推理的新突破,为大规模模型的节能高效推理提供了坚实技术支撑。
1. 离群点分布规律的重新认识
经过对LLaMA等主流大模型的深入研究,团队首次发现激活中的异常值并非随机分布,而是集中出现在大约10%的特定通道。这一发现为细粒度混合精度量化算法(FMPQ)的开发提供了重要依据。
在实际应用中,团队设计了一种兼顾模型精度与计算效率的创新方案:首先,通过将激活张量以128通道为单位进行分块划分,使其与现代GPU张量核心的64×64×32计算结构高度匹配,确保量化过程与硬件执行紧密结合。针对激活中非均匀分布的异常值,算法引入动态精度调节机制,赋予异常值密集区域8比特高精度表示,而其他部分则采用4比特量化,从而在保持模型性能的同时,大幅提升计算效率。
此外,FMPQ引入了通道重排技术,通过智能置换将分散的异常值通道聚集到同一计算块内,进一步降低对高精度计算资源的需求,实现更优的性能表现。
2. 动态调度混合精度计算
在实际应用中,结合W4A4和W4A8的混合精度运算常常导致计算负载分布不均,成为性能提升的瓶颈。为解决这一问题,团队设计了三层异步流水线结构,旨在充分发挥GPU的计算能力,提高整体资源利用率。
整个系统架构从数据传输与计算的重叠优化入手,采用双缓冲技术,在数据预加载的同时启动计算过程,有效隐藏了读取延迟,保证张量核心持续高效运行。团队创新性地打破了传统将一个计算任务固定分配给单个流式多处理器(SM)的模式,设计了更灵活的SM协同机制,使得闲置的计算单元能够动态接管邻近的任务,大幅度降低了局部计算资源的浪费。
此外,COMET框架在任务调度方面也实现了细致的优化。系统根据实时计算负载,动态调整INT4和INT8任务的分配比例,从而将不同SM之间的负载差异控制在5%以内,几乎实现了均衡的算力分配。这些设计共同提升了整体的计算效率和资源利用率。
3. 极致挖掘硬件指令集潜能
在4比特量化领域,真正的难点不在于量化算法本身,而是数据格式转换的效率限制。传统做法通常需要执行多达10条指令才能完成一次4比特到8比特的转换,极大地影响了部署性能。针对这一瓶颈,研究团队从硬件指令集入手,深度挖掘GPU的底层能力,通过重构数据存储格式和指令执行流程,将转换操作精简至仅需两条指令完成。
这一简化方案依赖两大关键创新:首先,团队充分利用GPU的PTX指令特性,实现了同时从16位数据块中提取4个4比特数值,显著提升了解包速度;其次,采用了“零扩展”替代常用的“符号扩展”技术,结合缩放参数的补偿机制,保证了整体精度稳定不受影响。凭借此设计,转换效率较传统方法提升了约5倍。
在4比特极限压缩的背景下,COMET框架通过细粒度混合精度量化,有效降低了量化对模型精度的影响。实际测试数据显示,LLaMA-7B模型在WikiText2数据集上的困惑度仅由5.68略微上升至5.95,误差保持在5%以内,展现出较强的稳定性。更为突出的是,针对LLaMA-3-70B等大型模型,困惑度的提升被严格控制在0.36以内,精度表现较传统方案提升了三倍以上。同时,COMET还创新性地将KV缓存压缩至4比特,显著减少了长序列任务中的内存需求,降低幅度达75%,且精度损失控制在极低的0.05左右,基本可忽略不计。整体来看,COMET的精度已接近FP16浮点基准,明显优于常见的W4A4及W8A8量化方法。
在性能方面,COMET同样实现了显著突破。以A100 GPU为例,经过深度优化的COMET-W4Ax计算内核在适配LLaMA、Mistral、Qwen等主流模型时,其运行速度是cuBLAS内核的2.75到2.97倍。更为关键的是,这个性能提升不仅体现在单次计算内核中,也得以在端到端推理流程中保持。在处理长序列任务(如1024个token输入)时,整体吞吐量提升达到了3.27倍;即使在较短序列(128个token)条件下,依然能实现1.63倍的加速。得益于异步流水线和动态调度技术的应用,GPU中流式多处理器(SM)的利用效率由原先的45%提升到了76%,有效挖掘了硬件的计算潜能。
COMET-W4Ax的内核和接口现已公开开源,开发者只需简单几行代码便能将其无缝集成到TensorRT-LLM框架中,实现从模型压缩到推理加速的全流程支持,极大推动大规模模型的普及与广泛应用。
相关论文《COMET: Towards Practical W4A4KV4 LLMs Serving》已正式发布,可通过链接 https://dl.acm.org/doi/pdf/10.1145/3676641.3716252 获取全文。论文中详细阐述了研究工作,并附带完整的开源代码,欢迎大家下载阅读并发表意见交流。
·
欢迎关注“福大大架构师每日一题”,让AI助力您的未来发展。
·