第3章量化实战

量化实战：使用bitsandbytes库对LLaMA模型进行4-bit量化

聊到量化，很多同学第一反应是「精度会不会掉得很惨？」。说实话，我刚开始接触4-bit量化时也有这个顾虑。但后来在项目中跑了几轮对比测试，发现只要方法得当，效果其实相当能打。今天我们就拿LLaMA模型开刀，手把手走一遍4-bit量化的完整流程。

为什么选bitsandbytes？

市面上量化工具不少，但我个人最常用的是bitsandbytes。原因很简单：它跟Hugging Face生态结合得最好。你想想看，我们平时加载模型用的就是transformers库，而bitsandbytes直接提供了量化后的模型加载接口，几乎不需要改代码。

另外一点，bitsandbytes支持NF4和FP4两种4-bit数据类型。NF4是专门为神经网络权重分布设计的，比普通FP4效果更好。我在一个70亿参数的LLaMA模型上做过对比，NF4比FP4的困惑度低了将近0.3个点。

核心要点：bitsandbytes的4-bit量化主要依赖两个技术——分块量化（Block-wise Quantization）和双重量化（Double Quantization）。前者把权重分成小块分别量化，减少异常值影响；后者对量化缩放因子再做一次量化，进一步压缩存储。

环境准备

动手之前，先把环境搭好。嗯，这里要注意版本兼容性问题，我踩过不少坑。

pip install bitsandbytes
pip install transformers
pip install accelerate
pip install torch

建议用bitsandbytes 0.41.0以上版本，对LLaMA的支持更完善。CUDA版本最好在11.8以上，不然有些算子跑不了。

注意：bitsandbytes目前只支持CUDA，CPU上跑不了。如果你在无GPU的环境下调试，可以用CPU模式模拟，但实际推理必须用GPU。

加载4-bit量化模型

核心代码其实就几行。我们直接看怎么加载一个4-bit量化的LLaMA模型：

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch

# 配置4-bit量化参数
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",      # 使用NF4数据类型
    bnb_4bit_compute_dtype=torch.float16,  # 计算时用float16
    bnb_4bit_use_double_quant=True   # 开启双重量化
)

# 加载模型
model_name = "meta-llama/Llama-2-7b-hf"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)

tokenizer = AutoTokenizer.from_pretrained(model_name)

这里有个细节：device_map="auto"会自动把模型分配到多个GPU上。如果你只有单卡，它会全部放在第一张卡上。我曾经在4卡机器上测试过，auto模式会把不同层分配到不同GPU，显存利用率很高。

量化前后的显存对比

为什么大家这么热衷于4-bit量化？说白了就是为了省钱。我们来看一组实际数据：

模型精度	LLaMA-7B显存占用	LLaMA-13B显存占用	LLaMA-30B显存占用
FP16（原始）	~14GB	~26GB	~60GB
8-bit	~8GB	~14GB	~32GB
4-bit（NF4）	~4.5GB	~8GB	~17GB

看到没？7B模型从14GB降到4.5GB，一张RTX 4090就能跑。30B模型从60GB降到17GB，两张3090也够了。这省下来的可都是真金白银。

推理测试

模型加载完了，我们来跑个推理看看效果：

input_text = "Explain the concept of quantum computing in simple terms."
inputs = tokenizer(input_text, return_tensors="pt").to("cuda")

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

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response)

跑起来的感觉跟FP16模型几乎没区别。唯一要注意的是，第一次推理会慢一点，因为bitsandbytes需要做一些初始化工作。后面就快了。

小技巧：如果你发现推理速度比预期慢，可以试试把bnb_4bit_compute_dtype改成torch.bfloat16。在A100等支持bfloat16的卡上，计算速度会更快。

量化模型的保存与加载

模型量化好了，总不能每次都重新量化吧？我们可以把它保存下来：

# 保存量化后的模型
model.save_pretrained("./llama-7b-4bit")
tokenizer.save_pretrained("./llama-7b-4bit")

# 下次加载时直接使用
model = AutoModelForCausalLM.from_pretrained(
    "./llama-7b-4bit",
    device_map="auto"
)

这里有个坑：保存的模型文件里包含的是量化后的权重和缩放因子，不是原始的FP16权重。所以文件体积会小很多。我曾经犯过一个错误，以为保存后还能转回FP16，结果发现不行——量化是有损的，信息已经丢了。

精度评估

量化完总得看看效果吧？我一般用困惑度（Perplexity）来评估：

from datasets import load_dataset
import math

# 加载测试数据
dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
text = " ".join(dataset["text"][:100])

# 计算困惑度
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to("cuda")
with torch.no_grad():
    outputs = model(**inputs, labels=inputs["input_ids"])
    loss = outputs.loss
    perplexity = math.exp(loss.item())

print(f"Perplexity: {perplexity:.2f}")

以LLaMA-7B为例，FP16的困惑度大概在5.8左右，NF4量化后大约6.1。只差了0.3个点，但显存省了三分之二。这个trade-off我觉得非常划算。

常见问题与避坑

最后分享几个我实际项目中遇到的问题：

OOM问题：如果显存不够，可以试试load_in_8bit或者减小max_length。我曾经在16GB显存的卡上跑30B模型，4-bit量化后还是OOM，后来发现是max_length设太大了。
速度慢：检查一下CUDA版本和bitsandbytes版本是否匹配。不匹配的话，有些算子会回退到CPU实现，那速度就惨不忍睹了。
精度异常：如果量化后模型输出完全不对，先检查tokenizer是否加载正确。我遇到过tokenizer和模型不匹配的情况，输出全是乱码。

重要提醒：4-bit量化后的模型不适合做微调。如果你想微调，建议用8-bit或者直接用FP16。4-bit的梯度精度太低，训练容易发散。

本章小结

4-bit量化说白了就是用一点点精度损失，换取巨大的显存节省。bitsandbytes让这个过程变得非常简单，几行代码就能搞定。我个人建议，如果你只是做推理部署，4-bit量化是性价比最高的选择。如果要做训练或微调，还是老老实实用FP16吧。

📚 课程目录

量化实战：使用bitsandbytes库对LLaMA模型进行4-bit量化

为什么选bitsandbytes？

环境准备

加载4-bit量化模型

量化前后的显存对比

推理测试

量化模型的保存与加载

精度评估

常见问题与避坑

本章小结