模型压缩：量化、剪枝、蒸馏与部署优化实战

# PyTorch 官方 PTQ 流程
import torch
import torch.quantization as tq

# 准备：校准模型（统计激活范围）
model.eval()
model.qconfig = tq.get_default_qconfig("fbgemm")   # 8-bit 对称量化
tq.prepare(model, inplace=True)

# 校准：喂入代表性样本，统计 min/max
with torch.no_grad():
    for batch in calib_loader:
        model(batch)

# 转换：实际量化
tq.convert(model, inplace=True)

# 推理时使用量化算子（fbgemm/x86 后端加速）
output = model(input_tensor)

# llama.cpp 使用 GGUF 量化（实际部署中最常用）
# 不同量化档位：Q4_K_M / Q5_K_M / Q8_0 等
# Q4_K_M：4bit 综合，体积约 4.5GB（7B 模型）
# Q8_0：8bit，体积约 7GB，精度接近 FP16
from llama_cpp import Llama

llm = Llama(
    model_path="models/Qwen2.5-7B-Instruct-Q4_K_M.gguf",
    n_ctx=8192, n_gpu_layers=-1,   # 全部层上 GPU
    verbose=False,
)
output = llm("中国的首都是", max_tokens=64)

# 用 llama.cpp 将模型转成 4bit GGUF
python convert_hf_to_gguf.py \
  --outfile models/Qwen2.5-7B-Instruct-Q4_K_M.gguf \
  --outtype q4_k_m \
  models/Qwen2.5-7B-Instruct/

# 幅值剪枝：置零最小幅度的权重
import torch

def magnitude_prune(weight: torch.Tensor, sparsity: float = 0.5):
    """按幅度剪枝：保留最大 sparsity 比例的权重。"""
    mask = torch.zeros_like(weight, dtype=torch.bool)
    k = int(weight.numel() * (1 - sparsity))
    flat = weight.view(-1).abs()
    _, top_idx = torch.topk(flat, k)
    mask.view(-1)[top_idx] = True
    return weight * mask          # 被 mask 为 False 的位置置零

weight = torch.randn(256, 256)
pruned = magnitude_prune(weight, sparsity=0.5)
print(f"非零比例: {(pruned != 0).float().mean():.3f}")   # ~0.5

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

def distillation_loss(
    student_logits, teacher_logits, labels, T=4.0, alpha=0.5
):
    """Hinton 蒸馏损失：软标签 KL + 硬标签 CE。"""
    soft_targets = F.softmax(teacher_logits / T, dim=-1)
    soft_pred = F.log_softmax(student_logits / T, dim=-1)
    kl_loss = F.kl_div(soft_pred, soft_targets, reduction="batchmean") * (T * T)
    ce_loss = F.cross_entropy(student_logits, labels)
    return alpha * kl_loss + (1 - alpha) * ce_loss

import torch

def low_rank_factorize(weight: torch.Tensor, r: int) -> tuple[torch.Tensor, torch.Tensor]:
    """SVD 低秩分解：返回两个小矩阵用于替换原权重。"""
    U, S, Vh = torch.linalg.svd(weight, full_matrices=False)
    # 只保留前 r 个奇异值
    U_r = U[:, :r]
    S_r = S[:r]
    Vh_r = Vh[:r, :]
    approx = (U_r * S_r) @ Vh_r
    rank_ratio = r * (weight.shape[0] + weight.shape[1]) / weight.numel()
    print(f"参数压缩比: {rank_ratio:.2%}")
    return U_r * S_r, Vh_r     # 两个低秩因子

# PyTorch → ONNX 导出（含动态维度）
import torch

class TinyBERT(nn.Module):
    def forward(self, input_ids, attention_mask):
        return self.bert(input_ids, attention_mask)[0]

model.eval()
dummy_ids = torch.randint(0, 30522, (1, 128))
dummy_mask = torch.ones(1, 128, dtype=torch.long)

torch.onnx.export(
    model,
    (dummy_ids, dummy_mask),
    "model.onnx",
    opset_version=17,
    input_names=["input_ids", "attention_mask"],
    output_names=["logits"],
    dynamic_axes={                       # 支持动态 batch 与序列长度
        "input_ids": {0: "batch", 1: "seq_len"},
        "attention_mask": {0: "batch", 1: "seq_len"},
    },
)

# ONNX Runtime 推理 + 启用 TensorRT EP
import onnxruntime as ort

providers = [("TensorrtExecutionProvider", {
    "trt_engine_cache_enable": True,    # 缓存 engine，避免重复构建
    "trt_fp16_enable": True,            # FP16 加速
}), "CUDAExecutionProvider"]

sess = ort.InferenceSession("model.onnx", providers=providers)
outputs = sess.run(["logits"], {
    "input_ids": input_ids.numpy(), "attention_mask": mask.numpy(),
})

# 示例：7B 模型不同精度的 KV cache 显存对比
# 序列长 4096、40 层、32 头、head_dim=128
# FP16 KV cache = 2 × 4096 × 40 × 32 × 128 × 2B = ~2.7GB
# INT8 KV cache = 一半，~1.3GB
# GQA（分组查询注意力）进一步把 K/V 头数减半，显存再降

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

def compute_perplexity(model, tokenizer, texts, stride=512):
    model.eval()
    total_nll, total_tokens = 0.0, 0
    for text in texts:
        encodings = tokenizer(text, return_tensors="pt")
        input_ids = encodings.input_ids
        with torch.no_grad():
            for i in range(0, input_ids.size(1) - 1, stride):
                inp = input_ids[:, i:i+stride]
                out = input_ids[:, i+1:i+stride+1]
                logits = model(inp).logits
                shift_logits = logits[:, :-1, :].reshape(-1, logits.size(-1))
                shift_labels = out[:, 1:].reshape(-1)
                nll = torch.nn.functional.cross_entropy(
                    shift_logits, shift_labels, reduction="sum")
                total_nll += nll.item()
                total_tokens += shift_labels.numel()
    return math.exp(total_nll / total_tokens)

# 对比 FP16 与 INT4 的 PPL，增幅 >5% 则回退档位

# vLLM 服务化部署 AWQ 量化模型（零代码改动启用量化）
from vllm import LLM, SamplingParams

llm = LLM(
    model="Qwen/Qwen2.5-7B-Instruct-AWQ",   # 已量化权重
    quantization="awq",                      # 指定量化方法
    tensor_parallel_size=1,
    gpu_memory_utilization=0.85,
    max_model_len=8192,
)
outputs = llm.generate(["量化部署如何保证精度？"], SamplingParams(temperature=0.7))
print(outputs[0].outputs[0].text)