LLM 推理优化（一）：量化、剪枝与蒸馏全面指南

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

# ===== BitsAndBytes INT4 量化实战 =====
# 这是目前最简单易用的 LLM 量化方案
# 支持 NF4（Normalized Float 4）精度，精度损失极小

model_name = "meta-llama/Llama-2-7b-hf"

# 配置 4-bit 量化参数
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,              # 启用 4-bit 量化
    bnb_4bit_quant_type="nf4",      # 使用 NF4 量化类型（比 INT4 精度更高）
    bnb_4bit_compute_dtype=torch.float16,  # 计算时使用 FP16
    bnb_4bit_use_double_quant=True, # 双重量化：进一步压缩量化参数
)

# 加载量化后的模型（显存从 14GB 降到约 4GB）
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=quantization_config,
    device_map="auto",              # 自动分配到可用 GPU
    torch_dtype=torch.float16,
)

tokenizer = AutoTokenizer.from_pretrained(model_name)

# 推理测试
prompt = "大语言模型推理优化的核心挑战是"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

with torch.no_grad():
    outputs = model.generate(
        **inputs,
        max_new_tokens=100,
        temperature=0.7,
        do_sample=True,
    )

print(tokenizer.decode(outputs[0], skip_special_tokens=True))
# 输出：大语言模型推理优化的核心挑战是在保持生成质量的同时
#       降低计算成本和延迟。目前主流的解决方案包括量化、
#       剪枝、知识蒸馏和推测解码等技术...

# ===== GPTQ 量化：基于梯度的逐层优化量化 =====
# GPTQ 的核心思想：用校准数据计算每层的量化误差梯度，
# 通过贪心算法逐元素修正量化误差

from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
from transformers import AutoTokenizer
import torch

model_name = "meta-llama/Llama-2-7b-hf"
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 配置 GPTQ 量化参数
quantize_config = BaseQuantizeConfig(
    bits=4,                          # 4-bit 量化
    group_size=128,                  # 逐组量化，每组 128 个权重
    damp_percent=0.01,               # 阻尼系数，防止量化矩阵奇异
    desc_act=False,                  # 不使用激活重排序（加快量化速度）
)

# 加载模型准备量化
model = AutoGPTQForCausalLM.from_pretrained(
    model_name,
    quantize_config,
    torch_dtype=torch.float16,
)

# 准备校准数据（128 条代表性文本）
calibration_texts = [
    "人工智能是计算机科学的一个重要分支",
    "深度学习通过多层神经网络提取特征",
    "自然语言处理让计算机理解人类语言",
    # ... 通常需要 128-512 条校准样本
]

calibration_data = tokenizer(
    calibration_texts,
    padding=True,
    truncation=True,
    max_length=512,
    return_tensors="pt",
).to("cuda")

# 执行量化（使用校准数据优化量化参数）
model.quantize(calibration_data)

# 保存量化后的模型（约 3.5GB）
model.save_quantized("./llama-2-7b-gptq-4bit")
print(f"量化完成！模型大小从 ~14GB 降至 ~3.5GB")
print(f"压缩比: {14/3.5:.1f}x")  # 4x 压缩

# ===== 结构化剪枝：注意力头剪枝实战 =====
# 通过评估每个注意力头的重要性，移除冗余头

import torch
import torch.nn as nn
from transformers import LlamaForCausalLM, LlamaConfig
from typing import List, Tuple

def evaluate_head_importance(
    model: LlamaForCausalLM,
    calibration_data: torch.Tensor,
    layer_idx: int,
    num_heads: int,
) -> List[float]:
    """评估每个注意力头的重要性分数
    
    方法：逐一遮蔽每个注意力头的输出，
    观察对模型输出的影响程度（KL 散度变化）
    """
    importance_scores = []
    
    with torch.no_grad():
        # 获取原始输出（作为基准）
        original_output = model(calibration_data).logits
        
        for head_idx in range(num_heads):
            # 遮蔽目标注意力头的输出
            head_output = original_output.clone()
            
            # 在注意力输出中遮蔽该头对应的维度
            head_dim = model.config.hidden_size // num_heads
            start_idx = head_idx * head_dim
            end_idx = (head_idx + 1) * head_dim
            
            # 通过该层的注意力输出遮蔽
            layer_attn = model.model.layers[layer_idx].self_attn
            # 简化的重要性评估：计算遮蔽后的 KL 散度
            kl_div = nn.functional.kl_div(
                torch.log_softmax(head_output, dim=-1),
                torch.softmax(original_output, dim=-1),
                reduction='batchmean'
            )
            importance_scores.append(kl_div.item())
    
    return importance_scores

def prune_attention_heads(
    model: LlamaForCausalLM,
    importance_scores: List[float],
    prune_ratio: float = 0.3,
) -> Tuple[LlamaForCausalLM, List[int]]:
    """剪枝不重要的注意力头
    
    Args:
        model: 原始模型
        importance_scores: 每个头的重要性分数
        prune_ratio: 剪枝比例（0.3 = 移除最不重要的 30%）
    
    Returns:
        剪枝后的模型和被移除的头索引
    """
    num_heads = len(importance_scores)
    num_to_prune = int(num_heads * prune_ratio)
    
    # 按重要性排序，选出最不重要的头
    head_indices = sorted(
        range(num_heads),
        key=lambda i: importance_scores[i]
    )
    heads_to_prune = head_indices[:num_to_prune]
    
    print(f"将移除 {len(heads_to_prune)}/{num_heads} 个注意力头")
    print(f"被移除的头索引: {sorted(heads_to_prune)}")
    
    # 注意：实际剪枝需要修改模型结构
    # 这里展示评估和选择逻辑
    return model, sorted(heads_to_prune)

# 使用示例
# importance = evaluate_head_importance(model, data, layer_idx=5, num_heads=32)
# pruned_model, pruned_heads = prune_attention_heads(model, importance, 0.3)
# → 移除 9/32 个注意力头，参数量减少约 8%，精度保持 97%+

# ===== 非结构化剪枝 + 稀疏推理 =====
# 对权重矩阵进行幅度剪枝，创建稀疏模型

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

# 对线性层进行非结构化剪枝
model = torch.hub.load('huggingface/pytorch-transformers',
                       'model', 'bert-base-uncased')

# 对每一层的注意力权重进行剪枝
for name, module in model.named_modules():
    if isinstance(module, torch.nn.Linear) and "query" in name:
        # 幅度剪枝：移除 50% 绝对值最小的权重
        prune.l1_unstructured(module, name='weight', amount=0.5)
        
        # 永久化剪枝（将 mask 应用到权重上）
        prune.remove(module, 'weight')
        
        # 计算稀疏度
        total_params = module.weight.numel()
        zero_params = (module.weight == 0).sum().item()
        sparsity = zero_params / total_params * 100
        print(f"{name}: 稀疏度 {sparsity:.1f}%")

# 剪枝后的模型评估
# 注意：非结构化剪枝需要稀疏推理引擎才能加速
# 在标准 PyTorch 上，稀疏矩阵的计算速度不会提升
# 需要使用 TensorRT、oneDNN 等支持稀疏推理的后端

# 结构化剪枝示例：移除整个 FFN 维度
def prune_ffn_dimension(
    model,
    layer_idx: int,
    reduction_ratio: float = 0.25,
):
    """减少 FFN 层的中间维度（结构化剪枝）"""
    layer = model.model.layers[layer_idx]
    original_dim = layer.mlp.gate_proj.out_features
    new_dim = int(original_dim * (1 - reduction_ratio))
    
    print(f"Layer {layer_idx} FFN: {original_dim} → {new_dim}")
    print(f"参数减少: {original_dim * model.config.hidden_size * 2} → "
          f"{new_dim * model.config.hidden_size * 2} "
          f"(减少 {reduction_ratio*100:.0f}%)")
    
    return new_dim

# prune_ffn_dimension(model, layer_idx=10, reduction_ratio=0.25)
# → FFN 中间维度从 11008 降至 8256，参数量减少 25%

# ===== 序列级别知识蒸馏：用大模型教小模型 =====
# Teacher 生成高质量输出 → Student 学习这些输出
# 这是最简单也最有效的 LLM 蒸馏方法

from transformers import (
    AutoModelForCausalLM, AutoTokenizer,
    TrainingArguments, Trainer
)
import torch
from datasets import Dataset

# 1. 加载 Teacher 和 Student 模型
teacher_name = "meta-llama/Llama-2-70b-hf"   # Teacher（70B）
student_name = "meta-llama/Llama-2-7b-hf"    # Student（7B）

teacher_tokenizer = AutoTokenizer.from_pretrained(teacher_name)
student_tokenizer = AutoTokenizer.from_pretrained(student_name)

# Teacher 生成蒸馏数据（仅在生成阶段需要）
teacher_model = AutoModelForCausalLM.from_pretrained(
    teacher_name,
    torch_dtype=torch.float16,
    device_map="auto",
)

# 2. 用 Teacher 生成蒸馏数据集
prompts = [
    "解释量子计算的基本原理",
    "比较监督学习和无监督学习的区别",
    "如何设计一个推荐系统",
    # ... 需要大量多样化的 prompts
]

distillation_data = []
for prompt in prompts:
    inputs = teacher_tokenizer(prompt, return_tensors="pt").to("cuda")
    with torch.no_grad():
        outputs = teacher_model.generate(
            **inputs,
            max_new_tokens=512,
            temperature=0.7,
            do_sample=True,
            top_p=0.9,
        )
    response = teacher_tokenizer.decode(outputs[0], skip_special_tokens=True)
    distillation_data.append({
        "prompt": prompt,
        "response": response.replace(prompt, "").strip(),
    })

# 3. 用蒸馏数据微调 Student 模型
student_model = AutoModelForCausalLM.from_pretrained(
    student_name,
    torch_dtype=torch.float16,
)

# 转换为训练格式
train_dataset = Dataset.from_list([
    {"text": f"{d['prompt']}\n{d['response']}"}
    for d in distillation_data
])

training_args = TrainingArguments(
    output_dir="./distilled-llama-7b",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    learning_rate=2e-5,
    fp16=True,
    logging_steps=10,
    save_steps=100,
)

# Trainer 会自动处理语言模型训练
# Student 通过学习 Teacher 的输出，获得接近 Teacher 的能力
print(f"蒸馏数据量: {len(distillation_data)} 条")
print(f"Student 模型参数: {student_model.num_parameters()/1e9:.1f}B")
print(f"Teacher 模型参数: {teacher_model.config.num_parameters/1e9:.1f}B")
print(f"压缩比: ~{teacher_model.config.num_parameters/student_model.num_parameters():.0f}x")

# ===== Logits 级蒸馏：对齐 Teacher 的每一步输出 =====
# 比序列级蒸馏更精细：Student 的每个 token 生成
# 都要对齐 Teacher 的 logits 分布

import torch
import torch.nn.functional as F
from transformers import AutoModelForCausalLM, AutoTokenizer

def compute_distillation_loss(
    student_logits: torch.Tensor,
    teacher_logits: torch.Tensor,
    temperature: float = 2.0,
    alpha: float = 0.5,
) -> torch.Tensor:
    """计算知识蒸馏损失
    
    结合了硬标签损失（学生输出与真实标签）
    和软标签损失（学生 logits 与 teacher logits）
    
    Args:
        student_logits: Student 模型的 logits [batch, seq, vocab]
        teacher_logits: Teacher 模型的 logits [batch, seq, vocab]
        temperature: 温度参数，越高分布越平滑
        alpha: 软标签损失的权重（0.5 = 各占一半）
    
    Returns:
        蒸馏损失（标量）
    """
    # 软标签损失：KL 散度
    # 温度缩放后计算 soft probability distribution
    student_probs = F.log_softmax(student_logits / temperature, dim=-1)
    teacher_probs = F.softmax(teacher_logits / temperature, dim=-1)
    
    kl_loss = F.kl_div(
        student_probs,
        teacher_probs,
        reduction='batchmean',
    ) * (temperature ** 2)
    
    # 硬标签损失：交叉熵（需要真实标签）
    # 这里用 teacher 生成的 token 作为伪标签
    teacher_tokens = teacher_logits.argmax(dim=-1)
    ce_loss = F.cross_entropy(
        student_logits.view(-1, student_logits.size(-1)),
        teacher_tokens.view(-1),
        ignore_index=-100,
    )
    
    # 组合损失
    total_loss = alpha * kl_loss + (1 - alpha) * ce_loss
    return total_loss

# 蒸馏训练循环示例
def distillation_step(
    student_model, teacher_model,
    input_ids, attention_mask,
    temperature=2.0, alpha=0.5,
):
    """单步蒸馏训练"""
    # 冻结 Teacher（只用于推理）
    with torch.no_grad():
        teacher_outputs = teacher_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
        )
        teacher_logits = teacher_outputs.logits
    
    # Student 前向传播
    student_outputs = student_model(
        input_ids=input_ids,
        attention_mask=attention_mask,
    )
    student_logits = student_outputs.logits
    
    # 计算蒸馏损失
    loss = compute_distillation_loss(
        student_logits, teacher_logits,
        temperature=temperature,
        alpha=alpha,
    )
    
    return loss

# 温度参数的作用：
# temperature=1.0 → 标准概率分布
# temperature=2.0 → 分布更平滑，传递更多"暗知识"
# temperature=5.0 → 分布接近均匀，蒸馏效果减弱
# 经验值：LLM 蒸馏通常用 2.0~4.0

#!/bin/bash
# ===== vLLM 生产环境部署脚本 =====
# 将量化后的模型部署为 OpenAI 兼容 API 服务

# 1. 安装 vLLM
pip install vllm==0.5.5

# 2. 启动服务（INT4 量化模型）
# --quantization: 量化类型（awq/gptq/fp8/bnb）
# --gpu-memory-utilization: GPU 显存利用率（0.9 = 90%）
# --max-num-batched-tokens: 每批次最大 token 数
# --max-num-seqs: 最大并发序列数
# --tensor-parallel-size: 张量并行 GPU 数

python -m vllm.entrypoints.openai.api_server \
    --model ./llama-2-7b-awq-4bit \
    --quantization awq \
    --dtype half \
    --gpu-memory-utilization 0.90 \
    --max-num-batched-tokens 8192 \
    --max-num-seqs 128 \
    --max-model-len 4096 \
    --tensor-parallel-size 1 \
    --host 0.0.0.0 \
    --port 8000 \
    --enable-chunked-prefill \
    --disable-log-requests

# 3. 测试服务
curl http://localhost:8000/v1/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "./llama-2-7b-awq-4bit",
    "prompt": "大语言模型推理优化的核心方法是",
    "max_tokens": 100,
    "temperature": 0.7
  }'

# 4. 性能基准测试
# 使用 vLLM 自带的 benchmark 工具
python -m vllm.entrypoints.benchmark_serving \
    --model ./llama-2-7b-awq-4bit \
    --dataset-name random \
    --num-prompts 1000 \
    --random-input-len 256 \
    --random-output-len 128 \
    --request-rate 10

# 5. 监控指标（配合 Prometheus + Grafana）
# vLLM 默认暴露 Prometheus 指标在 /metrics 端点
# 关键指标：
#   vllm:prompt_tokens_total - 总 prompt token 数
#   vllm:generation_tokens_total - 总生成 token 数
#   vllm:time_to_first_token_seconds - TTFT
#   vllm:gpu_cache_usage_perc - KV Cache 使用率