01 - 模型量化基础概念¶

import torch
import torch.quantization

def post_training_quantization(model, dataloader):
    """
    训练后量化

    Args:
        model: 要量化的模型
        dataloader: 用于校准的数据加载器
    """
    # 设置模型为评估模式
    model.eval()

    # 配置量化
    model.qconfig = torch.quantization.get_default_qconfig('fbgemm')

    # 准备量化
    model_prepared = torch.quantization.prepare(model)

    # 使用校准数据进行校准
    print("开始校准...")
    with torch.no_grad():
        for batch_x, _ in dataloader:
            model_prepared(batch_x)

    # 转换为量化模型
    print("转换为量化模型...")
    quantized_model = torch.quantization.convert(model_prepared)

    return quantized_model

# 使用示例
# quantized_model = post_training_quantization(model, dataloader)

import torch
import torch.nn as nn
import torch.quantization

def quantization_aware_training(model, dataloader, epochs=5):
    """
    量化感知训练

    Args:
        model: 要训练的模型
        dataloader: 数据加载器
        epochs: 训练轮数
    """
    # 配置量化
    model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')

    # 准备量化感知训练
    model_prepared = torch.quantization.prepare_qat(model, inplace=True)

    # 训练
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model_prepared.parameters(), lr=0.001)

    model_prepared.train()
    for epoch in range(epochs):
        for batch_x, batch_y in dataloader:
            optimizer.zero_grad()
            outputs = model_prepared(batch_x)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()
        print(f"Epoch {epoch+1} completed")

    # 转换为量化模型
    quantized_model = torch.quantization.convert(model_prepared)

    return quantized_model

# 使用示例
# quantized_model = quantization_aware_training(model, dataloader, epochs=5)

# 使用GPTQ库
from transformers import AutoModelForCausalLM, AutoTokenizer
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig

def gptq_quantization(model_name, calibration_data):
    """
    GPTQ量化

    Args:
        model_name: 模型名称
        calibration_data: 校准数据
    """
    # 加载分词器
    tokenizer = AutoTokenizer.from_pretrained(model_name)

    # 配置量化
    quantize_config = BaseQuantizeConfig(
        bits=4,  # 量化位数
        group_size=128,  # 组大小
        damp_percent=0.01,  # 阻尼百分比
        desc_act=False,  # 激活描述
        sym=True,  # 对称量化
        true_sequential=True,  # 顺序量化
        model_name_or_path=model_name,
        model_file_base_name="model"
    )

    # 加载并量化模型
    model = AutoGPTQForCausalLM.from_pretrained(
        model_name,
        quantize_config=quantize_config,
        calibration_data=calibration_data
    )

    return model, tokenizer

# 使用示例
# model, tokenizer = gptq_quantization("meta-llama/Llama-2-7b-hf", calibration_data)

# 练习: 实现简单的线性量化
def simple_quantize(tensor, bits=8):
    # 你的代码
    pass

# 练习: 计算量化误差
def quantization_error(original, quantized, scale, zero_point):
    # 你的代码
    pass

# 练习: 实现训练后量化
class PostTrainingQuantization:
    def __init__(self, model):
        # 你的代码
        pass

    def quantize(self, calibration_data):
        # 你的代码
        pass

# 练习: 实现量化感知训练
class QuantizationAwareTraining:
    def __init__(self, model):
        # 你的代码
        pass

    def train(self, dataloader, epochs):
        # 你的代码
        pass

def simple_quantize(tensor, bits=8):
    qmin = -2**(bits-1) if bits > 1 else 0
    qmax = 2**(bits-1) - 1 if bits > 1 else 1
    scale = (tensor.max() - tensor.min()) / (qmax - qmin)
    zero_point = qmin - round(tensor.min() / scale)
    quantized = torch.clamp(torch.round(tensor / scale) + zero_point, qmin, qmax)
    return quantized, scale, zero_point

def quantization_error(original, quantized, scale, zero_point):
    dequantized = (quantized - zero_point) * scale
    mae = torch.mean(torch.abs(original - dequantized)).item()
    mse = torch.mean((original - dequantized)**2).item()
    return {'mae': mae, 'mse': mse}