FSDP-QLoRA

FSDP-QLoRA 结合了数据并行（FSDP 允许跨 GPU 分片模型参数、优化器状态和梯度）、4 位量化和 LoRA，以便在双 24GB GPU 系统上训练高达 700 亿参数的 LLM。这项技术由 Answer.AI 与 bitsandbytes 合作发布，旨在使 LLM 训练更高效且人人可用。

本指南简要介绍了 bitsandbytes 如何支持存储量化权重以启用 FSDP-QLoRA，以及如何使用 Hugging Face 库运行训练。

bitsandbytes 支持 FSDP-QLoRA 所需的其他更改，例如从量化元数据重建权重以及防止在量化权重从 CPU 移动到 GPU 时再次量化，已在此Pull Request中记录，并在在消费级 GPU 上实现 700 亿参数微调博客文章中描述。我们强烈建议阅读这些资源，以便更好地理解 FSDP-QLoRA！

量化数据存储

FSDP 仅支持分片浮点数据类型，这可能会有问题，因为量化权重通常存储为整数数据类型 (uint8)。bitsandbytes 没有这个问题，因为它使用 StoreChar 读取和写入量化权重，而无需考虑数据类型存储。这使得向 Linear4bit 和 Params4bit 类添加 quant_storage 参数并将其设置为 torch.uint8 以保持与代码库的向后兼容性变得简单。使用 quant_storage 参数，您可以选择任何 FSDP 支持的数据类型来分片 Linear4bit，例如 bfloat16、float16 或 float32。

您通常会从 transformers.BitsAndBytesConfig 访问和配置此选项，方法是设置 bnb_4bit_quant_storage 参数。非常**重要**的是，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 更易于访问博客文章。

< > 在 GitHub 上更新

Bitsandbytes

FSDP-QLoRA

量化数据存储

训练

资源