模型量化与压缩：从 FP32 到 INT4 的完整指南（ML 全场景）

import numpy as np

def affine_quantize(x: np.ndarray, n_bits: int = 8):
    """将浮点数组量化为 n_bits 整数"""
    # 确定量化范围
    qmin = -(2 ** (n_bits - 1))      # INT8: -128
    qmax = 2 ** (n_bits - 1) - 1     # INT8: 127
    
    # 计算缩放参数
    x_min, x_max = x.min(), x.max()
    scale = (x_max - x_min) / (qmax - qmin)
    zero_point = round(-x_min / scale)
    
    # 量化
    q = np.clip(np.round(x / scale) + zero_point, qmin, qmax)
    return q.astype(np.int8), scale, zero_point

def dequantize(q: np.ndarray, scale: float, zero_point: int):
    """将量化值反量化回浮点"""
    return (q.astype(np.float32) - zero_point) * scale

# 示例：量化一个权重张量
weights = np.random.randn(64, 64).astype(np.float32)
q_weights, scale, zp = affine_quantize(weights, n_bits=8)
recovered = dequantize(q_weights, scale, zp)

# 精度损失
mse = np.mean((weights - recovered) ** 2)
print(f"量化 MSE: {mse:.6f}")  # 通常 < 0.001

import torch
import torch.quantization
from torch.utils.data import DataLoader

# 1. 加载预训练模型
model = load_pretrained_model()
model.eval()

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

# 3. 插入量化/反量化节点
torch.quantization.prepare(model, inplace=True)

# 4. 校准：在少量数据上运行推理
calib_loader = get_calibration_data(num_samples=500)
with torch.no_grad():
    for batch in calib_loader:
        model(batch)  # 收集激活值统计

# 5. 转换：将模型转换为量化格式
torch.quantization.convert(model, inplace=True)

# 6. 验证精度
accuracy = evaluate(model, test_loader)
print(f"PTQ INT8 精度: {accuracy:.4f}")

# 7. 保存量化模型
torch.jit.save(torch.jit.script(model), "model_int8.pt")

import torch
import torch.quantization

# 1. 加载预训练模型
model = load_pretrained_model()
model.train()

# 2. 设置 QAT 配置
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')

# 3. 准备 QAT：插入伪量化节点
torch.quantization.prepare_qat(model, inplace=True)

# 4. 量化感知微调（通常 5-10 轮就够了）
optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
for epoch in range(5):
    for batch, labels in train_loader:
        optimizer.zero_grad()
        output = model(batch)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()

# 5. 转换为推理模式并真正量化
model.eval()
torch.quantization.convert(model, inplace=True)

# 6. 验证
accuracy = evaluate(model, test_loader)
print(f"QAT INT8 精度: {accuracy:.4f}")

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
import torch

# 配置 INT4 量化
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",       # 归一化浮点 4 位
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,  # 双重量化，进一步压缩
)

# 加载 INT4 量化模型
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3-70B",
    quantization_config=bnb_config,
    device_map="auto",
)

# 模型已以 INT4 精度加载到 GPU
# 显存占用：FP16 需要 ~140GB，INT4 只需 ~35GB
print(f"模型显存: {model.get_memory_footprint() / 1e9:.1f} GB")

# 可以直接用于推理
output = model.generate(**inputs, max_new_tokens=100)

# 将 HuggingFace 模型转换为 GGUF 格式
python convert-hf-to-gguf.py meta-llama/Llama-3-8B \
    --outfile llama-3-8b-f16.gguf

# 量化到 INT4（Q4_K_M 是推荐的平衡方案）
./quantize llama-3-8b-f16.gguf llama-3-8b-q4_k_m.gguf Q4_K_M

# Q4_K_M vs Q4_0 vs Q5_K_M 对比：
# Q4_0: 最简单，精度最低
# Q4_K_M: 推荐方案，精度和体积平衡
# Q5_K_M: 精度更高，体积稍大

# 查看量化后大小
ls -lh llama-3-8b-*.gguf
# f16: ~16GB, q4_k_m: ~4.9GB (压缩 3.3x)

import torch
import torch.nn as nn
import torch.nn.functional as F

class DistillationLoss(nn.Module):
    """知识蒸馏的损失函数"""
    def __init__(self, temperature=3.0, alpha=0.7):
        super().__init__()
        self.temperature = temperature
        self.alpha = alpha  # 蒸馏损失的权重
    
    def forward(self, student_logits, teacher_logits, labels):
        # 任务损失（Student 对真实标签）
        hard_loss = F.cross_entropy(student_logits, labels)
        
        # 蒸馏损失（Student 模仿 Teacher 的 Soft Labels）
        T = self.temperature
        soft_student = F.log_softmax(student_logits / T, dim=1)
        soft_teacher = F.softmax(teacher_logits / T, dim=1)
        soft_loss = F.kl_div(soft_student, soft_teacher, reduction="batchmean") * (T * T)
        
        # 加权组合
        return self.alpha * soft_loss + (1 - self.alpha) * hard_loss

# 训练循环
teacher = load_teacher_model()
teacher.eval()
student = load_student_model()
student.train()
criterion = DistillationLoss(temperature=3.0, alpha=0.7)
optimizer = torch.optim.AdamW(student.parameters(), lr=2e-5)

for batch, labels in train_loader:
    optimizer.zero_grad()
    with torch.no_grad():
        teacher_logits = teacher(batch)
    student_logits = student(batch)
    loss = criterion(student_logits, teacher_logits, labels)
    loss.backward()
    optimizer.step()

import torch
import torch.nn.utils.prune as prune

# 加载预训练模型
model = load_pretrained_model()

# 对每个卷积层进行结构化剪枝
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Conv2d):
        # 按 L1 范数剪枝最不重要的通道
        prune.l1_unstructured(
            module, name="weight",
            amount=0.3,  # 剪掉 30% 的通道
            dim=[0]       # 按输出通道维度
        )
        
        # 将剪枝永久化（移除剪枝 mask）
        prune.remove(module, "weight")

# 验证剪枝后的精度
accuracy = evaluate(model, test_loader)
print(f"剪枝后精度: {accuracy:.4f}")

# 如果精度下降可接受，继续微调
# 否则减少 amount 参数重试

import torch
import torch.nn as nn

def analyze_quantization_error(fp32_model, int8_model, sample_input):
    """逐层比较 FP32 和 INT8 模型的输出差异"""
    fp32_outputs = {}
    int8_outputs = {}
    
    # 注册 hook 收集每层输出
    def make_hook(name, store_dict):
        def hook(module, input, output):
            store_dict[name] = output.detach()
        return hook
    
    for name, module in fp32_model.named_modules():
        if isinstance(module, (nn.Linear, nn.Conv2d)):
            module.register_forward_hook(make_hook(name, fp32_outputs))
    
    for name, module in int8_model.named_modules():
        if isinstance(module, (nn.Linear, nn.Conv2d)):
            module.register_forward_hook(make_hook(name, int8_outputs))
    
    # 运行推理
    with torch.no_grad():
        fp32_model(sample_input)
        int8_model(sample_input)
    
    # 计算每层的量化误差
    errors = {}
    for name in fp32_outputs:
        if name in int8_outputs:
            fp32_out = fp32_outputs[name]
            int8_out = int8_outputs[name].to(fp32_out.dtype)
            # 反量化后比较
            error = torch.mean((fp32_out - int8_out) ** 2).item()
            errors[name] = error
    
    # 按误差排序
    sorted_errors = sorted(errors.items(), key=lambda x: x[1], reverse=True)
    print("量化误差最大的层:")
    for name, error in sorted_errors[:5]:
        print(f"  {name}: MSE = {error:.6f}")
    
    return sorted_errors