PyTorch模型量化实战：INT8动态量化部署指南

大模型部署场景中，量化压缩是把推理成本压下来的核心手段。PyTorch原生提供的动态量化API，可以在不重训练模型的前提下，把Linear层权重从FP32降到INT8，显存占用减少约75%，推理延迟下降30%-60%。本教程以一个文本分类模型为例，演示完整的INT8动态量化流程。

动态量化与静态量化的区别

PyTorch支持两种量化模式。动态量化（torch.ao.quantization.quantize_dynamic）在推理时实时计算激活值的scale和zero_point，无需校准数据集，适合RNN、Transformer等动态输入模型。静态量化需要提前用代表性数据集做校准，精度更高但流程复杂。对于LLM推理部署，如果追求快速落地，动态量化是最直接的方案。

环境准备与依赖安装

需要PyTorch 2.0以上版本，建议2.1+以获得完整的量化算子支持：

pip install torch torchvision transformers accelerate
import torch
print(torch.__version__)  # 确认 >= 2.0.0

CPU推理场景下不需要额外安装量化后端，PyTorch自带x86量化引擎。GPU量化需要安装PyTorch的量化专用wheel包或使用FBGEMM。

完整动态量化代码实现

以下代码对HuggingFace的BERT模型做动态量化，适用于文本分类、特征提取等任务：

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import time, os

model_name = 'bert-base-chinese'
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
model.eval()

# 动态量化：只量化Linear层的权重
quantized_model = torch.ao.quantization.quantize_dynamic(
    model,
    {torch.nn.Linear},
    dtype=torch.qint8
)

# 对比模型大小
def get_model_size(m):
    torch.save(m.state_dict(), 'temp.pt')
    size = os.path.getsize('temp.pt') / 1024 / 1024
    os.remove('temp.pt')
    return f'{size:.1f} MB'

print(f'原始模型: {get_model_size(model)}')
print(f'量化模型: {get_model_size(quantized_model)}')

量化后模型体积从约411MB降到约104MB，压缩比接近4:1。

推理性能对比测试

用相同的输入做推理，对比FP32和INT8的延迟差异：

text = '这是一段用于测试量化的中文文本'
inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=512)

def benchmark(model, inputs, n=100):
    with torch.no_grad():
        for _ in range(10):
            _ = model(**inputs)
        start = time.perf_counter()
        for _ in range(n):
            _ = model(**inputs)
        elapsed = time.perf_counter() - start
    return (elapsed / n) * 1000

fp32_latency = benchmark(model, inputs)
int8_latency = benchmark(quantized_model, inputs)
print(f'FP32: {fp32_latency:.2f} ms/op')
print(f'INT8: {int8_latency:.2f} ms/op')
print(f'加速比: {fp32_latency/int8_latency:.2f}x')

CPU环境下，INT8推理通常能获得1.5x-2x的加速比。加速比受CPU指令集支持情况影响，AVX512-VNNI指令集的处理器能获得更好的量化加速。

量化精度评估方法

动态量化可能导致精度下降，需要在验证集上做评估：

from sklearn.metrics import accuracy_score

def evaluate(model, dataset, tokenizer, batch_size=32):
    model.eval()
    preds, labels = [], []
    with torch.no_grad():
        for batch in DataLoader(dataset, batch_size=batch_size):
            inputs = tokenizer(batch['text'], padding=True, truncation=True, return_tensors='pt')
            outputs = model(**inputs)
            pred = torch.argmax(outputs.logits, dim=-1)
            preds.extend(pred.tolist())
            labels.extend(batch['label'])
    return accuracy_score(labels, preds)

fp32_acc = evaluate(model, val_dataset, tokenizer)
int8_acc = evaluate(quantized_model, val_dataset, tokenizer)
print(f'FP32准确率: {fp32_acc:.4f}')
print(f'INT8准确率: {int8_acc:.4f}')
print(f'精度损失: {(fp32_acc - int8_acc) * 100:.2f}%')

BERT-base模型做INT8动态量化，准确率损失通常在0.3%以内。如果损失超过1%，考虑使用静态量化或量化感知训练（QAT）。

部署注意事项

量化模型保存和加载方式与普通模型略有不同，建议使用TorchScript序列化：

# 保存
torch.jit.save(torch.jit.script(quantized_model), 'quantized_bert.pt')
# 加载
loaded_model = torch.jit.load('quantized_bert.pt')
loaded_model.eval()

TorchScript序列化后的量化模型可以脱离HuggingFace框架独立部署，适合在Docker容器中做微服务。实际工程中，建议配合Triton Inference Server或vLLM做批处理推理，能进一步压榨量化模型的吞吐量。