Bitsandbytes 文档
FSDP-QLoRA
并获得增强的文档体验
开始使用
FSDP-QLoRA
FSDP-QLoRA 结合了数据并行(FSDP 能够跨 GPU 对模型参数、优化器状态和梯度进行分片)、4 位量化和 LoRA 技术,可在双 24GB GPU 系统上训练高达 70B 参数的 LLM。该技术由 Answer.AI 与 bitsandbytes 合作发布,旨在使训练 LLM 更高效,让每个人都能使用。
本指南简要介绍了 bitsandbytes 如何支持存储量化权重以实现 FSDP-QLoRA,以及如何使用 Hugging Face 库进行训练。
bitsandbytes 为支持 FSDP-QLoRA 所需的其他更改,例如从量化元数据中重构权重,以及防止在权重从 CPU 移至 GPU 时对已经量化的权重再次进行量化,都记录在此 Pull Request 中,并在 《在消费级 GPU 上实现 70B 模型微调》 博客文章中进行了描述。我们强烈建议阅读这些资源,以便更好地理解 FSDP-QLoRA!
量化数据存储
FSDP 仅支持对浮点数据类型进行分片,这可能会带来问题,因为量化权重通常存储为整数数据类型 (uint8)。bitsandbytes 不存在这个问题,因为它使用 StoreChar
来读取和写入量化权重,无论数据类型存储方式如何。这使得向 Linear4bit 和 Params4bit 类添加一个 quant_storage
参数并将其设置为 torch.uint8
变得很简单,从而保持与代码库的向后兼容性。通过 quant_storage
参数,您可以选择任何 FSDP 支持的数据类型来对 Linear4bit 进行分片,例如 bfloat16、float16 或 float32。
您通常会通过设置 bnb_4bit_quant_storage
参数从 transformers.BitsAndBytesConfig 访问和配置此选项。非常重要的是,quant_storage
的数据类型必须与整个模型中使用的数据类型相匹配,因为 FSDP 只能包装具有*相同浮点数据类型*的层和模块。确保数据类型一致将保证模型被正确分片。
compute_dtype
是在 CUDA 内核中用于计算的数据类型,其中 4 位量化权重会从 quant_storage
中的数据类型解包并反量化为 compute_dtype
。我们建议使用 torch.bfloat16(如果您的硬件支持)以获得更好的数值稳定性。
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,
)
请查看 PEFT 文档的此部分,以获取运行 FSDP-QLoRA 训练的配置文件和训练代码。
训练
FSDP 是一个分布式训练框架,需要使用像 Accelerate 或 torchrun 这样的库作为分布式训练作业来启动。本节中提供的启动命令使用 Accelerate 来启动训练脚本。
bitsandbytes 与 Hugging Face 生态系统深度集成,使其可以轻松地与 Transformers、PEFT 和 TRL 等库一起使用。
PEFT 提供了用于运行 FSDP-QLoRA 的配置文件(fsdp_config_qlora.yaml)、启动命令(run_peft_qlora_fsdp.sh)和训练脚本(train.py)。要了解更多信息,请查看《使用 PEFT QLoRA 和 FSDP 在多 GPU 上微调大型模型》文档。本节简要介绍了运行 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 传递给模型,以配置 FSDP-QLoRA。您应该将 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"
为 QLoRA 训练配置 ~peft.LoraConfig
类。
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()
资源
要了解更多关于 FSDP 和 QLoRA 的信息,请查看以下资源:
- AnswerDotAI/fsdp_qlora 仓库。
- Answer.AI 的介绍性博客文章 《现在你可以在家训练 70B 的语言模型了》。
- 要了解 FSDP 的入门知识,请阅读 《介绍 PyTorch 完全分片数据并行(FSDP)API》 这篇博客文章。
- 要了解更多关于 QLoRA 的细节,请查看 《通过 bitsandbytes、4 位量化和 QLoRA 让 LLM 更易于使用》 这篇博客文章。