量化技术对比：INT4/INT8/FP8

x_int8 = round(x_fp / scale)
scale = max(|x_fp|) / 127

x_int8 = round((x_fp - zero_point) / scale)
scale = (max(x) - min(x)) / 255
zero_point = min(x)

def per_group_quantize(x, group_size=128):
    """分组量化"""
    original_shape = x.shape
    # 重排为 [out_features // group_size, group_size, in_features]
    x_reshaped = x.view(-1, group_size, original_shape[-1])
    
    # 每组计算独立的scale
    scales = x_reshaped.abs().max(dim=1, keepdim=True).values / 127
    
    # 量化
    x_int8 = (x_reshaped / scales).round().clamp(-128, 127).to(torch.int8)
    
    return x_int8, scales

def per_group_dequantize(x_int8, scales, original_shape):
    """分组反量化"""
    x_fp = x_int8.float() * scales
    return x_reshaped.view(original_shape)

def llm_int8_matmul(A, B, threshold=6.0):
    """LLM.int8() 矩阵乘法"""
    # 识别异常值
    outlier_mask = A.abs() > threshold
    normal_mask = ~outlier_mask
    
    # 正常部分：INT8量化计算
    A_normal = A[normal_mask]
    B_quantized, scales = quantize_int8(B)
    result_normal = dequantize(A_normal_int8 @ B_quantized, scales)
    
    # 异常部分：FP16精度计算
    A_outlier = A[outlier_mask]
    result_outlier = A_outlier @ B
    
    # 合并结果
    result = torch.zeros_like(A @ B)
    result[normal_mask] = result_normal
    result[outlier_mask] = result_outlier
    
    return result

Y = (X · diag(s)^{-1}) · (diag(s) · W)

def gptq_quantize_layer(layer, calibration_data):
    """GPTQ单层量化"""
    W = layer.weight.data  # [out, in]
    
    # 用校准数据计算Hessian矩阵
    H = compute_hessian(layer, calibration_data)  # [in, in]
    H += torch.eye(H.shape[0]) * 0.01  # 正则化
    
    # 逐列量化
    for i in range(W.shape[1]):
        # 量化第i列
        w = W[:, i]
        w_q = quantize(w)
        
        # 计算量化误差
        err = (w - w_q) / H[i, i]
        
        # 将误差分摊到未量化的列
        W[:, i+1:] -= err.unsqueeze(1) @ H[i, i+1:].unsqueeze(0)
        
        W[:, i] = w_q
    
    layer.weight.data = W

def awq_quantize(W, X_calib, ratio=0.1):
    """AWQ量化"""
    # 找出重要权重通道（激活值大的通道）
    activation_magnitude = X_calib.abs().mean(dim=0)
    important_channels = activation_magnitude.topk(
        int(len(activation_magnitude) * ratio)
    ).indices
    
    # 对重要通道施加缩放
    scale = torch.ones(W.shape[1])
    scale[important_channels] = activation_magnitude[important_channels] / \
                                  activation_magnitude.mean()
    
    W_scaled = W * scale
    
    # 量化
    W_quantized = quantize_int4(W_scaled)
    
    # 反量化时除以scale
    W_dequantized = dequantize(W_quantized) / scale
    
    return W_dequantized

class FP8QuantizationAwareTraining:
    def __init__(self, model, calibration_data):
        self.model = model
        self.calibration_data = calibration_data
    
    def quantize_weights(self):
        """将权重量化到FP8（STE直通估计器）"""
        for param in self.model.parameters():
            param_fp8 = param.data.to(torch.float8_e4m3fn)
            # STE：前向用FP8，反向用FP16梯度
            param.data = param_fp8.to(param.dtype)

def apply_mixed_precision(model):
    """混合精度量化策略"""
    for name, module in model.named_modules():
        if 'embed' in name:
            # 嵌入层：保持FP16（对量化敏感）
            module.weight = module.weight.half()
        elif 'lm_head' in name:
            # 输出头：保持FP16
            module.weight = module.weight.half()
        elif 'attention' in name:
            # 注意力层：INT8（精度要求中等）
            quantize_module_int8(module)
        elif 'ffn' in name or 'mlp' in name:
            # FFN层：INT4（占参数主体，压缩收益最大）
            quantize_module_int4(module)
        else:
            # 其他层：INT8
            quantize_module_int8(module)