AI模型压缩与推理优化实战指南：从剪枝到蒸馏

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

# 1. 局部剪枝（单层）
model = resnet18()
module = model.conv1
prune.l1_unstructured(module, name='weight', amount=0.3)  # 剪掉 30%

# 2. 全局剪枝（跨层）
parameters_to_prune = [
    (model.conv1, 'weight'),
    (model.layer1[0].conv1, 'weight'),
    (model.layer1[0].conv2, 'weight'),
]

prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.2
)

# 3. 通道剪枝（更实用）
def channel_prune(model, amount=0.3):
    """按通道重要性剪枝"""
    for name, module in model.named_modules():
        if isinstance(module, nn.Conv2d):
            # 计算 BN 层的 gamma 作为重要性
            bn = find_corresponding_bn(model, name)
            importance = bn.weight.data.abs()
            threshold = torch.quantile(importance, amount)

            mask = importance > threshold
            module.weight.data = module.weight.data[mask]
            module.out_channels = mask.sum().item()

# 1. 训练原始模型
model = train_full_model()

# 2. 剪枝
prune_model(model, amount=0.3)

# 3. 微调恢复精度
finetune_model(model, epochs=10)

# 4. 重复（渐进式剪枝）
for iteration in range(5):
    prune_model(model, amount=0.1)
    finetune_model(model, epochs=5)

import torch
from torch.quantization import quantize_dynamic, quantize_fx

# 1. 动态量化（最简单，只量化权重）
quantized_model = quantize_dynamic(
    model,
    {nn.Linear},
    dtype=torch.qint8
)

# 2. 静态量化（权重 + 激活，需要校准数据）
from torch.quantization import prepare, convert

model.eval()
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')

# 插入观察者
prepared_model = prepare(model)

# 用校准数据跑前向
for batch in calibration_data:
    prepared_model(batch)

# 转换为量化模型
quantized_model = convert(prepared_model)

from torch.quantization import prepare_qat

# 训练时就模拟量化
model.train()
model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
qat_model = prepare_qat(model)

# 正常训练
for epoch in range(num_epochs):
    for batch in train_loader:
        output = qat_model(batch)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

# 转换为真正的量化模型
quantized_model = convert(qat_model.eval())

from transformers import AutoModelForCausalLM, BitsAndBytesConfig

# 4-bit 量化（QLoRA 用的）
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True
)

model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b",
    quantization_config=bnb_config,
    device_map="auto"
)

# 转换为 GGUF
python convert-hf-to-gguf.py model/ --outtype f16

# 量化为 Q4_K_M（推荐）
./quantize model-f16.gguf model-q4_k_m.gguf Q4_K_M

# 量化级别选择
# Q4_K_S: 最激进，4-bit
# Q4_K_M: 平衡（推荐）
# Q5_K_M: 更高精度
# Q8_0:   接近 FP16

class DistillationTrainer:
    def __init__(self, teacher, student, temperature=4.0, alpha=0.5):
        self.teacher = teacher.eval()  # 冻结
        self.student = student
        self.temperature = temperature
        self.alpha = alpha

    def compute_loss(self, x, y):
        # Teacher 输出（不计算梯度）
        with torch.no_grad():
            teacher_logits = self.teacher(x)

        # Student 输出
        student_logits = self.student(x)

        # 硬标签损失
        hard_loss = F.cross_entropy(student_logits, y)

        # 软标签损失（KL 散度）
        soft_loss = F.kl_div(
            F.log_softmax(student_logits / self.temperature, dim=-1),
            F.softmax(teacher_logits / self.temperature, dim=-1),
            reduction='batchmean'
        ) * (self.temperature ** 2)

        # 加权组合
        return self.alpha * hard_loss + (1 - self.alpha) * soft_loss

class FeatureDistillationModel(nn.Module):
    def __init__(self, teacher, student):
        super().__init__()
        self.teacher = teacher
        self.student = student
        # 适配层（对齐维度）
        self.adapter = nn.Linear(student_dim, teacher_dim)

    def forward(self, x):
        # Teacher 中间层特征
        with torch.no_grad():
            teacher_features = self.teacher.get_intermediate(x)

        # Student 中间层特征
        student_features = self.student.get_intermediate(x)
        adapted = self.adapter(student_features)

        # 特征蒸馏损失
        feature_loss = F.mse_loss(adapted, teacher_features)
        return feature_loss

# GPT-4 蒸馏到小模型
def generate_distillation_data(teacher, prompts):
    """用 Teacher 生成训练数据"""
    dataset = []
    for prompt in prompts:
        # Teacher 生成高质量回答
        response = teacher.generate(prompt)
        dataset.append({'instruction': prompt, 'output': response})
    return dataset

# 然后用这些数据 SFT 训练 Student

# Transformer 推理的 KV Cache
# 不用每次重新计算历史 token 的 K、V

class CachedAttention(nn.Module):
    def forward(self, x, past_kv=None):
        q, k, v = self.qkv(x).chunk(3, dim=-1)

        if past_kv is not None:
            past_k, past_v = past_kv
            k = torch.cat([past_k, k], dim=-2)
            v = torch.cat([past_v, v], dim=-2)

        new_kv = (k, v)

        # 只对当前 q 和累积的 k/v 计算 attention
        attn = attention(q, k, v)
        return attn, new_kv

# 标准 Attention：IO 密集（频繁读写 HBM）
# Flash Attention：分块计算，减少 HBM 读写

from flash_attn import flash_attn_func

attn_output = flash_attn_func(q, k, v, causal=True)

# vLLM 自动实现连续批处理
python -m vllm.entrypoints.api_server \
    --model meta-llama/Llama-2-7b \
    --enable-chunked-prefill \
    --max-num-batched-tokens 4096

# 1. 温度和采样参数
generation_config = GenerationConfig(
    temperature=0.7,
    top_p=0.9,
    top_k=50,
    do_sample=True,
    max_new_tokens=512,
    num_beams=1,        # 贪婪解码更快
    repetition_penalty=1.1
)

# 2. Speculative Decoding（推测解码）
# 用小模型先生成 draft，大模型验证
from transformers import AutoModelForCausalLM

draft_model = AutoModelForCausalLM.from_pretrained("small-model")
target_model = AutoModelForCausalLM.from_pretrained("large-model")

# 3. 模型合并（LoRA 合并）
merged_model = model.merge_and_unload()
merged_model.save_pretrained("./merged")

from transformers import LogitsProcessor, LogitsProcessorList

class NoBadWordsProcessor(LogitsProcessor):
    """禁止某些词"""
    def __init__(self, bad_words_ids):
        self.bad_words_ids = bad_words_ids

    def __call__(self, input_ids, scores):
        for bad_word_id in self.bad_words_ids:
            scores[:, bad_word_id] = -float('inf')
        return scores

class TemperatureLogitsProcessor(LogitsProcessor):
    """温度调节"""
    def __init__(self, temperature):
        self.temperature = temperature

    def __call__(self, input_ids, scores):
        return scores / self.temperature

class TopKLogitsProcessor(LogitsProcessor):
    """Top-K 采样"""
    def __init__(self, top_k):
        self.top_k = top_k

    def __call__(self, input_ids, scores):
        top_k = min(self.top_k, scores.size(-1))
        indices_to_remove = scores < torch.topk(scores, top_k)[0][..., -1, None]
        scores[indices_to_remove] = -float('inf')
        return scores

# 使用
logits_processor = LogitsProcessorList([
    NoBadWordsProcessor([bad_word_id_1, bad_word_id_2]),
    TemperatureLogitsProcessor(0.7),
    TopKLogitsProcessor(50),
])

output = model.generate(
    input_ids,
    logits_processor=logits_processor,
    max_new_tokens=100
)

from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer
import faiss

class HybridSearch:
    def __init__(self, documents):
        self.documents = documents

        # BM25 索引
        tokenized_docs = [doc.split() for doc in documents]
        self.bm25 = BM25Okapi(tokenized_docs)

        # 向量索引
        self.embedder = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
        embeddings = self.embedder.encode(documents)
        self.vector_index = faiss.IndexFlatIP(embeddings.shape[1])
        self.vector_index.add(embeddings.astype('float32'))

    def search(self, query, top_k=5, alpha=0.5):
        # BM25 搜索
        bm25_scores = self.bm25.get_scores(query.split())

        # 向量搜索
        query_emb = self.embedder.encode([query]).astype('float32')
        _, vector_indices = self.vector_index.search(query_emb, top_k * 2)
        vector_scores = np.zeros(len(self.documents))
        for rank, idx in enumerate(vector_indices[0]):
            vector_scores[idx] = 1.0 / (rank + 1)  # 倒数排名

        # 归一化
        bm25_norm = (bm25_scores - bm25_scores.min()) / (bm25_scores.max() - bm25_scores.min() + 1e-8)
        vector_norm = (vector_scores - vector_scores.min()) / (vector_scores.max() - vector_scores.min() + 1e-8)

        # 加权融合
        final_scores = alpha * vector_norm + (1 - alpha) * bm25_norm
        top_indices = np.argsort(final_scores)[::-1][:top_k]

        return [(self.documents[i], final_scores[i]) for i in top_indices]

from sentence_transformers import CrossEncoder

# 用 Cross Encoder 重排序（更准但更慢）
reranker = CrossEncoder('BAAI/bge-reranker-base')

def hybrid_search_with_rerank(query, top_k=5):
    # 1. 初步召回（快）
    candidates = hybrid_search.search(query, top_k=top_k * 4)

    # 2. 重排序（准）
    pairs = [(query, doc) for doc, _ in candidates]
    rerank_scores = reranker.predict(pairs)

    # 3. 返回 top_k
    ranked = sorted(zip(candidates, rerank_scores), key=lambda x: x[1], reverse=True)
    return ranked[:top_k]

# 动态批处理：等待一小段时间收集请求
import asyncio
from collections import deque

class DynamicBatcher:
    def __init__(self, model, max_batch=32, max_wait=0.05):
        self.model = model
        self.max_batch = max_batch
        self.max_wait = max_wait
        self.queue = deque()

    async def predict(self, input_data):
        future = asyncio.Future()
        self.queue.append((input_data, future))
        return await future

    async def batch_worker(self):
        while True:
            if len(self.queue) >= self.max_batch:
                batch = [self.queue.popleft() for _ in range(self.max_batch)]
            elif self.queue:
                await asyncio.sleep(self.max_wait)
                batch = list(self.queue)
                self.queue.clear()
            else:
                await asyncio.sleep(0.01)
                continue

            # 批量推理
            inputs = [item[0] for item in batch]
            outputs = self.model(inputs)

            # 返回结果
            for (_, future), output in zip(batch, outputs):
                future.set_result(output)

# FastAPI 异步推理
from fastapi import FastAPI
from fastapi.concurrency import run_in_threadpool

app = FastAPI()
model = load_model()

@app.post("/predict")
async def predict(input_data: dict):
    # 模型推理放到线程池，不阻塞事件循环
    result = await run_in_threadpool(model.predict, input_data)
    return {"result": result}

from functools import lru_cache
import hashlib

# 结果缓存
@lru_cache(maxsize=1000)
def cached_predict(input_hash):
    return model.predict(input_hash)

def predict(input_data):
    input_hash = hashlib.md5(str(input_data).encode()).hexdigest()
    return cached_predict(input_hash)

# 语义缓存（相似输入复用）
class SemanticCache:
    def __init__(self, threshold=0.95):
        self.embedder = SentenceTransformer(...)
        self.cache = []

    def get(self, query):
        query_emb = self.embedder.encode([query])
        for cached_query, cached_emb, result in self.cache:
            sim = cosine_similarity(query_emb, cached_emb)
            if sim > self.threshold:
                return result
        return None