模型量化

bitsandbytes 集成

Accelerate 将 bitsandbytes 量化引入您的模型。您现在只需几行代码即可加载任何 8 位或 4 位的 PyTorch 模型。

如果您想使用带有 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")

最后，您需要使用 BnbQuantizationConfig 设置您的量化配置。

以下是 8 位量化的示例

from accelerate.utils import BnbQuantizationConfig
bnb_quantization_config = BnbQuantizationConfig(load_in_8bit=True, llm_int8_threshold = 6)

以下是 4 位量化的示例

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")

要使用所选配置量化您的空模型，您需要使用 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 中的 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 演示，了解如何在 GTP2 模型上运行量化模型。GPT2-1.5B 模型检查点为 FP32，使用 6GB 内存。量化后，使用 8 位模块时使用 1.6GB，使用 4 位模块时使用 1.2GB。

< > 更新在 GitHub 上

加速