跳转到内容
Jason's Study

AWQ 2023 — 让 70B 大模型住进 RTX 4090

是什么

AWQ（Activation-aware Weight Quantization）是 MIT HAN Lab 2023 年提出的训练后 4-bit 权重量化方法，目标是把超大模型塞进单卡甚至手机里跑。

日常类比：搬家时所有家具要塞进小货车，但有几件传家宝不能压。AWQ 不是按”哪件最大”挑，而是看”哪件配着最重的客人坐”——客人重才是要保护的椅子。

更精确说：AWQ 用输入激活的幅值判断哪 1% 权重通道是关键通道，对它们做数学等价的缩放，再统一量化成 INT4——全 INT4 存储、无混合精度、硬件直接读，所以又快又省。

为什么重要

  • 70B 模型从 140GB 降到 35GB：消费级 GPU（RTX 4090 24GB）能跑，移动端经过裁剪也能跑
  • 比 GPTQ 更快：GPTQ 要逐列反向重建 + Hessian 求逆，AWQ 只跑前向 grid search，单卡几分钟出结果
  • 跨域不掉精度：GPTQ 容易过拟合校准集，AWQ 因为不重建权重所以泛化更稳，指令微调 / 视觉模型都能用
  • TinyChat 推理框架：作者团队自配 INT4 GEMM kernel，端侧吞吐 1.5–3x 提升
  • vLLM / TensorRT-LLM / HuggingFace 默认 4-bit 路径

如果你想理解今天为什么 Llama-2 70B 能在一台笔记本上跑出 token，AWQ 是绕不开的环节。

核心要点

AWQ 的逻辑链可以拆成 三步：

  1. 不是所有权重都同样重要：随机量化全部权重会掉很多点；保留 0.1%–1% 的”显著权重”用 FP16，模型几乎不掉点。但 FP16+INT4 混合存储硬件解码慢——这是 LLM.int8() 的痛点。

  2. 显著看激活，不看权重：哪些权重显著？答案在输入激活里——激活值大的输入通道对应的整列权重就是要保护的。直觉是 y = W·x 里 x[j] 大就放大整列 W 的误差。

  3. 缩放替代混合精度：把显著列权重乘 s>1，对应激活除以 s——数学等价，但量化时这些权重占据更多 INT4 表达范围，相对误差被压缩。s = s_X^α，α∈[0,1] 网格搜索。最终全 INT4，硬件友好。

实践案例

案例 1：为什么”看激活而不是看权重”

设 y = W·x，输出 y_i = Σ_j W[i,j] · x[j]。

如果 x[j] 是 1000，那 W[i,j] 哪怕量化误差只有 0.01，乘上去也变成 10——直接放大输出误差。所以保护对象不是 W 自己最大的位置，而是 x 大的那一列对应的整列 W。

这是 AWQ 与 LLM.int8() 等”按权重大小选关键位置”方法的根本区别。

案例 2：缩放数学等价的把戏

原始：y = W · x

AWQ 缩放：

y = (W · diag(s)) · (diag(1/s) · x) = W · x  ✓

数学完全等价。但是 W·diag(s) 这个新矩阵的”显著列”被放大，量化后相对误差被压缩。激活的除法可以离线吸进前一层 LayerNorm 或线性层——运行时零开销。

案例 3：在 vLLM 里加载 AWQ 模型

from vllm import LLM
llm = LLM(model="TheBloke/Llama-2-70B-AWQ", quantization="awq")
out = llm.generate(["你好"])

vLLM 把 INT4 权重读进来，前向时反量化（dequant）成 FP16 再做矩阵乘。显存省 4x，吞吐升 1.5–3x——这是 70B 模型能跑在单 RTX 4090 上的根据。

案例 4：α 的搜索过程

对一层线性层，AWQ 内层流程：

  1. 跑校准数据，记录每通道激活均值 s_X[j]
  2. 对 α ∈ {0, 0.05, 0.1, ..., 1.0} 这 21 个值各试一次
  3. 用候选 s = s_X^α 做缩放，量化后反量化，算输出 MSE
  4. 取 MSE 最小那个 α 作为该层的最终值

整个过程纯前向，单层不到 1 秒，70B 模型几百层全量 grid search 单卡几分钟搞定。

案例 5：移动端部署的可行性

把 70B 模型量化到 INT4 后只占 35GB，但手机内存远不够。AWQ 的真正贡献在于：激活仍然 FP16 不需要重新设计算子——配合 Llama-2 7B / 13B 这种中等模型，INT4 后 4–7GB 内存就能跑。

TinyChat（作者配套框架）在 RTX 4090 / Apple M2 / Jetson Orin 三种硬件上都给出了完整 INT4 推理路径。这意味着开发者不用自己写 kernel，直接 pip install 就能在端侧得到 1.5–3x 吞吐提升。

踩过的坑

  1. 校准集不是越大越好：AWQ 论文实测 16–512 条样本就够；继续加反而引入噪声。要的是”激活分布形状”，不是覆盖所有领域。

  2. α 不是固定值：每层各算各的最优 α，不能全模型共享。

  3. “看激活” ≠ “激活也量化”：AWQ 是 W4A16——激活仍然 FP16。如果想 W4A8 / W8A8，请看 SmoothQuant，那是另一条路线。

  4. per-channel vs per-group：实践中 AWQ 用 group-wise 量化（每 128 个权重一组共享 scale），group=128 是常见 sweet spot。

  5. scale 不能太大：理论上 s 越大保护越强，但 s 极大会让其他列被压扁。α≤1 是必须的硬约束。

  6. 校准域偏移：通用文本校准、部署跑代码生成，激活分布偏移会让 α 失准。建议用部署域小样本重新校准。

适用 vs 不适用场景

适用：

  • 大模型推理部署，显存是瓶颈（7B/13B/70B）
  • 需要快速量化、没有训练资源
  • 跨领域 / 多模态模型（指令微调、VLM）
  • W4A16 路径下要稳精度
  • 端侧部署（笔记本、手机）

不适用：

  • 训练阶段（AWQ 是 PTQ-only）
  • 需要激活也量化（请看 SmoothQuant / LLM.int8()）
  • 极端 2-bit 场景——AWQ 对 INT3/INT4 设计，2-bit 误差爆炸
  • 模型本身就小（<1B）——量化收益有限

历史小故事（可跳过）

  • 2022.08：LLM.int8() 发现激活有”异常通道”，提出按列分两路计算（FP16 + INT8）
  • 2022.10：GPTQ 用 OBS（最优脑外科）思路，Hessian 反向重建权重——精度好但只能拟合校准集
  • 2023.06：AWQ arXiv 上线，提出”激活决定显著、缩放替代混合精度”两步走
  • 2023 末：vLLM / TensorRT-LLM 接入 AWQ 作为默认 4-bit 路径
  • 2024.05：MLSys 2024 Best Paper，正式工业落地

之后 AWQ 与 GPTQ 形成事实上的双标杆，前者偏前向 + 速度，后者偏反向重建 + 校准精度。

学到什么

  1. 重要性不在自己身上，在它和别人的关系里——这个思想超越量化（搜索里的 idf、推荐里的协同过滤本质都是相对量）
  2. 数学等价的缩放是个万能工具：训练里有 batchnorm 折叠，推理里有 AWQ scaling
  3. PTQ 三角：精度、速度、复杂度——AWQ 选了”前向校准 + 简单 grid search”的甜点
  4. 硬件友好优先于学术最优：混合精度数学上能赢，但工程上输给了”全 INT4 + 缩放”
  5. per-channel scaling 是工程最简形式——比 GPTQ 的 Hessian 反向重建简单一个量级，精度相当甚至更好
  6. 激活统计 = 语义先验：前向跑校准本身就是用模型自己告诉你”哪里重要”——比启发式都准

延伸阅读

关联

  • vllm —— 推理引擎，AWQ 是它的默认 4-bit 路径
  • paged-attention —— vLLM 显存管理，与权重压缩互补
  • gptq-2023 —— 同期对手，反向重建 vs 前向 grid search 的两条路线
  • llm-int8-2022 —— 异常通道路由的早期混合精度方案
  • product-quantization-2011 —— 古典 PQ 思想：用查表代替全精度乘法
  • ampere-architecture-2020 —— INT4 / INT8 张量核心的硬件支撑