把感知换到无损时踩过的坑

from transformers import AutoModelForCausalLM, AutoTokenizer
from optimum.bettertransformer import BetterTransformer

model_path = "your-model-path"
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16)
tokenizer = AutoTokenizer.from_pretrained(model_path)

# 转换为 BetterTransformer 以优化推理
model = BetterTransformer.transform(model)

# 简单的 INT8 量化
from optimum.quanto import quantize, freeze

quantize(model, weights=torch.int8, activations=torch.int8)
freeze(model)

prompt = "解释一下什么是量化感知训练？"
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(**inputs, max_length=200)
print(tokenizer.decode(outputs[0]))

import torch
import torch.nn as nn
from transformers import Trainer, TrainingArguments

# 定义一个带量化感知的层
class QATLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        self.weight = nn.Parameter(torch.randn(out_features, in_features))
        self.bias = nn.Parameter(torch.zeros(out_features))
        self.scale = nn.Parameter(torch.ones(out_features, 1))
        self.zero_point = nn.Parameter(torch.zeros(out_features, 1))

    def forward(self, x):
        # 模拟量化操作
        weight_quant = torch.clamp(
            torch.round(self.weight / self.scale) + self.zero_point,
            -128, 127
        ) * self.scale - self.zero_point
        return nn.functional.linear(x, weight_quant, self.bias)

# 替换模型的线性层
def replace_with_qat(model):
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            qat_linear = QATLinear(module.in_features, module.out_features)
            qat_linear.weight.data = module.weight.data
            qat_linear.bias.data = module.bias.data
            parent = model
            for part in name.split('.')[:-1]:
                parent = getattr(parent, part)
            setattr(parent, name.split('.')[-1], qat_linear)

replace_with_qat(model)

# 微调模型
training_args = TrainingArguments(
    output_dir="./qat_finetuned",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=2,
    learning_rate=2e-5,
    num_train_epochs=3,
    save_steps=500,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=your_dataset,
)

trainer.train()

quantize(model, weights=torch.int8, activations=torch.int8)
freeze(model)

prompt = "解释一下什么是量化感知训练？"
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(**inputs, max_length=200)
print(tokenizer.decode(outputs[0]))

# 定义敏感层名单（这些层用更高精度）
sensitive_layers = [
    "model.layers.0.mlp.gate_proj",
    "model.layers.0.mlp.up_proj",
    # ... 其他敏感层
]

def mixed_precision_quantize(model):
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            if any(sens in name for sens in sensitive_layers):
                # 保留 FP16
                continue
            else:
                # 量化为 INT8
                quantize(module, weights=torch.int8, activations=torch.int8)

mixed_precision_quantize(model)

def per_channel_quantize(model):
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            # 为每个输出通道单独计算 scale 和 zero_point
            weight = module.weight.data  # [out_features, in_features]
            scale = weight.abs().max(dim=1, keepdim=True)[0] / 127
            zero_point = torch.zeros_like(scale)

            # 应用量化
            weight_quant = torch.clamp(
                torch.round(weight / scale) + zero_point,
                -128, 127
            ) * scale - zero_point

            module.weight.data = weight_quant

# 动态调整量化参数
def dynamic_quantization_params(tensor):
    # 计算张量的统计信息
    mean = tensor.mean()
    std = tensor.std()
    max_val = tensor.abs().max()

    # 根据分布调整 scale 和 zero_point
    if std > max_val / 2:
        # 分布较广，用更保守的量化
        scale = max_val / 127
        zero_point = torch.tensor(0)
    else:
        # 分布较窄，可以用更激进的量化
        scale = (max_val + 3 * std) / 127
        zero_point = -torch.tensor(mean / scale).round()

    return scale, zero_point

def safe_quantize(tensor, scale, zero_point):
    # 量化前 clip
    tensor_clipped = torch.clamp(tensor, -scale * 128, scale * 127)

    # 量化
    quantized = torch.round(tensor_clipped / scale) + zero_point

    # 反量化后再次 clip
    dequantized = (quantized - zero_point) * scale
    dequantized = torch.clamp(dequantized, -scale * 128, scale * 127)

    return dequantized