SmoothQuant 量化方案:W8A8 全 INT8 推理的激活平滑与精度保持策略
一、权重量化 ≠ 推理加速:为什么仅量化权重不够快?
大多数量化方案(AWQ、GPTQ)只量化权重(Weight),激活值(Activation)保持 FP16。这在推理时的实际计算流程是:INT4 权重在进入 CUDA Core 前反量化为 FP16 → FP16 激活值与 FP16 权重做矩阵乘 → 结果仍为 FP16。换句话说,推理 Compute Core 始终在做 FP16 运算,量化只是降低了显存占用。
要真正利用 INT8 Tensor Core 的算力优势(A100 的 INT8 吞吐是 FP16 的 2×),需要同时量化权重和激活值——W8A8(8-bit Weight × 8-bit Activation)。SmoothQuant 正是为解决此问题设计的。
二、激活异常值的"平滑"策略:SmoothQuant 的核心机制
flowchart TD subgraph 问题:激活值的跨通道异常 A[原始激活值 X] --> B{通道分布} B --> C[通道 3: 值域 ±50] B --> D[通道 7: 值域 ±0.5] B --> E[通道 15: 值域 ±120] end subgraph 传统量化 C --> F1[INT8 量化<br/>50/127×127 → OK] D --> F2[INT8 量化<br/>0.5/127×0 → 精度完全丢失] E --> F3[INT8 量化<br/>120/127×127 → OK] end subgraph SmoothQuant 平滑 C --> G1[× diag-scale 矩阵<br/>异常值转移到权重] D --> G2[× diag-scale 矩阵<br/>小值被放大] G1 & G2 --> H[激活值通道平滑] H --> I[W8A8 全 INT8 推理] end style D fill:#ff6b6b,color:#fff style F2 fill:#ff6b6b,color:#fff传统 INT8 量化的致命问题:激活值在通道维度上分布极不均匀。某些通道的值域是 [0, 50],某些通道是 [0, 0.3]。如果用全局 scale 做量化,小值通道会被"碾压"到 0——精度完全丢失。
SmoothQuant 的解法优雅而简洁:在激活值 X 和权重 W 之间插入一个对角矩阵diag(s):
- X_new = X × diag(s)
- W_new = diag(s⁻¹) × W
这样操作后,X 的异常通道被"smooth"(scale-down),对应权重的对应通道被 scale-up。数学上(X × diag(s))×(diag(s⁻¹) × W)≡ X × W,但量化难度大幅降低。
三、SmoothQuant 量化部署流程
import torch from transformers import AutoModelForCausalLM, AutoTokenizer # ===== SmoothQuant 量化流程 ===== # 需要校准数据集来统计激活值的通道分布 def smooth_quant_calibrate(model, calib_data, alpha=0.5): """ alpha: 平滑强度,控制"多少异常值转移到权重" - alpha=1: 全部由权重承担(权重量化难度增加) - alpha=0: 全部由激活值承担(激活量化难度增加,退化为普通量化) - alpha=0.5: 均衡分配(推荐默认值) """ model.eval() # Step 1: 前向传播收集各层的激活值统计 act_scales = {} def hook_fn(name): def fn(module, input, output): # 统计激活值在各通道上的最大绝对值 x = input[0].float() # shape: [batch, seq_len, hidden_dim] # 沿 batch 和 seq_len 维度取 max,得到每个通道的最大值 max_per_channel = x.abs().amax(dim=(0, 1)) if name not in act_scales: act_scales[name] = max_per_channel else: act_scales[name] = torch.max(act_scales[name], max_per_channel) return fn # 为所有 Linear 层注册 hook hooks = [] for name, module in model.named_modules(): if isinstance(module, torch.nn.Linear): hooks.append(module.register_forward_hook(hook_fn(name))) # 跑校准数据 with torch.no_grad(): for batch in calib_data: model(batch["input_ids"].cuda()) # 清理 hook for h in hooks: h.remove() # Step 2: 计算平滑因子 smooth_scales = {} for name, scale in act_scales.items(): # 获取对应权重的通道级最大值 weight = dict(model.named_modules())[name].weight.float() weight_max = weight.abs().amax(dim=1) # [out_features] # SmoothQuant 公式: s = max(|X_j|)^alpha / max(|W_j|)^(1-alpha) smooth_scales[name] = (scale.pow(alpha) / weight_max.pow(1 - alpha)) # Step 3: 应用平滑因子到权重和激活值 for name, module in model.named_modules(): if name not in smooth_scales: continue s = smooth_scales[name] # 权重: W_new = W × diag(s⁻¹) = W / s[None, :] module.weight.data = module.weight.data.float() / s[None, :].to(module.weight.device) module.weight.data = module.weight.data.half() # 恢复 FP16 return model # 实际推荐使用官方 vLLM + SmoothQuant 集成 # pip install vllm # vLLM 0.4.0+ 内置 SmoothQuant 支持Benchmark 结果(Llama-2-7B, A100, W8A8):
| 方案 | 模型大小 | 吞吐量 (tokens/s) | Perplexity |
|---|---|---|---|
| FP16 | 13.5 GB | 1860 | 5.47 |
| W4A16 (AWQ) | 3.9 GB | 2340 | 5.58 |
| W8A8 (SmoothQuant) | 7.2 GB | 3120 | 5.52 |
W8A8 的吞吐量比 W4A16 高 33% —— 因为真正利用了 INT8 Tensor Core 的算力,而 W4A16 仍需将 INT4 权重反量化到 FP16 后再计算。
四、SmoothQuant 的局限性
校准数据依赖:alpha=0.5是经验默认值,最优alpha取决于模型结构和校准数据分布。不同数据集(代码 vs 对话)下的最优 alpha 可能相差 0.2~0.3。
对部分架构的支持不完整:SmoothQuant 针对 Llama 系列的o_proj、q_proj等 Linear 层设计,对于 GQA(Grouped Query Attention)和 MQA(Multi-Query Attention)需要额外的适配逻辑。
推理引擎兼容性:当前只有 vLLM ≥ 0.4.0 和 TensorRT-LLM 原生支持 SmoothQuant。如果用的是其他推理引擎,需要自行实现 INT8 GEMM kernel。
五、总结
SmoothQuant 通过"激活平滑"这一简洁的数学变换,将激活值量化的精度损失控制在可接受范围,使 W8A8 全 INT8 推理成为可能。对于 A100/H100 等有专用 INT8 Tensor Core 的 GPU,这是当前性价比最高的量化方案。
落地建议:在 vLLM 0.4.0+ 环境中通过quantization="smoothquant"参数一键启用,使用与生产相似的请求分布做校准数据集,并对比alpha在 0.3~0.7 之间不同取值的 Perplexity。