深度学习模型轻量化实战：AWQ量化技术入门指南

1次阅读

共计 2900 个字符，预计需要花费 8 分钟才能阅读完成。

在将深度学习模型（尤其是大语言模型如 LLaMA-7B）部署到边缘设备时，我们常常遇到两个主要问题：

内存瓶颈：一个 7B 参数的 FP16 模型需要占用约 14GB 内存，远超大多数移动设备的承载能力
推理延迟：高精度浮点计算在边缘设备上执行缓慢，难以满足实时性要求

传统解决方案如剪枝 (Pruning) 和知识蒸馏 (Knowledge Distillation) 往往需要重新训练模型，而量化 (Quantization) 技术可以直接对预训练模型进行压缩——其中 AWQ(Activation-aware Weight Quantization)因其独特的激活感知特性成为当前研究热点。

方法	精度损失	计算开销	是否需要校准数据	适用场景
PTQ(Post-Training Quantization)	中	低	否	快速部署
GPTQ(GPT Quantization)	小	高	是	高精度需求
AWQ	极小	中	是	边缘设备部署

AWQ 的核心优势在于：
– 通过分析激活值 (Activation) 分布自动识别重要权重
– 对异常值 (Outlier) 进行特殊处理，减少量化误差传播

AWQ 的核心公式描述权重缩放策略：

s^* = \arg\min_s \|Wx - \hat{W}(s)x\|^2

其中：
– $W$ 是原始权重
– $\hat{W}$ 是量化后权重
– $s$ 是逐通道 (per-channel) 的缩放因子

对于包含异常值的特征通道，AWQ 采用非均匀量化策略：
1. 计算激活值的百分位数 (如 99.9%) 作为截断阈值
2. 对超出阈值的权重保留更高精度

def preprocess_weights(weight, act_scale, quant_bit=4):
    """
    weight: 原始权重 [out_features, in_features]
    act_scale: 激活值缩放系数 [in_features]
    """
    # 计算逐通道缩放因子
    channel_scale = torch.clamp(weight.abs().max(dim=1)[0] / act_scale, min=1e-5)

    # 应用缩放
    scaled_weight = weight / channel_scale.unsqueeze(1)

    # 执行量化
    max_val = scaled_weight.abs().max()
    quant_step = max_val / (2**(quant_bit-1)-1)
    quantized = torch.clamp(torch.round(scaled_weight/quant_step), 
                           -2**(quant_bit-1), 2**(quant_bit-1)-1)

    return quantized, channel_scale, quant_step

def calibrate_model(model, calib_loader, num_samples=512):
    """收集激活统计量"""
    act_stats = {}

    with torch.no_grad():
        for i, (inputs, _) in enumerate(calib_loader):
            if i * inputs.size(0) >= num_samples:
                break

            outputs = model(inputs)

            # 记录每层激活的 99.9% 百分位数
            for name, module in model.named_modules():
                if isinstance(module, nn.Linear):
                    act = outputs[name+'.act']  # 假设已注册 forward hook
                    if name not in act_stats:
                        act_stats[name] = []
                    act_stats[name].append(act.abs().quantile(0.999))

    # 计算每层的缩放因子            
    act_scales = {k: torch.stack(v).mean() for k,v in act_stats.items()}
    return act_scales

def build_awq_engine(onnx_path, quant_cfg):
    builder = trt.Builder(TRT_LOGGER)
    network = builder.create_network()
    parser = trt.OnnxParser(network, TRT_LOGGER)

    # 加载 ONNX 模型
    with open(onnx_path, "rb") as f:
        parser.parse(f.read())

    # 配置 AWQ 量化参数
    config = builder.create_builder_config()
    config.set_flag(trt.BuilderFlag.FP16)
    config.set_flag(trt.BuilderFlag.INT8)
    config.quantization_flags = trt.QuantizationFlag.CALIBRATE_BEFORE_FUSION

    # 设置逐层量化参数
    for layer in network:
        if layer.type == trt.LayerType.FULLY_CONNECTED:
            layer.precision = trt.int8
            layer.set_output_type(0, trt.int8)
            layer.set_input_quantization(0, quant_cfg[layer.name])

    return builder.build_engine(network, config)