概述
bitsandbytes 是一个基于 PyTorch 的轻量级库,专注于通过量化技术和低精度计算来优化大规模模型的训练和推理效率。它通过自定义 CUDA 函数实现了高效的 8-bit 和 4-bit 量化操作,显著降低了模型的内存占用和计算成本,使得在资源受限的设备(如消费级 GPU)上运行大语言模型(LLM)成为可能。
核心功能
- 8-bit 量化:
- 支持将模型权重和激活值压缩为 8-bit 整数,减少内存占用约 50%。
- 提供 LLM.int8() 推理模式,允许在 8-bit 精度下高效运行大模型(如 BERT、GPT 等)。
- 内置 8-bit 优化器(如 AdamW、Lion),节省训练时的显存消耗(可减少 75% 以上)。
- 4-bit 量化(FP4/NF4):
- 支持更激进的 4-bit 量化,进一步压缩模型体积(4-bit 比 8-bit 再减少 50% 内存)。
- 提供动态量化和**双量化(Double Quantization)**技术,在低比特下保持较高精度。
- 与 Hugging Face Transformers 深度集成,可一键加载 4-bit 模型(如
load_in_4bit=True)。
- 混合精度训练:
- 支持混合精度训练(FP16/BF16 + 8-bit/4-bit),平衡计算效率和数值稳定性。
- 提供高效的嵌入层(Embedding Layer)优化,减少训练中的内存波动。
- 零样本量化(Zero-shot Quantization):
- 不需要额外的校准数据即可完成量化,简化部署流程。
- 适用于任何包含
torch.nn.Linear层的模型(如 Whisper、ViT、Blip2 等)。
主要作用
- 降低显存占用:
- 通过 8-bit 或 4-bit 量化,显存需求可减少 50%~75%,使得大模型(如 10B 参数以上)在消费级 GPU 上也能运行。
- 例如,Bloomz-3B 模型在 8-bit 量化后,显存占用从 12GB 降至 6GB。
- 加速推理和训练:
- 低精度计算(如 8-bit/4-bit)比全精度(FP32)更快,尤其在支持 Tensor Core 的 GPU 上效果显著。
- 优化器(如 AdamW 8-bit)减少训练时的显存消耗,提升训练吞吐量。
- 兼容性和易用性:
- 与 Hugging Face Transformers 原生集成,只需简单配置即可启用量化(无需修改模型代码)。
- 支持 PyTorch 的
device_map="auto"功能,自动分配模型到多卡或 CPU。
- 跨模态支持:
- 不仅适用于 NLP 模型(如 GPT、BERT),还支持语音(Whisper)、视觉(ViT)、多模态(Blip2)等模型的量化。
量化降低了哪些精度
- 权重(Weights):
- 模型的参数(如神经网络中的权重矩阵)会被压缩为低精度(如 8-bit 或 4-bit 整数),从而减少存储空间。
- 例如,在 8-bit 量化中,权重从 32-bit 浮点数(FP32)或 16-bit 浮点数(FP16)转换为 8-bit 整数(INT8)。
- 异常值处理:bitsandbytes 会保留部分异常值(如超出正常分布范围的权重)以 FP16 精度计算,避免精度损失过大。
- 激活值(Activations):
- 在推理或训练过程中,中间激活值(如神经元的输出)也会被动态量化为低精度(如 8-bit 或 4-bit)。
- 激活值的量化通常是动态的(运行时根据输入范围调整量化参数),而非静态校准。
- 优化器状态(Optimizers):
- bitsandbytes 提供 8-bit 优化器(如 AdamW 8-bit),将优化器的状态(如动量、梯度)也压缩为 8-bit,显著减少训练时的显存占用。
量化对模型性能的影响
- 精度损失(Accuracy):
- 8-bit 量化:通常对模型精度影响较小,尤其是对于大型语言模型(LLM),精度下降可忽略不计(如 BERT 的精度下降 < 0.5%)。
- 4-bit 量化:精度损失可能稍大(如 1-2%),但通过双量化(Double Quantization)和异常值处理可缓解(例如,将部分权重保留为 FP16 计算)。
- 与 HQQ/EETQ 的对比:bitsandbytes 的 4-bit 量化精度略低于 HQQ(HQQ 的 4-bit 精度更高),但 bitsandbytes 更注重易用性和与 Hugging Face 的兼容性。
- 推理速度提升:
- 低精度计算(如 8-bit/4-bit)比全精度(FP32/FP16)更快,尤其在支持 Tensor Core 的 GPU 上效果显著。
- 例如,8-bit 量化可使 LLM 推理速度提升 2-3 倍。
- 显存占用减少:
- 8-bit 量化可减少显存占用约 50%,4-bit 量化可减少 75% 以上(如 BLOOMZ-3B 在 4-bit 量化后显存需求从 12GB 降至 3GB)。
量化对模型训练的影响
- 支持训练的量化类型:
- 8-bit 量化:bitsandbytes 提供 8-bit 优化器(如 AdamW 8-bit),允许在 8-bit 精度下进行训练,显著减少显存占用。
- 4-bit 量化:目前主要支持推理,但可通过 QLoRA(Quantized Low-Rank Adaptation)进行微调(例如,冻结量化后的主模型,仅训练低秩适配层)。
- 量化感知训练(QAT)的对比:
- bitsandbytes 的量化是后训练量化(Post-Training Quantization, PTQ),即在训练完成后进行量化。
- 如果需要更高的精度,可结合量化感知训练(Quantization-Aware Training, QAT)(如 PyTorch 的
torch.quantization),在训练过程中模拟量化噪声,使模型适应低精度计算。
- 注意事项:
- 动态量化:bitsandbytes 的激活值量化是动态的(运行时根据输入范围调整),因此在训练时需确保量化参数的稳定性。
- 硬件兼容性:4-bit 量化对硬件要求较高,需使用支持低精度计算的 GPU(如 NVIDIA A100/H100)。
总结
| 特性 | bitsandbytes 量化 |
|---|---|
| 量化目标 | 权重、激活值、优化器状态 |
| 精度损失 | 8-bit 影响小,4-bit 约 1-2%(可通过双量化和异常值处理缓解) |
| 性能提升 | 显存减少 50-75%,推理速度提升 2-3 倍 |
| 是否支持训练 | 支持(8-bit 优化器);4-bit 量化可通过 QLoRA 微调 |
| 适用场景 | 大模型部署(如 LLM)、资源受限设备、快速原型开发 |
示例代码:加载 4-bit 量化模型
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
# 配置 4-bit 量化
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4", # 使用 NF4 量化
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=torch.bfloat16
)
# 加载量化模型
model = AutoModelForCausalLM.from_pretrained(
"bigscience/bloomz-3b",
quantization_config=quantization_config
)
安装
# 前置依赖
pip install "transformers>=4.45.1"
# pip安装
pip install bitsandbytes
8 位优化器
使用 8 位优化器,与使用标准 32 位优化器进行训练相比,可以在不损失任何精度的前提下,将大型模型的微调 GPU 内存需求减少 75%。
减少的内存需求意味着 8 位优化器比标准优化器快 4 倍,且无需进行超参数调整。
8 位优化器是常规优化器的直接替代品,这意味着它们也接受与常规优化器相同的参数。对于 NLP 模型,建议使用 StableEmbedding 类以提高稳定性和结果。
import bitsandbytes as bnb
# 使用8位优化器替代常规优化器
# adam = torch.optim.Adam(...) # 替换前
adam = bnb.optim.Adam8bit(...) # 替换后
# NLP models推荐使用StableEmbedding
# torch.nn.Embedding(...) # 替换前
# bnb.nn.StableEmbedding(...) # 替换后
默认情况下,即使使用 8 位优化器,所有元素数量少于 4096 的参数张量都保持在 32 位。这是因为小张量量化不会节省太多内存,并且它们通常包含高度可变的参数(偏差)或需要高精度的参数(批量归一化、层归一化)。
可以通过 min_8bit_size 参数进行调整。例如,如果想在最小值达到 16384 个值时,才将参数优化为 8 位(建议使用 4096 的倍数):
import bitsandbytes as bnb
adam = bnb.optim.Adam8bit(model.parameters(), min_8bit_size=16384)
可以配置的其他参数包括学习率(lr)、衰减率(betas)、优化器状态的位数(optim_bits)以及百分位数裁剪(percentile_clipping),这些可以提高稳定性。
例如:初始化一个带有 5 百分位阈值裁剪极端梯度值的 32 位 Adam 优化器:
import bitsandbytes as bnb
adam = bnb.optim.Adam(
model.parameters(),
lr=0.001,
betas=(0.9, 0.995),
optim_bits=32,
percentile_clipping=5
)
参数说明:
1. lr=0.001
- 含义:学习率(Learning Rate),控制参数更新的步长大小。
- 作用:较大的学习率可能导致训练不稳定,而较小的学习率会减缓收敛速度。默认值通常设为 0.001。
2. betas=(0.9, 0.995)
- 含义:两个动量参数,
(beta1, beta2),用于计算梯度的一阶矩估计(均值)和二阶矩估计(未中心化的方差)。beta1=0.9:控制一阶矩估计的指数衰减率。beta2=0.995:控制二阶矩估计的指数衰减率。
- 作用:
beta1和beta2的值越接近 1,历史梯度的权重越大,优化器对梯度变化的响应越平滑。
3. optim_bits=32
- 含义:指定优化器状态(如动量、梯度平方等)的存储精度(位数)。
- 作用:
optim_bits=32表示使用 32-bit 浮点数存储优化器状态(即全精度)。- 在 bitsandbytes 中,默认的 Adam8bit 会将优化器状态压缩为 8-bit 整数,以减少显存占用。但通过设置
optim_bits=32,可以强制使用全精度,从而提升数值稳定性(代价是增加显存消耗)。默认使用 8-bit(如 Adam8bit)。
4. percentile_clipping=5
- 含义:梯度裁剪的百分位阈值,用于限制梯度的范围。
- 作用:
- 梯度裁剪(Gradient Clipping)是防止梯度爆炸的常用技术。
percentile_clipping=5表示将梯度限制在分布的 5% 和 95% 百分位之间,即裁剪掉极端值。 - 这种方法比传统的全局梯度裁剪(如
clipnorm或clipvalue)更温和,仅对分布尾部的异常值进行限制。 - 这是 bitsandbytes 特有的功能,专门针对低精度优化器设计,通过动态百分位裁剪增强训练稳定性。
- 例如,在 8-bit 量化训练中,梯度可能因低精度计算而出现噪声或异常值,
percentile_clipping可以缓解这一问题。
- 梯度裁剪(Gradient Clipping)是防止梯度爆炸的常用技术。
优化不稳定参数:使用 32 位 Adam 优化某些不稳定参数,使用 8 位 Adam 优化其他参数
a. 当参数位于 CPU 上时进行注册
import torch
import bitsandbytes as bnb
mng = bnb.optim.GlobalOptimManager.get_instance()
model = MyModel()
mng.register_parameters(model.parameters())
b. 使用新的期望超参数覆盖配置。 例如,覆盖 model.fc1.weight 层以使用 32 位 Adam:
model = model.cuda()
# 对所有参数使用8位进行初始化
adam = bnb.optim.Adam(model.parameters(), lr=0.001, optim_bits=8)
# model.fc1.weight的参数使用32位初始化
mng.override_config(model.fc1.weight, "optim_bits", 32)
# override_config可以一次性覆盖多个层的精度,
# 通过将指定模型层的精度放在一个列表中,优化器参数放在一个字典中,一次性覆盖多个层的精度
# 例如,让我们将 model.special.weight 和 model.also_special.weight 层覆盖为使用稀疏优化和较低的学习率及衰减率
mng.override_config(
[model.special.weight, model.also_special.weight],
key_value_dict={'is_sparse': True, 'lr': 1e-5, 'betas': (0.9, 0.98)}
)
FSDP-QLoRA
FSDP-QLoRA 结合了数据并行(FSDP 能够在 GPU 之间分片模型参数、优化器状态和梯度)、4 位量化以及 LoRA 技术,可以在两张 24GB GPU 上训练高达 70B 参数的 LLMs。
量化数据存储
FSDP 仅支持分片浮点数据类型,但是量化权重通常存储为整数数据类型(uint8)。
为了解决数据类型不兼容的问题,bitsandbytes 使用 StoreChar 对任意数据类型读取和写入量化权重。通过在 Linear4bit 和 Params4bit 类中添加 quant_storage 参数并设置为 torch.uint8 以保持与代码库的向后兼容。
通过 quant_storage 参数,可以选择 FSDP 支持的任何数据类型来分片 Linear4bit,例如 bfloat16、float16 或 float32。
from transformers import BitsAndBytesConfig, AutoModelForCausalLM
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_quant_storage=torch.bfloat16,
)
model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-2-70b",
quantization_config=bnb_config,
torch_dtype=torch.bfloat16,
)
训练
bitsandbytes 与 Hugging Face 生态系统深度集成,使其能够轻松与 Transformers、PEFT 和 TRL 等库配合使用。
PEFT 提供了一个配置文件(fsdp_config_qlora.yaml)、一个启动命令(run_peft_qlora_fsdp.sh)以及一个训练脚本(train.py),用于运行 FSDP-QLoRA。
安装依赖:
pip install -U bitsandbytes accelerate transformers peft trl
启用 FSDP-QLoRA 训练: 调整 BitsAndBytesConfig 类中的 bnb_4bit_quant_storage 参数,将量化权重的存储数据类型设置为浮点数据类型。
from transformers import BitsAndBytesConfig
bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_storage=torch.bfloat16,
)
将 BitsAndBytesConfig 传递给模型。
注意: 将
torch_dtype参数设置为与bnb_4bit_quant_storage匹配,以便Linear4bit层与Linear层包装在一起。如果存储类型不匹配,则每个Linear4bit层将单独包装。
from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained(
"meta-llama/Llama-2-70b",
quantization_config=bnb_config,
torch_dtype=torch.bfloat16,
)
设置 target_modules="all-linear" 配置 peft.LoraConfig 类以进行 QLoRA 训练:
from peft import LoraConfig
peft_config = LoraConfig(
lora_alpha=16,
lora_dropout=0.1,
r=64,
bias="none",
task_type="CAUSAL_LM",
target_modules="all-linear",
)
通过 SFTTrainer 进行训练:
from trl import SFTTrainer
trainer = SFTTrainer(
model=model,
train_dataset=dataset,
peft_config=peft_config,
processing_class=tokenizer,
args=training_arguments,
)
trainer.train()
HuggingFace 框架集成
Transformers
使用 Transformers,可以以 4 位或 8 位加载任何模型,并在运行时进行量化。
示例:加载模型并量化为 4 位,使用 bfloat16 数据类型进行推理
from transformers import AutoModelForCausalLM, BitsAndBytesConfig
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16
)
model_4bit = AutoModelForCausalLM.from_pretrained(
"bigscience/bloom-1b7",
device_map=device_map,
quantization_config=quantization_config,
)
说明:
bnb_4bit_compute_dtype参数默认为float32,float32可以确保向后兼容、数值稳定,但是占用空间大、计算速度慢,float16可以更小更快,但是数值不稳定。bfloat16结合了float32和float16的最佳特性:提供了float32的数值稳定性和float16的空间小计算快。- 在硬件支持的前提下,建议将
bnb_4bit_compute_dtype指定为torch.bfloat16。
指定优化器类型
初始化 Trainer 类时可以通过 optim 指定优化器类型,optim 支持的优化器类型如下:
| 优化器名称 | 描述 |
|---|---|
adamw_hf | Hugging Face 默认的 AdamW 实现(基于 PyTorch)。 |
adamw_torch | PyTorch 原生的 AdamW 优化器。 |
paged_adamw_32bit | 使用 32-bit 精度的 Paged AdamW(分页存储优化器状态,节省显存)。 |
paged_adamw_8bit | 8-bit 量化版本的 Paged AdamW(结合 bitsandbytes 实现低精度优化器状态)。 |
adafactor | Adafactor 优化器(适合大模型训练,内存效率高)。 |
adamw_apex_fused | 使用 NVIDIA Apex 的融合 AdamW(需安装 Apex 库)。 |
adamw_torch_fused | PyTorch 的融合 AdamW(需 PyTorch 2.0+)。 |
adamw_torch_xla | 在 TPU 上使用的 AdamW(需 XLA 支持)。 |
from transformers import TrainingArguments, Trainer
training_args = TrainingArguments(
output_dir="./results",
optim="paged_adamw_32bit", # 使用 32-bit Paged AdamW
learning_rate=2e-5,
per_device_train_batch_size=16,
num_train_epochs=3,
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
)
trainer.train()
说明: optim 参数涉及的分页存储优化器状态,是指按需加载优化器状态到显存的技术,以减少显存峰值占用。AdamW 优化器状态包含三部分:
- 动量(
exp_avg):梯度的移动平均值 - 梯度平方(
exp_avg_sq):梯度平方的移动平均值 - 偏差校正项(
step):用于校正初始偏置的计数器
这些状态通常与模型参数一一对应,且每个参数的优化器状态都需要额外的显存存储。分页存储的核心思想是只将当前计算所需的优化器状态保留在 GPU 显存中,其余状态暂时存储在 CPU 内存中,并按需加载。
分页存储的工作原理:
a. 状态分页:
- 分页粒度:将优化器状态划分为固定大小的”页”,例如:128MB/页
- 存储位置:
- GPU 显存:缓存当前需要计算的页,例如最近使用的页
- CPU 内存:存储剩余的页
b. 缓存管理:
- 缓存命中(Cache Hit):如果当前计算所需的优化器状态页已经在 GPU 显存中,直接使用
- 缓存未命中(Cache Miss):
- 换出(Page Out):从 GPU 显存中选择一个不常用的页(例如使用 LRU 策略),将其写回 CPU 内存
- 换入(Page In):从 CPU 内存中加载所需页到 GPU 显存
- 异步传输:利用 CUDA 流实现换入/换出操作与 GPU 计算的并行执行,减少等待时间
c. 异步 IO 优化:
- Pinned Memory:使用页锁定内存(Pinned Memory)加速 CPU 与 GPU 之间的数据传输
- 重叠计算与传输:在 GPU 计算当前页的同时,后台异步加载/写回其他页,最大化硬件利用率
分页存储的效果: 只保留当前需要的优化器页,显存占用减少 50%~90%,训练速度略微下降(增加 5%~20% 训练时间)。
optim 与 BitsAndBytesConfig 的区别
| 特性 | optim 参数 | BitsAndBytesConfig 量化策略 |
|---|---|---|
| 作用对象 | 优化器状态(动量、梯度平方等) | 模型权重、激活值、优化器状态 |
| 目标 | 控制参数更新方式(如 AdamW、SGD) | 降低显存占用(通过低精度量化) |
| 精度影响 | 控制优化器计算的精度(如 32-bit、8-bit) | 量化模型参数和计算(如 8-bit、4-bit) |
| 显存优化 | 使用分页技术(Paged AdamW)减少峰值显存 | 量化压缩模型参数和优化器状态(如 8-bit AdamW) |
| 兼容性 | 与 Hugging Face 原生模型和训练流程集成 | 需结合 bitsandbytes 库和量化配置 |
关键差异示例:
optim="paged_adamw_32bit":使用 32-bit 精度的 Paged AdamW 优化器,优化器状态以分页方式存储,减少显存峰值。BitsAndBytesConfig(load_in_8bit=True):将模型权重量化为 8-bit,同时使用 8-bit AdamW 优化器(optim="paged_adamw_8bit")。
optim 与 BitsAndBytesConfig 结合使用
两者关系:
optim参数:直接控制优化器的类型和精度(如paged_adamw_32bit、paged_adamw_8bit)。BitsAndBytesConfig:控制模型参数的量化策略(如load_in_8bit、load_in_4bit)和计算时的数据类型(如bnb_4bit_compute_dtype)。
优先级规则:
- 如果
optim和BitsAndBytesConfig的配置不冲突,两者会协同工作。 - 如果配置冲突(例如
optim="paged_adamw_8bit"但BitsAndBytesConfig(load_in_4bit=True)),则以optim的配置为准。
协同工作示例:
场景 1:8-bit 优化器 + 8-bit 模型量化
from transformers import TrainingArguments, BitsAndBytesConfig
# 使用 8-bit 优化器
training_args = TrainingArguments(
optim="paged_adamw_8bit", # 8-bit AdamW 优化器
...
)
# 使用 8-bit 模型量化
quantization_config = BitsAndBytesConfig(
load_in_8bit=True # 将模型权重量化为 8-bit
)
# 加载模型时应用量化配置
model = AutoModelForCausalLM.from_pretrained(
"model_name",
quantization_config=quantization_config
)
说明:
- 模型权重以 8-bit 存储(
load_in_8bit=True)。- 优化器状态(动量、梯度平方等)也以 8-bit 存储(
optim="paged_adamw_8bit")。- 显存占用大幅降低,但训练稳定性可能略低于 32-bit。
场景 2:8-bit 优化器 + 4-bit 模型量化
# 使用 8-bit 优化器
training_args = TrainingArguments(
optim="paged_adamw_8bit",
...
)
# 使用 4-bit 模型量化
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_use_double_quant=True,
bnb_4bit_compute_dtype=torch.float16
)
说明:
- 模型权重以 4-bit 存储(
load_in_4bit=True)。- 优化器状态仍以 8-bit 存储(
optim="paged_adamw_8bit")。- 计算时使用 float16 数据类型(
bnb_4bit_compute_dtype)。- 显存占用进一步降低,但需确保 4-bit 量化与 8-bit 优化器兼容。
冲突示例:
- 冲突示例 1:8-bit 优化器 vs 4-bit 模型量化:
optim="paged_adamw_8bit"需要模型权重为 8-bit,但load_in_4bit=True将模型权重量化为 4-bit。若使用load_in_8bit=True,则optim="paged_adamw_8bit"是合理的选择。 - 冲突示例 2:32-bit 优化器 vs 8-bit 模型量化:
optim="adamw_hf"(32-bit)与load_in_8bit=True可能导致显存浪费。优先使用 8-bit 优化器:将optim改为paged_adamw_8bit,以匹配模型的 8-bit 量化策略。
PEFT
PEFT 基于 bitsandbytes Transformers 集成,并在此基础上扩展,用于训练时增加几个步骤。
a. 调用 peft.prepare_model_for_kbit_training 方法来准备模型进行训练。仅适用于 Transformers 模型。
from peft import prepare_model_for_kbit_training
model_4bit = prepare_model_for_kbit_training(model_4bit)
b. 设置 peft.LoraConfig 以使用 QLoRA。
from peft import LoraConfig
config = LoraConfig(
r=16,
lora_alpha=8,
target_modules="all-linear",
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM"
)
c. 在模型和配置上调用 peft.get_peft_model 函数来创建可训练的 PeftModel。
from peft import get_peft_model
model = get_peft_model(model_4bit, config)
Accelerate
通过传递 BnbQuantizationConfig 来量化任何 PyTorch 模型,然后调用 load_and_quantize_model 函数进行量化。
from accelerate import init_empty_weights
from accelerate.utils import BnbQuantizationConfig, load_and_quantize_model
from mingpt.model import GPT
model_config = GPT.get_default_config()
model_config.model_type = 'gpt2-xl'
model_config.vocab_size = 50257
model_config.block_size = 1024
with init_empty_weights():
empty_model = GPT(model_config)
bnb_quantization_config = BnbQuantizationConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16, # optional
bnb_4bit_use_double_quant=True, # optional
bnb_4bit_quant_type="nf4" # optional
)
quantized_model = load_and_quantize_model(
empty_model,
weights_location=weights_location,
bnb_quantization_config=bnb_quantization_config,
device_map="auto"
)