AWQ量化中的激活感知为什么不是所有权重都应该被同等量化一、量化误差不是均匀分布的——激活异常值导致的重要通道效应传统训练后量化方法如RTN、GPTQ的隐含假设是所有权重对输出误差的贡献均等因此量化策略对所有通道一视同仁。但2023年的多项研究LLM.int8(), SmoothQuant, AWQ揭示了一个关键事实在大语言模型中某些通道的激活值比均值大100倍以上。这些激活异常通道上即使微小的权重量化误差也会被放大100倍后传导到输出。AWQActivation-aware Weight Quantization的核心洞察是与其均匀地保护所有权重不如识别出那些与激活异常通道相关的权重并给予额外的精度预算。这可以通过对权重矩阵进行逐通道的缩放来实现——在量化前将异常通道的权重放大使它们在量化网格上占据更多有效位在推理时通过对应的反向缩放来补偿。flowchart TD A[权重矩阵 W] -- B[统计激活分布] B -- C[识别异常通道: 激活幅值 均值] C -- D[逐通道缩放] D -- D1[异常通道: scale 1 → 权重放大] D -- D2[正常通道: scale ≈ 1 → 不变] D1 -- E[等效变换: W W × diag(s)] D2 -- E E -- F[统一量化 W → INT4] F -- G[推理时: 输入 × diag(1/s)] G -- H[结果: 异常通道获得更多有效Bit] style C fill:#fff3e0 style D fill:#e8f5e9 style H fill:#e1f5fe二、AWQ的等效变换原理缩放不改变计算结果但改变量化误差AWQ的操作在数学上是等价的——通过引入缩放因子 s 并分别在权重和输入侧做变换前向传播结果不变原始计算Y X × WAWQ变换Y (X × diag(1/s)) × (diag(s) × W) X × W但在量化视角下diag(s) × W中异常通道的权重被放大量化网格的相对分辨率更高有效比特数更大。缩放因子 s 的选择基于激活幅值s_i ∝ ||X_i||^α其中 α 是需要搜索的超参数。import torch import torch.nn as nn from typing import Tuple, Optional def awq_scale_weights( W: torch.Tensor, activation_stats: torch.Tensor, # [in_features] 每个输入通道的平均激活幅值 alpha: float 0.5, n_grid: int 20 ) - Tuple[torch.Tensor, torch.Tensor]: AWQ 权重缩放的核心实现。 alpha 参数控制缩放力度 - alpha0: 不缩放退化为标准均匀量化 - alpha1: 缩放比例 激活幅值异常通道被过度保护 - alpha0.5: 论文推荐的折中值 为什么需要搜索 alpha 而非固定值 不同模型和不同层的激活异常程度不同。 在OPT系列中 alpha≈0.3 效果最好 在LLaMA系列中 alpha≈0.5-0.7 更好。 这表明激活异常的模式依赖模型架构。 in_features W.shape[1] device W.device # 步骤1: 计算每个输入通道的缩放因子 # s mean(|X|)^alpha然后在网格上搜索最优 s act_magnitude activation_stats.float() # 搜索最优 alpha best_alpha alpha best_error float(inf) for candidate_alpha in torch.linspace(0.1, 1.0, n_grid): s torch.pow(act_magnitude, candidate_alpha) # 归一化保证整体数值尺度不变 s s / s.max() s s.to(device) # 应用缩放 W_scaled W.float() * s.unsqueeze(0) # 模拟量化-反量化 # 对称INT4量化[-7, 7] max_val W_scaled.abs().max(dim1, keepdimTrue).values scale_q max_val / 7.0 W_quant torch.clamp( torch.round(W_scaled / scale_q), -7, 7 ) W_dequant W_quant * scale_q # 恢复原始尺度 W_dequant W_dequant / s.unsqueeze(0) # 计算量化误差 error (W.float() - W_dequant).abs().mean().item() if error best_error: best_error error best_alpha candidate_alpha.item() # 使用最优 alpha 重算 s torch.pow(act_magnitude, best_alpha) s s / s.max() s s.to(device) return s, best_alpha def apply_awq_to_linear( linear: nn.Linear, calibration_data: torch.Tensor, # [N, in_features] 代表性校准数据 alpha: float 0.5, n_grid: int 20 ) - nn.Linear: 将 AWQ 应用到 PyTorch 的 nn.Linear 层。 修改后的层在推理时自动处理输入侧的缩放补偿。 with torch.no_grad(): # 步骤1: 统计激活分布 # 使用一批校准数据前向传播收集输入激活 # 在实际应用中校准数据应来自训练集的代表性子集 activation_stats calibration_data.abs().mean(dim0) # [in_features] # 步骤2: 计算最优缩放因子 s, best_alpha awq_scale_weights( linear.weight.data, activation_stats, alphaalpha, n_gridn_grid ) # 步骤3: 变换权重 # W W * diag(s) scaled_weight linear.weight.data.float() * s.unsqueeze(0) # 步骤4: 量化 max_val scaled_weight.abs().max(dim1, keepdimTrue).values scale_q max_val / 7.0 quantized_weight torch.clamp( torch.round(scaled_weight / scale_q), -7, 7 ).to(linear.weight.dtype) linear.weight.data quantized_weight # 步骤5: 存储量化参数和缩放因子作为 buffer linear.register_buffer(awq_scale_q, scale_q.to(linear.weight.dtype)) linear.register_buffer(awq_s, s.to(linear.weight.dtype)) # 修改 forward 方法在输入侧除以 s original_forward linear.forward def awq_forward(x): # 输入缩放补偿 x_scaled x / linear.awq_s # 反量化 计算 weight_dequant linear.weight * linear.awq_scale_q return nn.functional.linear(x_scaled, weight_dequant, linear.bias) linear.forward awq_forward return linear三、AWQ vs GPTQ两种误差处理策略GPTQ 和 AWQ 都是训练后权重量化方法但它们的误差处理策略截然不同GPTQ逐列贪心量化 误差在未量化列上补偿。优点是不需要校准激活分布缺点是量化顺序固定不能并行。AWQ通道级缩放 统一量化。优点是可以并行量化所有列速度更快缺点是需要校准数据来统计激活分布且缩放因子的搜索增加了预处理时间。在LLaMA-7B上的对比实验显示两者在WikiText-2困惑度上差距在0.05以内AWQ略好但AWQ的量化速度约快5-10倍得益于并行量化。四、AWQ的适用前提和失效场景需要校准数据AWQ 依赖校准数据来统计激活分布。如果下游任务的数据分布与校准数据差异很大激活异常通道的识别可能不准确导致精度退化。激活异常值依赖模型架构AWQ 的效果受模型架构影响。在 LLaMA 系列上效果很好因为 SwiGLU 激活函数导致明显的通道间激活差异但在某些架构如使用 GeLU 且具有较强正则化的模型上激活差异可能不够显著AWQ的缩放收益有限。混合精度推理的复杂性AWQ 的缩放-量化-反缩放流程需要在线推理时执行额外的逐通道乘除。虽然开销远小于反量化本身但对极致低延迟场景P99 1ms可能有可测量的影响。五、总结AWQ 的核心贡献是识别并利用了激活异常值的不均匀分布不是所有权重都应该被同等量化——与激活异常通道关联的权重需要更多精度预算。通过数学等效变换缩放-量化-反缩放在保持前向传播等价的前提下优化量化精度。相比 GPTQAWQ 通过并行量化实现了更快的量化速度。方法依赖校准数据在跨域部署时需要重新统计激活分布。