【大模型】- 量化

import torch
import numpy as np

def linear_quantize(tensor, num_bits=8):
    """线性量化：将浮点数映射到整数范围"""
    qmin = 0
    qmax = 2**num_bits - 1
    
    # 计算缩放因子和零点
    min_val = tensor.min()
    max_val = tensor.max()
    
    scale = (max_val - min_val) / (qmax - qmin)
    zero_point = qmin - min_val / scale
    
    # 量化
    quantized = torch.round(tensor / scale + zero_point)
    quantized = torch.clamp(quantized, qmin, qmax).to(torch.int8)
    
    return quantized, scale, zero_point

def linear_dequantize(quantized, scale, zero_point):
    """线性反量化"""
    return (quantized.float() - zero_point) * scale

def symmetric_quantize(tensor, num_bits=8):
    """对称量化：零点为 0"""
    qmax = 2**(num_bits - 1) - 1
    
    max_abs = tensor.abs().max()
    scale = max_abs / qmax
    
    quantized = torch.round(tensor / scale)
    quantized = torch.clamp(quantized, -qmax - 1, qmax).to(torch.int8)
    
    return quantized, scale

def block_quantize(tensor, block_size=128, num_bits=8):
    """分块量化：每块独立量化"""
    original_shape = tensor.shape
    tensor = tensor.view(-1, block_size)
    
    # 每块独立量化
    quantized_blocks = []
    scales = []
    
    for block in tensor:
        q, s = symmetric_quantize(block, num_bits)
        quantized_blocks.append(q)
        scales.append(s)
    
    return quantized_blocks, scales, original_shape

from transformers import GPTQConfig, AutoModelForCausalLM

# 使用 GPTQ 量化
quantization_config = GPTQConfig(
    bits=4,
    dataset="c4",
    group_size=128,
    desc_act=True,
    damp_percent=0.01,
)

model = AutoModelForCausalLM.from_pretrained(
    "model_name",
    quantization_config=quantization_config,
    device_map="auto"
)

class GPTQQuantizer:
    def __init__(self, model, num_bits=4, group_size=128):
        self.model = model
        self.num_bits = num_bits
        self.group_size = group_size
    
    def quantize(self, calibration_data):
        """使用校准数据进行 GPTQ 量化"""
        # 收集激活统计
        activations = self.collect_activations(calibration_data)
        
        # 逐层量化
        for name, module in self.model.named_modules():
            if isinstance(module, nn.Linear):
                self.quantize_layer(name, module, activations[name])
        
        return self.model
    
    def quantize_layer(self, name, module, activations):
        """量化单个层"""
        weight = module.weight.data
        
        # 计算 Hessian 矩阵
        H = activations.T @ activations
        
        # 逐列量化
        for col in range(weight.shape[1]):
            # 计算量化误差
            w = weight[:, col]
            q, scale = symmetric_quantize(w, self.num_bits)
            
            # 更新剩余列以补偿误差
            error = w - q * scale
            if col + 1 < weight.shape[1]:
                weight[:, col+1:] += error.unsqueeze(1) * H[col, col+1:] / H[col, col]
        
        # 存储量化权重
        module.weight = nn.Parameter(q)
        module.scale = scale

from transformers import AwqConfig

awq_config = AwqConfig(
    bits=4,
    fuse_max_seq_len=2048,
    do_fuse=True,
)

model = AutoModelForCausalLM.from_pretrained(
    "model_name",
    quantization_config=awq_config,
    device_map="auto"
)

def evaluate_quantization(original_model, quantized_model, eval_dataset):
    """评估量化对模型质量的影响"""
    # 计算困惑度
    original_ppl = compute_perplexity(original_model, eval_dataset)
    quantized_ppl = compute_perplexity(quantized_model, eval_dataset)
    
    # 计算困惑度增加
    ppl_increase = (quantized_ppl - original_ppl) / original_ppl
    
    print(f"原始困惑度: {original_ppl:.4f}")
    print(f"量化困惑度: {quantized_ppl:.4f}")
    print(f"困惑度增加: {ppl_increase:.2%}")
    
    return ppl_increase < 0.1  # 如果增加小于 10%，则可接受

def compute_perplexity(model, dataset):
    """计算模型困惑度"""
    total_loss = 0
    total_tokens = 0
    
    model.eval()
    with torch.no_grad():
        for batch in dataset:
            input_ids = batch["input_ids"].to(model.device)
            attention_mask = batch["attention_mask"].to(model.device)
            
            outputs = model(input_ids, attention_mask=attention_mask, labels=input_ids)
            total_loss += outputs.loss.item() * input_ids.shape[1]
            total_tokens += input_ids.shape[1]
    
    avg_loss = total_loss / total_tokens
    perplexity = torch.exp(torch.tensor(avg_loss))
    
    return perplexity.item()

def select_quantization_precision(model_size_gb, target_size_gb, quality_requirement):
    """根据模型大小和目标选择量化精度"""
    if quality_requirement == "high":
        if target_size_gb >= model_size_gb / 2:
            return "fp16"
        else:
            return "int8"
    elif quality_requirement == "medium":
        if target_size_gb >= model_size_gb / 4:
            return "int8"
        else:
            return "int4"
    else:  # low
        if target_size_gb >= model_size_gb / 8:
            return "int4"
        else:
            return "int3"

def mixed_precision_quantization(model, sensitivity_threshold=0.1):
    """混合精度量化：敏感层使用高精度"""
    quantized_model = {}
    
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            # 评估层的敏感性
            sensitivity = evaluate_layer_sensitivity(module)
            
            if sensitivity > sensitivity_threshold:
                # 敏感层使用 FP16
                quantized_model[name] = module.half()
            else:
                # 不敏感层使用 INT4
                quantized_model[name] = quantize_to_int4(module)
    
    return quantized_model

# 使用 bitsandbytes 量化
from transformers import BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model = AutoModelForCausalLM.from_pretrained(
    "model_name",
    quantization_config=bnb_config,
    device_map="auto"
)

# 推理
inputs = tokenizer("Hello, world!", return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=50)