模型量化
bitsandbytes
集成
Accelerate 将 bitsandbytes
量化引入到你的模型中。现在,你可以用几行代码将任何 PyTorch 模型加载为 8 位或 4 位。
如果你希望使用 bitsandbytes
与 Transformers 模型,你应该参考此 文档。
要了解 bitsandbytes
量化的工作原理,可以查看关于 8 位量化 和 4 位量化 的博客文章。
先决条件
你需要安装以下依赖:
- 安装
bitsandbytes
库
pip install bitsandbytes
- 从源代码安装最新版
accelerate
pip install git+https://github.com/huggingface/accelerate.git
- 安装
minGPT
和huggingface_hub
以运行示例
git clone https://github.com/karpathy/minGPT.git
pip install minGPT/
pip install huggingface_hub
工作原理
首先,我们需要初始化模型。为了节省内存,我们可以使用上下文管理器 [init_empty_weights
] 来初始化一个空模型。
让我们以 minGPT 库中的 GPT2 模型为例。
from accelerate import init_empty_weights
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)
然后,我们需要获取模型权重的路径。路径可以是 state_dict 文件(例如 "pytorch_model.bin")或包含分片检查点的文件夹。
from huggingface_hub import snapshot_download
weights_location = snapshot_download(repo_id="marcsun13/gpt2-xl-linear-sharded")
最后,你需要使用 [~utils.BnbQuantizationConfig
] 设置量化配置。
这是一个 8 位量化的示例:
from accelerate.utils import BnbQuantizationConfig
bnb_quantization_config = BnbQuantizationConfig(load_in_8bit=True, llm_int8_threshold = 6)
这是一个 4 位量化(4-bit quantization)的例子:
from accelerate.utils import BnbQuantizationConfig
bnb_quantization_config = BnbQuantizationConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4")
要使用所选配置量化你的空模型,你需要使用 [~utils.load_and_quantize_model
]。
from accelerate.utils import load_and_quantize_model
quantized_model = load_and_quantize_model(empty_model, weights_location=weights_location, bnb_quantization_config=bnb_quantization_config, device_map = "auto")
保存和加载 8 位模型
你可以使用 accelerate
中的 [~Accelerator.save_model
] 保存你的 8 位模型。
from accelerate import Accelerator
accelerate = Accelerator()
new_weights_location = "path/to/save_directory"
accelerate.save_model(quantized_model, new_weights_location)
quantized_model_from_saved = load_and_quantize_model(empty_model, weights_location=new_weights_location, bnb_quantization_config=bnb_quantization_config, device_map = "auto")
请注意,4 位模型序列化目前不受支持。
将模块卸载到 CPU 和磁盘
如果你的 GPU 没有足够的空间来存储整个模型,可以将一些模块卸载到 CPU 或磁盘。 这在底层使用了大模型推理。更多详细信息请参阅此文档。
对于 8 位量化,选定的模块将被转换为 8 位精度。
对于 4 位量化,选定的模块将保持在用户在 BnbQuantizationConfig
中传递的 torch_dtype
中。我们将在 4 位序列化成为可能时添加支持,以将这些卸载的模块转换为 4 位。
你只需要传递一个自定义的 device_map
以将模块卸载到 CPU 或磁盘。当需要时,卸载的模块将被调度到 GPU。以下是一个示例:
device_map = {
"transformer.wte": 0,
"transformer.wpe": 0,
"transformer.drop": 0,
"transformer.h": "cpu",
"transformer.ln_f": "disk",
"lm_head": "disk",
}
微调量化模型
无法对这些模型进行纯8位或4位训练。但是,你可以通过利用参数高效的微调方法(PEFT)来训练这些模型,例如在它们之上训练适配器。请参阅 peft 库以获取更多详细信息。
目前,你不能在任何量化模型上添加适配器。但是,随着适配器对 Transformers 模型的官方支持,你可以微调量化模型。如果你想要微调一个 Transformers 模型,请参阅此 文档。查看此 演示,了解如何微调一个4位的 Transformers 模型。
请注意,在加载模型进行训练时,你不需要传递 device_map
。它会自动将模型加载到你的 GPU 上。请注意,device_map=auto
仅应用于推理。
示例演示 - 在 Google Colab 上运行 GPT2 1.5B
查看此 Google Colab 演示,了解如何在 GPT2 模型上运行量化模型。GPT2-1.5B 模型检查点为 FP32,占用 6GB 内存。量化后,使用 8 位模块占用 1.6GB,使用 4 位模块占用 1.2GB。