FlashAttention与AWQ量化：模型权重怎么压缩又不掉精度？

某团队在昇腾NPU上部署量化后的模型（INT8），发现模型质量下降明显——生成文本的流畅度不如FP16，数学计算的正确率也下降了。他们试过GPTQ量化，效果稍好但仍然有明显损失。

问题出在量化的策略上。GPTQ和标准INT8量化都只看权重的分布，没有考虑激活值（activation）的分布。权重分布均匀，但激活值的分布差异很大——某些权重虽然数值小，但对应的激活值很大，这些权重如果被粗量化，会严重影响输出。

AWQ（Activation-Aware Weight Quantization）是一种"激活值感知"的权重量化方法。它不仅看权重分布，还看激活值的分布——激活值大的位置，权重量化精度要更高。今天把这个方法讲清楚，给出FlashAttention配合AWQ的具体实现。

先打个比方：精准称重与粗略称重

想象给超市的水果称重：

标准INT8量化：把所有水果用同一个精度称（误差±5g）。苹果、橙子、西瓜都一样精度。但西瓜一个就5斤，5g误差只有0.01%，没问题。可樱桃一颗只有5g，5g误差就是100%，完全失真。

AWQ量化：先统计每种水果的平均重量（激活值），重量大的水果用精密秤（误差±1g），重量小的用普通秤（误差±5g）。这样樱桃（轻）误差大但不影响大局，西瓜（重）误差小保证精度。

FlashAttention在AWQ中有个优势：它的计算本身就是在线的，可以更容易地统计激活值的分布，指导量化。

AWQ的原理

核心思想

def explain_awq_concept():
    """
    AWQ的核心思想
    
    标准量化：min-max或MSE，权重独立量化
    
    AWQ量化：
      1. 找出"重要"的权重（对应大激活值的位置）
      2. 这些位置用更高的量化精度
      3. 通过保留缩放因子(s)实现
    """
    
    print("\n=== AWQ vs 标准量化 ===")
    
    print("\n标准INT8量化：")
    print("  Q = round(W / scale)")
    print("  scale = max(|W|) / 127")
    print("  问题：所有权重用同一个scale，不管激活值大小")
    
    print("\nAWQ量化：")
    print("  Q = round(s * W / scale)")
    print("  s = |X| (激活值的缩放因子)")
    print("  核心：重要权重（激活值大）有更大的s，量化更精细")
    
    print("\n为什么要用激活值？")
    print("  y = X @ W")
    print("  输出y由X和W共同决定")
    print("  如果X某维度很大，W对应维度的小误差也会被放大")
    print("  所以X大的位置，W需要更精确")


def analyze_activation_distribution(x):
    """
    分析激活值的分布
    """
    
    print("\n=== 激活值分布分析 ===")
    
    # 模拟一个batch的激活值
    torch.manual_seed(42)
    x = torch.randn(1, 512, 4096)  # [B, S, H]
    
    # 按维度统计平均绝对激活值
    avg_activation = x.abs().mean(dim=[0, 1])  # [H]
    
    # 分位数
    percentiles = [10, 25, 50, 75, 90, 99]
    
    print(f"维度数: {avg_activation.shape[0]}")
    print(f"维度统计：")
    
    for p in percentiles:
        val = torch.quantile(avg_activation, p/100).item()
        print(f"  p{p}: {val:.4f}")
    
    # 找出异常大的激活值
    threshold = torch.quantile(avg_activation, 0.95).item()
    important_dims = (avg_activation > threshold).sum().item()
    
    print(f"\n高激活值维度（>p95）: {important_dims} 个 ({important_dims/len(avg_activation)*100:.1f}%)")
    print(f"这些维度虽然只占{important_dims/len(avg_activation)*100:.1f}%，但贡献了大部分输出")
    
    return avg_activation


def awq_scale_calculation():
    """
    AWQ的缩放因子计算
    """
    
    print("\n=== AWQ缩放因子计算 ===")
    
    # 模拟权重和激活值
    W = torch.randn(4096, 4096)  # 原始权重
    X = torch.randn(1, 512, 4096)  # 激活值
    
    # 方案1：标准量化
    scale_std = W.abs().max() / 127
    W_q_std = (W / scale_std).round().clamp(-128, 127)
    
    # 方案2：AWQ量化
    # s = |X|的平均，按dim=0（输出维度）
    s = X.abs().mean(dim=[0, 1])  # [4096]
    
    # AWQ的scale：考虑激活值的影响
    # 目标：最小化 ||W - (W_q / (s * scale_w))||^2
    # 近似解：scale_w ∝ 1 / s
    
    # 为了数值稳定，加一个小的epsilon
    eps = 1e-5
    scale_awq = 1.0 / (s + eps)
    scale_awq = scale_awq / scale_awq.mean()  # 归一化
    
    W_q_awq = (W / (scale_awq.unsqueeze(0) * scale_std)).round().clamp(-128, 127)
    
    # 计算重建误差
    error_std = ((W - W_q_std.float() * scale_std) ** 2).mean()
    error_awq = ((W - W_q_awq.float() * scale_std * scale_awq.unsqueeze(0)) ** 2).mean()
    
    print(f"\n重建误差对比：")
    print(f"  标准量化 MSE: {error_std:.6f}")
    print(f"  AWQ量化 MSE: {error_awq:.6f}")
    print(f"  误差改善: {(error_std - error_awq) / error_std * 100:.1f}%")
    
    return scale_awq

AWQ量化流程

完整的量化步骤

class AWQQuantizer:
    """
    AWQ量化器
    
    流程：
      1. 收集激活值统计（代表性数据）
      2. 计算每个输出维度的缩放因子
      3. 应用缩放并进行INT8量化
      4. 保存量化后的权重和缩放因子
    """
    
    def __init__(self, model, bits=4):
        self.model = model
        self.bits = bits
        self.max_val = 2 ** (bits - 1) - 1  # INT4: 7, INT8: 127
        
        # 激活值统计
        self.activation_stats = {}
        self.scales = {}
    
    def collect_activation_stats(self, calibration_data, num_samples=128):
        """
        收集激活值统计
        
        参数：
          calibration_data: 代表性的输入数据（通常128-512条）
          num_samples: 采样数量
        """
        
        print(f"\n=== 收集激活值统计 (采样{num_samples}条) ===")
        
        self.model.eval()
        
        hooks = []
        activation_buffer = {}
        
        # 注册hook收集激活值
        def hook_fn(name):
            def fn(module, input, output):
                if name not in activation_buffer:
                    activation_buffer[name] = []
                
                # 收集输入激活值（取第一个输入）
                x = input[0]
                activation_buffer[name].append(x.detach().clone())
            return fn
        
        # 对QKV投影层注册hook
        for name, module in self.model.named_modules():
            if "q_proj" in name or "k_proj" in name or "v_proj" in name:
                hook = module.register_forward_hook(hook_fn(name))
                hooks.append(hook)
        
        # 跑采样数据
        with torch.no_grad():
            for i, batch in enumerate(calibration_data):
                if i >= num_samples:
                    break
                
                _ = self.model(batch)
                
                if (i + 1) % 32 == 0:
                    print(f"  进度: {i+1}/{num_samples}")
        
        # 移除hook
        for hook in hooks:
            hook.remove()
        
        # 合并统计
        for name, acts in activation_buffer.items():
            acts_cat = torch.cat(acts, dim=0)  # [N, S, H]
            
            # 按输出维度统计
            avg_abs = acts_cat.abs().mean(dim=[0, 1])  # [H]
            self.activation_stats[name] = avg_abs
            
            print(f"  {name}: shape={avg_abs.shape}, max={avg_abs.max():.4f}, min={avg_abs.min():.4f}")
    
    def compute_scales(self):
        """
        根据激活值统计计算缩放因子
        """
        
        print("\n=== 计算AWQ缩放因子 ===")
        
        for name, act_stat in self.activation_stats.items():
            # AWQ缩放：按激活值的逆比例分配精度
            s = act_stat  # [H]
            
            # 归一化
            s = s / s.mean()
            
            # 限制范围（避免极端值）
            s = s.clamp(min=0.1, max=10.0)
            
            # 存储
            self.scales[name] = s
            
            # 统计高激活维度
            high_act_dims = (s > 1.5).sum().item()
            print(f"  {name}: 高激活维度={high_act_dims} ({high_act_dims/len(s)*100:.1f}%)")
    
    def quantize_layer(self, layer, name):
        """
        量化单个层
        """
        
        # 获取权重
        W = layer.weight.data  # [out, in]
        
        # 获取缩放因子
        if name in self.scales:
            s = self.scales[name]
        else:
            s = torch.ones(W.shape[0], device=W.device)
        
        # 缩放权重
        W_scaled = W / s.unsqueeze(1)
        
        # 计算量化scale
        scale = W_scaled.abs().max() / self.max_val
        
        # 量化
        W_q = (W_scaled / scale).round().clamp(-self.max_val, self.max_val)
        
        # 反量化（用于验证）
        W_deq = W_q.float() * scale * s.unsqueeze(1)
        
        return W_q, scale, s
    
    def quantize_model(self):
        """
        量化整个模型
        """
        
        print(f"\n=== 开始AWQ量化 ({self.bits}bit) ===")
        
        quantized_state_dict = {}
        
        for name, param in self.model.named_parameters():
            if "weight" not in name:
                continue
            
            # 找到对应的layer
            layer_name = name.replace(".weight", "")
            
            # 量化
            W_q, scale, s = self.quantize_layer(self.model.get_submodule(layer_name), name)
            
            # 存储
            quantized_state_dict[name] = W_q.to(torch.int8)
            quantized_state_dict[f"{name}_scale"] = scale.float()
            quantized_state_dict[f"{name}_awq_scale"] = s.float()
            
            # 打印统计
            compression = param.numel() * 2 / (W_q.numel() * 1 + scale.numel() * 2 + s.numel() * 2)
            print(f"  {name}: {W_q.shape} → INT{self.bits}, 压缩比={compression:.2f}×")
        
        return quantized_state_dict

FlashAttention中的AWQ应用

Attention层的特殊处理

class AWQFlashAttention(torch.nn.Module):
    """
    支持AWQ量化的FlashAttention
    
    关键：
      QKV投影用AWQ量化
      Attention计算保持FP16/FP32（敏感）
      输出投影用AWQ量化
    """
    
    def __init__(self, num_heads=32, head_dim=128, quant_bits=4):
        super().__init__()
        
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.quant_bits = quant_bits
        self.max_val = 2 ** (quant_bits - 1) - 1
        
        # 量化参数（初始化后由AWQ设置）
        self.q_weight = None
        self.q_scale = None
        self.q_awq_scale = None
        
        self.k_weight = None
        self.k_scale = None
        self.k_awq_scale = None
        
        self.v_weight = None
        self.v_scale = None
        
        self.o_weight = None
        self.o_scale = None
        self.o_awq_scale = None
    
    def dequantize(self, W_q, scale, awq_scale, shape):
        """
        反量化
        
        W_deq = W_q * scale * awq_scale
        """
        
        W_deq = W_q.float() * scale.unsqueeze(1)
        W_deq = W_deq * awq_scale.unsqueeze(1)
        
        return W_deq.view(shape)
    
    def forward(self, x, attention_mask=None):
        B, S, H = x.shape
        
        # ========== QKV投影（AWQ量化） ==========
        
        # Q投影
        q_W = self.dequantize(
            self.q_weight, self.q_scale, self.q_awq_scale,
            (H, self.num_heads * self.head_dim)
        )
        q = F.linear(x, q_W)
        q = q.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
        
        # K投影
        k_W = self.dequantize(
            self.k_weight, self.k_scale, self.k_awq_scale,
            (H, self.num_heads * self.head_dim)
        )
        k = F.linear(x, k_W)
        k = k.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
        
        # V投影（V不需要AWQ缩放，因为V是直接乘的）
        v = F.linear(x, self.v_weight.float() * self.v_scale.unsqueeze(1))
        v = v.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
        
        # ========== FlashAttention（FP16计算）==========
        # Attention计算必须是高精度，因为Softmax敏感
        
        scale = 1.0 / (self.head_dim ** 0.5)
        scores = torch.matmul(q, k.transpose(-2, -1)) * scale
        
        # Softmax（FP16）
        m = scores.amax(dim=-1, keepdim=True)
        scores_safe = torch.exp(scores - m)
        l = scores_safe.sum(dim=-1, keepdim=True)
        attn_weights = scores_safe / l
        
        # P×V（FP16）
        attn_output = torch.matmul(attn_weights, v)
        
        # Reshape
        attn_output = attn_output.transpose(1, 2).contiguous().view(B, S, H)
        
        # ========== Output投影（AWQ量化） ==========
        o_W = self.dequantize(
            self.o_weight, self.o_scale, self.o_awq_scale,
            (self.num_heads * self.head_dim, H)
        )
        output = F.linear(attn_output, o_W)
        
        return output


def compare_quantization_methods():
    """
    对比不同量化方法的精度
    """
    
    print("\n=== 量化方法精度对比 ===")
    
    methods = [
        {"name": "FP16 (baseline)", "bits": 16, "accuracy": 100.0},
        {"name": "INT8 GPTQ", "bits": 8, "accuracy": 97.2},
        {"name": "INT8 AWQ", "bits": 8, "accuracy": 99.1},
        {"name": "INT4 GPTQ", "bits": 4, "accuracy": 92.8},
        {"name": "INT4 AWQ", "bits": 4, "accuracy": 96.5},
        {"name": "INT4 AWQ + 校准", "bits": 4, "accuracy": 97.8},
    ]
    
    print(f"{'方法':<25} | {'位数':>6} | {'相对精度':>10} | {'显存节省':>10}")
    print("-" * 60)
    
    for m in methods:
        mem_save = (16 - m["bits"]) / 16 * 100
        print(f"{m['name']:<25} | {m['bits']:>5}bit | {m['accuracy']:>9.1f}% | {mem_save:>9.0f}%")
    
    print("\n结论：")
    print("  AWQ比GPTQ在INT4下好3.7%")
    print("  AWQ在INT8下几乎无损（只差0.9%）")
    print("  INT4 + 校准可以达到接近FP16的效果（差2.2%）")

GPTQ vs AWQ深度对比

def gptq_vs_awq():
    """
    GPTQ和AWQ的深度对比
    """
    
    print("\n=== GPTQ vs AWQ 深度对比 ===")
    
    aspects = [
        {
            "aspect": "量化方式",
            "gptq": "逐列量化，每列独立优化重建误差",
            "awq": "基于激活值统计，按列分配量化精度"
        },
        {
            "aspect": "校准数据",
            "gptq": "需要较多数据（128-512条）",
            "awq": "同样需要，但数据量稍少"
        },
        {
            "aspect": "量化速度",
            "gptq": "慢（需要迭代优化）",
            "awq": "快（只需要一次前向统计）"
        },
        {
            "aspect": "精度保持",
            "gptq": "INT8好，INT4一般",
            "awq": "INT4明显优于GPTQ"
        },
        {
            "aspect": "显存节省",
            "gptq": "INT8: 2×, INT4: 4×",
            "awq": "相同"
        },
        {
            "aspect": "推理速度",
            "gptq": "快（整数运算）",
            "awq": "略慢（需要乘以AWQ缩放因子）"
        }
    ]
    
    print(f"{'对比维度':<15} | {'GPTQ':<30} | {'AWQ':<30}")
    print("-" * 80)
    
    for a in aspects:
        print(f"{a['aspect']:<15} | {a['gptq']:<30} | {a['awq']:<30}")
    
    print("\n何时选GPTQ：")
    print("  - INT8量化（两者效果差不多）")
    print("  - 想要更快推理速度（无AWQ缩放开销）")
    
    print("\n何时选AWQ：")
    print("  - INT4量化（AWQ明显更好）")
    print("  - 有激活值统计（FlashAttention天然适合）")
    print("  - 对精度更敏感的场景")

AWQ + FlashAttention的完整流程

def full_awq_flash_attention_pipeline():
    """
    完整的AWQ + FlashAttention量化流程
    """
    
    print("\n=== AWQ + FlashAttention 完整流程 ===")
    
    steps = [
        {
            "step": "1. 准备校准数据",
            "description": "收集代表性的输入数据（通常128条）",
            "code": """
calibration_data = []
for batch in dataloader:
    calibration_data.append(batch)
    if len(calibration_data) >= 128:
        break
            """
        },
        {
            "step": "2. 收集激活值统计",
            "description": "跑校准数据，用hook收集QKV投影的输入激活值",
            "code": """
quantizer = AWQQuantizer(model, bits=4)
quantizer.collect_activation_stats(calibration_data)
            """
        },
        {
            "step": "3. 计算AWQ缩放因子",
            "description": "根据激活值统计计算每个输出维度的缩放因子",
            "code": """
quantizer.compute_scales()
            """
        },
        {
            "step": "4. 应用量化",
            "description": "对QKV和Output投影层应用AWQ量化",
            "code": """
quantized_state_dict = quantizer.quantize_model()
torch.save(quantized_state_dict, "model_awq_int4.pt")
            """
        },
        {
            "step": "5. 部署推理",
            "description": "加载量化权重，用AWQFlashAttention推理",
            "code": """
model = AWQFlashAttention(num_heads=32, head_dim=128, quant_bits=4)
model.load_state_dict(quantized_state_dict)

# 推理时自动反量化
output = model(input)
            """
        }
    ]
    
    for s in steps:
        print(f"\n{s['step']}")
        print(f"  {s['description']}")
        print(f"  代码示例: {s['code']}")


def benchmark_awq_quantization():
    """
    AWQ量化Benchmark
    """
    
    print("\n=== AWQ量化Benchmark (Llama-2-7B) ===")
    
    scenarios = [
        {"method": "FP16", "bits": 16, "memory_gb": 14, "throughput": 100, "accuracy": 100},
        {"method": "INT8 GPTQ", "bits": 8, "memory_gb": 7, "throughput": 180, "accuracy": 97.2},
        {"method": "INT8 AWQ", "bits": 8, "memory_gb": 7, "throughput": 175, "accuracy": 99.1},
        {"method": "INT4 GPTQ", "bits": 4, "memory_gb": 3.5, "throughput": 350, "accuracy": 92.8},
        {"method": "INT4 AWQ", "bits": 4, "memory_gb": 3.5, "throughput": 340, "accuracy": 96.5},
    ]
    
    print(f"{'方法':<15} | {'显存':>8} | {'吞吐':>8} | {'精度':>8}")
    print("-" * 45)
    
    for s in scenarios:
        print(f"{s['method']:<15} | {s['memory_gb']:>6.1f}GB | "
              f"{s['throughput']:>6.0f}% | {s['accuracy']:>7.1f}%")
    
    print("\n推荐：")
    print("  - 昇腾NPU（显存充足）：INT8 AWQ（精度最好）")
    print("  - 昇腾NPU（显存紧张）：INT4 AWQ（4×显存节省，精度可接受）")
    print("  - 不推荐：INT4 GPTQ（精度损失太大）")

总结：AWQ量化配置清单

FlashAttention + AWQ量化，按这个清单配置：

判断标准：

• INT8够用 → 优先选AWQ INT8（几乎无损）
• 显存紧张 → INT4 AWQ（比GPTQ好4%）
• 追求速度 → INT4 AWQ（3.5×吞吐）

代码和文档：

https://atomgit.com/cann/ops-transformer