大模型压缩常用方法

一、为什么需要压缩

Llama 3.1 405B FP16 ≈ 810 GB 显存。单张 H100 仅 80 GB。

压缩目标：减少参数量 / 降低精度 / 迁移知识，在精度损失可控的前提下降低部署门槛。

三种核心方法：

方法	原理	适用阶段
剪枝	移除贡献小的权重或神经元	训练后 / 训练中
量化	降低数值精度（FP32→INT8/INT4）	推理部署
蒸馏	大模型（teacher）指导小模型（student）训练	训练阶段

二、剪枝（Pruning）

2.1 原理

训练后的网络中存在大量接近零的权重，对输出贡献可忽略。剪枝即移除这些冗余连接。

两种粒度：

非结构化剪枝：逐权重置零(小于某个阈值)，矩阵变稀疏。精度损失小，但稀疏矩阵在 GPU 上的加速有限。

结构化剪枝：按神经元 / attention head / 通道整块移除。矩阵仍为稠密但尺寸减小，GPU 加速明显，但精度下降更大。

2.2 代码

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

fc = torch.nn.Linear(512, 256)
print(f"原始非零参数: {(fc.weight != 0).sum().item()}")

# 非结构化：L1 范数最小的 30% 权重置零
prune.l1_unstructured(fc, name="weight", amount=0.3)
print(f"剪枝后: {(fc.weight != 0).sum().item()}")
prune.remove(fc, "weight")  # 固化 mask

# 结构化：按 L2 范数剪掉 20% 的输出神经元
fc2 = torch.nn.Linear(512, 256)
prune.ln_structured(fc2, name="weight", amount=0.2, n=2, dim=1)

# 迭代剪枝：多次剪枝 + 微调
def iterative_prune(model, rate=0.2, rounds=3):
    for r in range(rounds):
        for m in model.modules():
            if isinstance(m, torch.nn.Linear):
                prune.l1_unstructured(m, name="weight", amount=rate)
        # trainer.train(epochs=2)  ← 每次剪枝后微调
    for m in model.modules():
        if isinstance(m, torch.nn.Linear):
            try: prune.remove(m, "weight")
            except: pass

三、量化（Quantization）

3.1 原理

将 FP32 参数映射到低比特整数，核心计算：

q = round(w / scale) + zero_point
scale = (max - min) / (qmax - qmin)

两种策略：

• PTQ（训练后量化）：模型训完后直接转换，无需重训。GPTQ、AWQ、bitsandbytes 均属此类。
• QAT（量化感知训练）：训练时模拟量化误差，模型自适应。精度更高但需重训。

3.2 代码

import torch

# ═══ PyTorch 原生动态量化（INT8）═══
class SimpleModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = torch.nn.Linear(512, 256)
        self.fc2 = torch.nn.Linear(256, 128)
    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

model = SimpleModel()
quantized = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)
# 权重大小约降至 1/4（FP32→INT8 + 量化参数开销）

# ═══ bitsandbytes 4-bit 量化（LLM 推理用）═══
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",            # NormalFloat4
    bnb_4bit_compute_dtype=torch.bfloat16, # 反量化后计算精度
    bnb_4bit_use_double_quant=True,        # 量化因子也量化
)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-1.5B-Instruct",
    quantization_config=bnb_config,
    device_map="auto",
)

四、知识蒸馏（Knowledge Distillation）

4.1 原理

用大模型（teacher）的 soft label 指导小模型（student）训练。核心在于 温度参数 T：

soft_label = softmax(logits / T)

T 越高，输出分布越平缓，类别间的相对关系（"暗知识"）越明显。T=1 退化为标准 softmax。

蒸馏 loss 由两部分加权：

loss = α · CrossEntropy(student, hard_label)  +  (1-α) · KL(student/T, teacher/T) · T²

4.2 代码

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

# Teacher（冻结）
teacher = nn.Sequential(
    nn.Linear(128, 256), nn.ReLU(),
    nn.Linear(256, 256), nn.ReLU(),
    nn.Linear(256, 10)
)
teacher.eval()
for p in teacher.parameters():
    p.requires_grad = False

# Student（目标小模型）
student = nn.Sequential(
    nn.Linear(128, 64), nn.ReLU(),
    nn.Linear(64, 10)
)

def distillation_loss(s_logits, t_logits, labels, T=4.0, alpha=0.7):
    soft_s = F.log_softmax(s_logits / T, dim=1)
    soft_t = F.softmax(t_logits / T, dim=1)
    kd = F.kl_div(soft_s, soft_t, reduction="batchmean") * (T * T)
    ce = F.cross_entropy(s_logits, labels)
    return alpha * ce + (1 - alpha) * kd

# 训练循环
optimizer = torch.optim.Adam(student.parameters(), lr=1e-3)
x, y = torch.randn(32, 128), torch.randint(0, 10, (32,))

for epoch in range(10):
    optimizer.zero_grad()
    with torch.no_grad():
        t_logits = teacher(x)
    s_logits = student(x)
    loss = distillation_loss(s_logits, t_logits, y)
    loss.backward()
    optimizer.step()

tp = sum(p.numel() for p in teacher.parameters())
sp = sum(p.numel() for p in student.parameters())
print(f"Teacher: {tp:,}  Student: {sp:,}  压缩比: {tp/sp:.1f}x")

五、低秩分解（Low-Rank Factorization）

将大矩阵分解为两个小矩阵的乘积：

W ∈ ℝ^(d×k) →  W ≈ A @ B,   A ∈ ℝ^(d×r), B ∈ ℝ^(r×k),  r ≪ min(d,k)
参数量：d×k  →  r×(d+k)

LoRA 即基于此原理——不更新原权重矩阵 W，而是训练两个低秩矩阵 A 和 B：

h = W₀x + (α/r) · BAx     # W₀ 冻结，只训 A 和 B

六、方法对比

 

方法	压缩率	精度损失	推理加速	是否需重训
非结构化剪枝	30~90%	低	不明显	否（可迭代微调）
结构化剪枝	20~50%	中	明显	建议微调
INT8 量化	~4×	极低	明显	否
INT4 量化	~8×	中	明显	否
知识蒸馏	由 student 规模决定	中	明显	是
低秩分解（LoRA）	可训参数 ↓99%+	低	推理时合并，无损	是

七、工程建议

推理部署：4-bit 量化（GPTQ / AWQ / bitsandbytes）是成本最低的加速方案，几分钟完成，通常精度损失 <1%。

端侧部署：蒸馏得到小模型 → INT8 量化 → 移动端推理。蒸馏决定能力上限，量化决定部署下限。

微调：LoRA（低秩分解）可训参数仅占 0.1%~1%，配合 4-bit 量化（QLoRA）可在单卡微调 70B 模型。