AI模型压缩进阶踩坑记录

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch.quantization as quant

# 加载原始模型
model_path = "/models/llama-2-7b-chat"
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.float16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_path)

# 准备校准数据（需要真实输入，不是随便生成的）
def prepare_calibration_data(tokenizer, num_samples=100):
    samples = []
    with open("/data/calibration_text.txt", "r", encoding="utf-8") as f:
        text = f.read()
        sentences = text.split("\n")
        for i in range(min(num_samples, len(sentences))):
            encoded = tokenizer(sentences[i], return_tensors="pt", truncation=True, max_length=512)
            samples.append(encoded["input_ids"])
    return torch.cat(samples, dim=0)

calibration_data = prepare_calibration_data(tokenizer)

# 动态量化（只量化权重，不量化激活）
model_quantized = quant.quantize_dynamic(
    model,
    {torch.nn.Linear},  # 只量化 Linear 层
    dtype=torch.qint8
)

# 保存量化后的模型
torch.save(model_quantized.state_dict(), "/models/llama-2-7b-chat-dynamic-int8.pt")

# 静态量化配置
quantization_config = quant.get_default_qconfig("x86")
model.qconfig = quantization_config
model_prepared = quant.prepare(model, inplace=True)

# 校准过程
with torch.no_grad():
    for i in range(0, len(calibration_data), 4):
        batch = calibration_data[i:i+4].to(model.device)
        _ = model(batch)
        if i % 20 == 0:
            print(f"Calibration progress: {i}/{len(calibration_data)}")

# 转换为量化模型
model_quantized = quant.convert(model_prepared, inplace=True)

def get_layer_name(module):
    for name, child in module.named_modules():
        yield child, name

# 为不同层设置不同的量化策略
for module, name in get_layer_name(model):
    if "LayerNorm" in name or "layernorm" in name.lower():
        module.qconfig = None  # 不量化 LayerNorm

import torch.nn.utils.prune as prune

def compute_importance_scores(model, calibration_data):
    """计算每个层的重要性分数"""
    importance_scores = {}
    
    with torch.no_grad():
        for batch_idx in range(0, len(calibration_data), 8):
            batch = calibration_data[batch_idx:batch_idx+8].to(model.device)
            output = model(batch, output_hidden_states=True)
            
            # 使用梯度信息或激活值统计作为重要性指标
            # 这里用输出方差作为简化示例
            hidden_states = output.hidden_states
            for layer_idx, hidden_state in enumerate(hidden_states[1:]):
                variance = torch.var(hidden_state, dim=[0, 1]).mean().item()
                importance_scores[f"layer_{layer_idx}"] = variance
    
    return importance_scores

importance_scores = compute_importance_scores(model, calibration_data)

# 对重要性低的层进行剪枝
threshold = 0.3  # 剪掉重要性分数低于阈值的层
for name, module in model.named_modules():
    if "layers" in name:
        layer_num = int(name.split(".")[3])
        if importance_scores.get(f"layer_{layer_num}", 0) < threshold:
            prune.l1_unstructured(module, name="weight", amount=0.4)  # 剪掉 40% 参数

from transformers import TrainingArguments, Trainer

# 准备少量训练数据（剪枝微调不需要太多）
train_dataset = prepare_training_data("/data/fine_tune_samples.jsonl", max_samples=1000)

training_args = TrainingArguments(
    output_dir="/models/pruned_finetuned",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    learning_rate=1e-5,
    logging_steps=10,
    save_steps=50,
    fp16=True,
)

trainer = Trainer(
    model=model_quantized,
    args=training_args,
    train_dataset=train_dataset,
)

trainer.train()

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling
)

# 教师模型
teacher_model = AutoModelForCausalLM.from_pretrained(
    "/models/llama-2-7b-chat",
    torch_dtype=torch.float16,
    device_map="auto"
)
teacher_model.eval()

# 学生模型（可以从头训练或用预训练的小模型）
student_model = AutoModelForCausalLM.from_pretrained(
    "/models/tinyllama-3b",
    torch_dtype=torch.float16
)

# 蒸馏损失函数
class DistillationLoss:
    def __init__(self, temperature=2.0, alpha=0.5):
        self.temperature = temperature
        self.alpha = alpha
        self.kl_div = torch.nn.KLDivLoss(reduction="batchmean")
        self.ce_loss = torch.nn.CrossEntropyLoss()
    
    def __call__(self, student_outputs, teacher_outputs, labels):
        # 软标签损失（教师输出）
        student_log_probs = torch.nn.functional.log_softmax(
            student_outputs.logits / self.temperature, dim=-1
        )
        teacher_probs = torch.nn.functional.softmax(
            teacher_outputs.logits / self.temperature, dim=-1
        )
        distill_loss = self.kl_div(student_log_probs, teacher_probs) * (self.temperature ** 2)
        
        # 硬标签损失（真实标签）
        ce_loss = self.ce_loss(
            student_outputs.logits.view(-1, student_outputs.logits.size(-1)),
            labels.view(-1)
        )
        
        return self.alpha * distill_loss + (1 - self.alpha) * ce_loss

distillation_loss = DistillationLoss(temperature=3.0, alpha=0.7)

# 训练参数
training_args = TrainingArguments(
    output_dir="/models/distilled_student",
    num_train_epochs=5,
    per_device_train_batch_size=8,
    gradient_accumulation_steps=2,
    learning_rate=5e-5,
    warmup_steps=100,
    logging_steps=20,
    save_steps=200,
    fp16=True,
    deepspeed="ds_config_zero3.json",  # 使用 DeepSpeed 优化显存使用
)

class DistillationTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        # 教师前向传播
        with torch.no_grad():
            teacher_outputs = teacher_model(**inputs)
        
        # 学生前向传播
        student_outputs = model(**inputs)
        
        # 计算蒸馏损失
        loss = distillation_loss(
            student_outputs, 
            teacher_outputs, 
            inputs["labels"]
        )
        
        return (loss, student_outputs) if return_outputs else loss

# 准备训练数据
train_dataset = prepare_training_data("/data/distillation_data.jsonl")

trainer = DistillationTrainer(
    model=student_model,
    args=training_args,
    train_dataset=train_dataset,
    data_collator=DataCollatorForLanguageModeling(
        tokenizer=tokenizer, 
        mlm=False
    ),
)

trainer.train()

# 转换为 ONNX 格式
python -m transformers.onnx \
  --model /models/llama-2-7b-chat-int8 \
  --feature_map causal-lm-with-past \
  /models/llama-2-7b-chat-int8.onnx

# 转换为 GGML 格式（适合 CPU 部署）
python convert_llama_weights_to_ggml.py \
  /models/llama-2-7b-chat-int8 \
  /models/llama-2-7b-chat-int8-ggml \
  q8_0

# 转换为 GGUF 格式（新版 llama.cpp）
python convert.py /models/llama-2-7b-chat-int8 \
  --outfile /models/llama-2-7b-chat-int8.gguf \
  --outtype q8_0

# 使用 vLLM 加速推理
from vllm import LLM, SamplingParams

llm = LLM(
    model="/models/llama-2-7b-chat-int8",
    quantization="awq",  # 使用 AWQ 量化
    tensor_parallel_size=2,  # 使用 2 张 GPU
    max_model_len=2048,
    gpu_memory_utilization=0.9,
)

sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=512,
)

def generate_response(prompt):
    outputs = llm.generate([prompt], sampling_params)
    return outputs[0].outputs[0].text

# 使用 TensorRT-LLM（需要先转换模型）
# python convert_checkpoint.py --model_dir /models/llama-2-7b-chat-int8 \
#   --output_dir /models/llama-2-7b-chat-int8-trtllm

import gc
import torch

def cleanup_memory():
    gc.collect()
    torch.cuda.empty_cache()
    torch.cuda.synchronize()

# 在推理循环中定期清理
request_count = 0
while True:
    # 处理请求
    process_request()
    
    request_count += 1
    if request_count % 100 == 0:
        cleanup_memory()

# 限制 KV Cache 大小
llm = LLM(
    model="/models/llama-2-7b-chat-int8",
    quantization="awq",
    max_model_len=1024,  # 从 2048 减到 1024
    enable_prefix_caching=True,  # 启用前缀缓存
)

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="Optimized LLM API")

class GenerationRequest(BaseModel):
    prompt: str
    max_tokens: int = 256
    temperature: float = 0.7

@app.post("/generate")
async def generate(request: GenerationRequest):
    try:
        logger.info(f"Processing request: {request.prompt[:50]}...")
        response = generate_response(request.prompt)
        return {"response": response}
    except Exception as e:
        logger.error(f"Error processing request: {e}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    return {"status": "healthy", "model": "llama-2-7b-chat-int8"}

if __name__ == "__main__":
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8000,
        workers=2,  # 根据 GPU 数量调整
        log_level="info"
    )

gunicorn -w 4 -k uvicorn.workers.UvicornWorker \
  --bind 0.0.0.0:8000 \
  --timeout 300 \
  --access-logfile - \
  --error-logfile - \
  api_server:app