bitsandbytes: 让大型语言模型更易接近的量化工具

bitsandbytes

bitsandbytes: 让大型语言模型触手可及

在人工智能和深度学习领域,大型语言模型(Large Language Models, LLMs)的出现和发展引起了广泛关注。这些模型在自然语言处理任务中展现出惊人的能力,但同时也带来了巨大的计算和存储需求。为了让更多研究者和开发者能够使用这些强大的模型,bitsandbytes应运而生。

bitsandbytes简介

bitsandbytes是一个轻量级的Python库,旨在通过k位量化技术为PyTorch提供访问大型语言模型的便捷方式。它主要包含以下几个方面的功能:

8位优化器:可以将优化器的内存占用降低75%。
矩阵乘法(LLM.int8()):支持8位整数矩阵乘法,大幅减少内存使用。
量化函数:提供8位和4位量化操作。

通过使用bitsandbytes,开发者可以在有限的硬件资源下运行和微调大型语言模型,极大地扩展了这些模型的应用场景。

bitsandbytes logo

主要功能

1. 8位优化器

bitsandbytes提供了多种8位优化器,包括Adam、AdamW、RMSProp、LARS和LAMB等。使用这些优化器可以将参数存储所需的内存降低75%,同时保持与32位优化器相当的性能。

使用方法非常简单,只需要将原来的优化器替换为bitsandbytes中对应的8位版本即可:

import bitsandbytes as bnb

# 原来的优化器
# optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 替换为8位优化器
optimizer = bnb.optim.Adam8bit(model.parameters(), lr=0.001)

2. 8位矩阵乘法(LLM.int8())

LLM.int8()是bitsandbytes的一个重要特性,它允许在8位精度下进行矩阵乘法运算。这种方法可以显著减少内存使用,同时保持计算精度。

使用LLM.int8()非常直接:

import bitsandbytes as bnb

# 创建8位线性层
linear = bnb.nn.Linear8bitLt(input_size, output_size, bias=True, has_fp16_weights=False, threshold=6.0)

# 使用8位矩阵乘法
output = linear(input.to(torch.float16))

3. 量化函数

bitsandbytes提供了多种量化方法,包括线性量化、动态量化和分位数量化等。这些函数可以将模型权重和激活值量化为8位或4位,从而大幅减少模型大小和内存占用。

import bitsandbytes as bnb

# 8位量化
quantized_weights = bnb.nn.Params8bit(original_weights)

# 4位量化
quantized_weights = bnb.nn.Params4bit(original_weights)

安装和使用

bitsandbytes的安装非常简单,可以通过pip直接安装:

pip install bitsandbytes

需要注意的是,bitsandbytes目前主要支持Linux系统和NVIDIA GPU。对于Windows用户,可能需要一些额外的步骤来解决兼容性问题。

安装完成后,就可以在Python代码中导入和使用bitsandbytes了:

import bitsandbytes as bnb

# 使用8位优化器
optimizer = bnb.optim.Adam8bit(model.parameters(), lr=0.001)

# 使用8位线性层
linear = bnb.nn.Linear8bitLt(input_size, output_size)

# 使用量化函数
quantized_weights = bnb.nn.Params8bit(original_weights)

与其他框架的集成

bitsandbytes可以与多个流行的深度学习框架和工具集成,特别是在自然语言处理领域。例如,它可以与Hugging Face的Transformers库无缝集成,让用户能够轻松地对大型语言模型进行量化和微调。

from transformers import AutoModelForCausalLM
import bitsandbytes as bnb

# 加载并量化模型
model = AutoModelForCausalLM.from_pretrained("gpt2", 
                                             load_in_8bit=True,
                                             device_map="auto")

# 使用8位优化器进行微调
optimizer = bnb.optim.Adam8bit(model.parameters(), lr=5e-5)