LLM 量化技术实战：从 FP16 到 INT4 的压缩艺术

# PTQ vs QAT 对比示例
import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

# ===== 方法 1：PTQ - 使用 bitsandbytes 进行 INT4 量化 =====
bnb_config_int4 = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",      # 归一化浮点 4-bit
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,  # 双重量化，进一步压缩
)

model_4bit = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3-8B",
    quantization_config=bnb_config_int4,
    device_map="auto",
)
# 优点：无需训练，直接加载量化后的模型
# 缺点：精度略低于 QAT

# ===== 方法 2：QAT - 量化感知训练 =====
class QuantizedModel(nn.Module):
    """量化感知训练的简化示例"""
    def __init__(self, base_model):
        super().__init__()
        # 插入伪量化节点（FakeQuantize）
        # 训练时模拟量化效果，推理时应用真实量化
        self.model = base_model
        self._insert_fake_quantizers()
    
    def _insert_fake_quantizers(self):
        for name, module in self.model.named_modules():
            if isinstance(module, nn.Linear):
                # 为权重和激活添加 FakeQuantize
                module.weight_fake_quant = torch.ao.nn.quantized.FakeQuantize(
                    observer=torch.ao.quantization.MinMaxObserver(),
                    quant_min=-128, quant_max=127,
                    dtype=torch.qint8, qscheme=torch.per_tensor_symmetric
                )

# QAT 训练流程：
# 1. 加载预训练模型
# 2. 插入 FakeQuantize 节点
# 3. 用校准数据集微调（通常 100-1000 步）
# 4. 导出量化模型（torch.quantization.convert）

print("PTQ 适合快速部署，QAT 适合追求极致精度")

# 使用 AutoGPTQ 量化模型
# 安装：pip install auto-gptq optimum

# Step 1: 下载模型
python -c "
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained('meta-llama/Llama-3-8B')
tokenizer.save_pretrained('./llama-3-8b')
"

# Step 2: 执行 GPTQ 量化（需要 ~20GB VRAM）
python -c "
from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
from transformers import AutoTokenizer

model_name_or_path = 'meta-llama/Llama-3-8B'
quantized_model_dir = 'llama-3-8b-gptq-int4'

quantize_config = BaseQuantizeConfig(
    bits=4,                # INT4 量化
    group_size=128,        # 分组大小，128 是最佳平衡点
    desc_act=False,        # 是否使用描述性激活（设为 False 加速推理）
)

# 加载模型
model = AutoGPTQForCausalLM.from_pretrained(
    model_name_or_path, quantize_config
)

# 准备校准数据（只需要 128 个样本）
from datasets import load_dataset
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
calibration_data = load_dataset('wikitext', 'wikitext-2-raw-v1', split='train')
examples = [tokenizer(text, return_tensors='pt') for text in calibration_data['text'][:128]]

# 执行量化
model.quantize(examples, batch_size=4)

# 保存量化模型
model.save_quantized(quantized_model_dir)
print(f'量化模型已保存到 {quantized_model_dir}')
"

# Step 3: 加载量化后的模型进行推理
python -c "
from auto_gptq import AutoGPTQForCausalLM
from transformers import AutoTokenizer

model = AutoGPTQForCausalLM.from_quantized(
    'llama-3-8b-gptq-int4',
    device='cuda:0',
    use_triton=True,  # 使用 Triton 加速
)
tokenizer = AutoTokenizer.from_pretrained('llama-3-8b-gptq-int4')

inputs = tokenizer('解释一下联邦学习的核心思想', return_tensors='pt').to('cuda')
output = model.generate(**inputs, max_new_tokens=200)
print(tokenizer.decode(output[0], skip_special_tokens=True))
"

# 使用 AWQ 量化模型
# 安装：pip install autoawq

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = "meta-llama/Llama-3-8B"
quant_path = "llama-3-8b-awq-int4"

# 加载模型和分词器
model = AutoAWQForCausalLM.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# 配置量化参数
quant_config = {
    "zero_point": True,     # 是否使用零点
    "q_group_size": 128,    # 分组大小（与 GPTQ 类似）
    "w_bit": 4,            # 4-bit 量化
    "version": "GEMM",      # 推理内核（GEMM 或 GEMV）
}

# 准备校准数据
calib_data = [
    "The quick brown fox jumps over the lazy dog.",
    "Machine learning is a subset of artificial intelligence.",
    # ... 更多校准文本（通常 128 条）
]

# 执行量化
model.quantize(tokenizer, quant_config=quant_config, calib_data=calib_data)

# 保存
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)

# ===== 推理 =====
from awq import AutoAWQForCausalLM

model = AutoAWQForCausalLM.from_quantized(
    quant_path,
    device_map="auto",
    fuse_layers=True,     # 融合层，加速推理
)

prompt = "解释一下 AWQ 量化的核心思想"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
output = model.generate(**inputs, max_new_tokens=200)
print(tokenizer.decode(output[0], skip_special_tokens=True))

# 内存对比
# FP16:  ~16GB
# AWQ INT4: ~4.5GB（含推理开销）
# 压缩比: ~3.5x

# 生产环境量化模型部署模板
import time
from transformers import AutoTokenizer, AutoModelForCausalLM
from vllm import LLM, SamplingParams

class QuantizedLLMService:
    """量化 LLM 生产服务"""
    
    def __init__(
        self,
        model_path: str,
        quantization: str = "gptq",  # gptq, awq, fp8
        gpu_memory_utilization: float = 0.9,
        max_model_len: int = 4096,
    ):
        # 使用 vLLM 部署量化模型
        self.llm = LLM(
            model=model_path,
            quantization=quantization,
            gpu_memory_utilization=gpu_memory_utilization,
            max_model_len=max_model_len,
            tensor_parallel_size=1,  # 多卡时调整
        )
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.sampling_params = SamplingParams(
            temperature=0.7,
            top_p=0.9,
            max_tokens=512,
        )
        
        # 性能监控
        self.request_count = 0
        self.total_latency = 0.0
    
    def generate(self, prompt: str) -> dict:
        """生成响应并记录性能指标"""
        start_time = time.time()
        
        outputs = self.llm.generate(
            [prompt], 
            self.sampling_params,
            use_tqdm=False
        )
        
        latency = time.time() - start_time
        self.request_count += 1
        self.total_latency += latency
        
        result = outputs[0].outputs[0].text
        tokens = outputs[0].outputs[0].token_ids
        
        return {
            "text": result,
            "tokens_generated": len(tokens),
            "latency_ms": latency * 1000,
            "tokens_per_sec": len(tokens) / latency,
        }
    
    def get_metrics(self) -> dict:
        """获取服务性能指标"""
        avg_latency = (
            self.total_latency / self.request_count 
            if self.request_count > 0 else 0
        )
        return {
            "total_requests": self.request_count,
            "avg_latency_ms": avg_latency * 1000,
        }


# ===== 量化方案决策辅助 =====
def choose_quantization(
    vram_gb: int,
    model_size_b: int,
    need_finetuning: bool = False,
    quality_priority: bool = False,
) -> str:
    """根据约束条件推荐量化方案"""
    fp16_vram = model_size_b * 2  # FP16 需要 ~2GB/1B参数
    
    if fp16_vram <= vram_gb and not quality_priority:
        return "fp16"  # 显存够用，不量化
    
    ratio = vram_gb / fp16_vram
    
    if ratio >= 0.5:
        return "int8"      # INT8（2x 压缩）
    elif need_finetuning:
        return "nf4_qlora" # QLoRA 微调
    elif quality_priority:
        return "gptq_int4" # GPTQ INT4（精度最优）
    else:
        return "awq_int4"  # AWQ INT4（速度最优）

# 使用示例
scheme = choose_quantization(
    vram_gb=24, model_size_b=30, 
    need_finetuning=True
)
print(f"推荐方案: {scheme}")  # → nf4_qlora