llm-awq
MIT 开发的 LLM 激活感知训练后量化方法(INT3/4),获 MLSys 2024 最佳论文,被 vLLM/TensorRT-LLM 等主流平台广泛采用
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
MIT 开发的 LLM 激活感知训练后量化方法(INT3/4),获 MLSys 2024 最佳论文,被 vLLM/TensorRT-LLM 等主流平台广泛采用
加载项目详情…
本应用为开源项目,仅供学习研究,请遵守其开源协议。
想象一下,你手里有一台普通的游戏显卡 RTX 4090,想跑一个 70B 参数的大语言模型——按照常理,这简直是痴人说梦。MIT 的一群研究者偏偏说"我可以",他们开发的 AWQ(Activation-aware Weight Quantization) 技术,就像给大模型做了一次精准的"减肥手术",让它能在一张消费级显卡上流畅运行。
AWQ 在 2024 年 MLSys 顶会上斩获最佳论文奖,被 Google Vertex AI、Amazon SageMaker、NVIDIA TensorRT-LLM、Intel Neural Compressor、vLLM 等主流平台纷纷采纳,成为了 LLM 量化领域的事实标准。
图1:AWQ 整体架构——激活感知权重量化流程
大语言模型(LLM)能力强,但"体重"也惊人。一个 70B 参数的模型,FP16 精度下需要约 140GB 显存,光是显存就足以让大多数人望而却步。传统量化方案(比如 INT8)压缩效果有限,而更激进的 INT4 量化又往往让模型"变傻"——输出质量断崖式下滑。
MIT Han Lab 的研究团队从一个问题出发:量化过程中,哪些权重最重要?
他们的发现出人意料:只需要保护约 1% 的"显著权重"(salient weights),就能在 INT4 量化下几乎无损模型的表达能力。这个发现奠定了 AWQ 的核心思想——**"激活感知"**意味着在量化搜索过程中,以激活分布的统计特性为依据,智能识别并保护最重要的权重。
图2:多模态模型量化效果——AWQ-INT4 与 FP16 几乎无差异
AWQ 的核心技术路径分为三步:
第一步:激活感知扫描(AWQ Search)。在少量校准数据(通常 100 条文本样本)上运行模型,收集各层激活分布的统计数据。系统据此计算每个权重通道的"显著性分数",识别出对模型输出影响最大的那些权重。
第二步:保护性缩放(Scaling)。对识别出的关键权重,将其缩放倍数调高;对非关键权重则允许更激进的量化。整个过程不需要反向传播,也不需要重新训练模型,因此 AWQ 是一种**训练后量化(Post-Training Quantization, PTQ)**方法,大幅降低了使用门槛。
第三步:高效推理。量化后的权重配合 AWQ 团队自研的 W4A16 CUDA 内核(在 awq/kernels/ 目录下基于 CUTLASS 实现),在推理时以 INT4 权重 + FP16 激活的混合精度模式运行,兼顾了显存节省与计算效率。
这套方案的核心优势在于:
如果 AWQ 是引擎,TinyChat 就是装好引擎的整车。
TinyChat 是基于 AWQ 构建的端侧推理框架,当前版本为 2.0。相比初代版本,TinyChat 2.0 在预填充(prefilling)速度上提升了 1.5-1.7 倍,已支持 LLaMA-3-8B、DeepSeek-R1-Distill 系列、VILA-1.5、NVILA 等主流模型。
图3:TinyChat 在 RTX 4090 上,LLaMA-3-8B 量化推理速度达到 FP16 的 2.7 倍
图4:TinyChat 在 Jetson Orin 边缘设备上,推理速度达到 FP16 的 2.9 倍
值得注意的是,Jetson Orin 是一块功耗仅 15-60W 的嵌入式 GPU,通常用于机器人、无人机等场景。能在这样的设备上跑通 8B 级别的大模型,本身就是一件了不起的事。
TinyChat 还支持 Gradio Web UI,提供了一个图形化交互界面,用户不需要写代码,在浏览器里就能和大模型对话。对于想快速体验量化效果的开发者,这个功能非常友好。
AWQ 的模型支持列表非常全面,涵盖了主流的开源大模型:
| 模型类别 | 代表模型 | AWQ 支持 |
|---|---|---|
| 纯语言模型 | LLaMA-1/2/3, OPT, DeepSeek-R1-Distill | INT3/4 |
| 代码模型 | CodeLlama, StarCoder | INT3/4 |
| 视觉语言模型 | VILA-1.5, NVILA, LLaVA | INT4 |
| 对话模型 | Vicuna, Qwen-2.5 | INT4 |
所有预计算的 AWQ 量化权重托管在 HuggingFace Model Zoo,用户下载后可以直接加载使用,无需自己跑量化流程。
AWQ 的代码结构清晰,主要分为三层:
TinyChat 则是独立部署模块,包含针对不同模型类型的推理脚本(demo.py/vila10_demo.py/nvila_demo.py 等)和流式生成器(stream_generators/)。
AWQ 并非完美无缺,实际使用中有几个值得关注的点:
1. CUDA 内核编译门槛:虽然核心 Python 代码依赖不多,但 awq/kernels/setup.py 需要从源码编译 CUDA 扩展。边缘设备(如 Jetson Orin)需要手动安装预编译的 PyTorch wheel,增加了部署复杂度。
2. Flash Attention 依赖:高速推理依赖 Flash Attention 2,但在某些环境(尤其是非标准 CUDA 版本)下安装过程繁琐,容易出现 ABI 不匹配问题。
3. 量化精度对校准数据敏感:AWQ 搜索结果的质量依赖校准数据的多样性。对于垂直领域应用(如医疗、金融),使用默认的通用文本校准集可能导致领域特定任务精度下降。
4. INT3 量化支持有限:相比 INT4,INT3 量化在更多模型上缺乏预计算权重,用户需要自己跑 AWQ search,成本较高。
AWQ 的出现,让"大模型民主化"从口号变成了现实。它解决了三个关键问题:
AWQ 获得 MLSys 2024 最佳论文,表明学术界对"高效部署大模型"这个方向的持续关注。随着多模态大模型(VLM)成为主流,像 AWQ 这样能同时支持 LLM 和 VLM 的量化工具,价值将进一步凸显。
一句话总结:AWQ 是一种激活感知的训练后量化方法,通过保护 1% 的关键权重实现 INT3/4 量化下近乎无损的精度,被广泛集成于主流 AI 推理平台,是当前大模型端侧部署的首选方案之一。