【Python】bitsandbytes库：PyTorch低精度（k-bit）优化与量化功能详解

bitsandbytes 是一个轻量级的 Python 库，专注于为 PyTorch 提供低精度（k-bit）优化和量化功能，主要用于降低大型语言模型（LLM）在训练和推理过程中的内存消耗，同时保持性能。它通过 CUDA 自定义函数实现 8 位和 4 位量化、优化器和矩阵乘法操作，特别适合在 GPU 上运行深度学习任务。以下是对 bitsandbytes 库的详细介绍，包括其功能、主要模块、用法和注意事项。

1. 核心功能

bitsandbytes 提供以下三大主要功能，用于大幅减少内存占用并加速计算：

此外，库还支持：

2. 主要模块和功能

bitsandbytes 库的 Python 接口主要包括以下模块和类：

a. bitsandbytes.optim 模块

提供 8 位优化器，替代 PyTorch 的标准优化器，支持常见的优化算法：

import bitsandbytes as bnb
model = MyModel().cuda()
# 替换标准 Adam 为 8 位 Adam
optimizer = bnb.optim.Adam8bit(model.parameters(), lr=0.001, betas=(0.9, 0.995))
# 调整小参数的 8 位阈值
optimizer = bnb.optim.Adam8bit(model.parameters(), min_8bit_size=16384)

b. bitsandbytes.nn 模块

提供量化的神经网络层，替代 PyTorch 的标准层：

import bitsandbytes as bnb
linear = bnb.nn.Linear8bitLt(dim1, dim2, bias=True, has_fp16_weights=False, threshold=6.0)
out = linear(x.to(torch.float16))  # 输入需为 FP16

import torch.nn as nn
from bitsandbytes.nn import Embedding4bit
fp16_module = nn.Embedding(128, 64)
quantized_module = Embedding4bit(128, 64)
quantized_module.load_state_dict(fp16_module.state_dict())
quantized_module = quantized_module.to(0)  # 移动到 GPU

c. bitsandbytes.functional 模块（部分功能已废弃）

提供底层的量化和矩阵乘法函数，例如：

d. bitsandbytes.optim.GlobalOptimManager

用于混合精度优化，允许部分参数使用 32 位优化器，其他参数使用 8 位优化器：

import torch
import bitsandbytes as bnb
model = MyModel()
mng = bnb.optim.GlobalOptimManager.get_instance()
mng.register_parameters(model.parameters())
model = model.cuda()
optimizer = bnb.optim.Adam(model.parameters(), lr=0.001, optim_bits=8)
mng.override_config(model.fc1.weight, 'optim_bits', 32)  # 指定 fc1.weight 使用 32 位

3. 安装和环境要求

pip install bitsandbytes

pip install bitsandbytes-cudaXXX  # XXX 为 110、111、113 等

python -m bitsandbytes

wget https://gist.githubusercontent.com/TimDettmers/1f5188c6ee6ed69d211b7fe4e381e713/raw/check_bnb_install.py
python check_bnb_install.py

git clone https://github.com/bitsandbytes-foundation/bitsandbytes.git
cd bitsandbytes
CUDA_VERSION=117 make cuda11x
python setup.py install

4. 典型用法

以下是一些常见的使用场景和代码示例：

a. 加载 4 位量化模型（QLoRA）

结合 Hugging Face 的 transformers 库加载 4 位模型：

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
model_name = "facebook/opt-350m"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,  # 双重量化，节省额外 0.4 位/参数
    bnb_4bit_compute_dtype=torch.float16  # 计算使用 FP16
)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto"
)

b. 使用 8 位优化器训练

替换标准优化器为 8 位优化器：

import bitsandbytes as bnb
model = MyModel().cuda()
optimizer = bnb.optim.Adam8bit(model.parameters(), lr=0.001)
for data in dataloader:
    optimizer.zero_grad()
    outputs = model(data)
    loss = outputs.loss
    loss.backward()
    optimizer.step()

c. LLM.int8() 推理

使用 8 位线性层进行高效推理：

import torch
import bitsandbytes as bnb
linear = bnb.nn.Linear8bitLt(64, 64, bias=True, has_fp16_weights=False, threshold=6.0)
x = torch.randn(32, 64).to(torch.float16).cuda()
out = linear(x)

d. 结合 PEFT 进行 QLoRA 微调

结合 Hugging Face 的 PEFT 库进行 4 位模型微调：

from transformers import AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model
model_name = "facebook/opt-350m"
bnb_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16)
model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=bnb_config, device_map="auto")
lora_config = LoraConfig(r=8, lora_alpha=16, target_modules=["q_proj", "v_proj"])
model = get_peft_model(model, lora_config)
# 进行微调

5. 注意事项

pip install -U bitsandbytes transformers accelerate

6. 许可和引用

7. 最新动态

8. 参考资源

9. 常见问题

作者：彬彬侠