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 更易於使用》這篇部落格文章。

< > 在 GitHub 上更新

Bitsandbytes

FSDP-QLoRA

量化資料儲存

訓練

資源