在 24GB 消費者 GPU 上使用 RLHF 微調 20B 大型語言模型

TL;DR

Hugging Face 宣佈 trl 函式庫現在可與 peft 以及 8 位元量化一起使用，讓 RLHF 微調 20 B 參數的大型語言模型在單一 24 GB 消費者 GPU 上成為可能。這使得大規模的 RL 微調變得負擔得起且易於取得，無需多 GPU 模型平行設定。

為何 RLHF 需要高效的微調

強化學習與人類回饋 (RLHF) 通常分為三個階段：(1) 針對指令的監督式微調，(2) 從人工標註訓練獎勵模型，(3) 使用獎勵模型進行基於 PPO 的 RL 微調。RL 步驟需要在每個 GPU 上放置兩個模型副本（主動模型與參考模型），對於超過 10 B 參數的模型，這很快就會超出單一設備的記憶體容量。

TRL：用於 PPO 基礎 RL 的函式庫

trl 提供了用於語言模型 PPO 訓練的高階 API。它利用 🤗 Accelerate 在單設備或分散式環境中執行。PPO 迴圈需要同時擁有一個主動模型（正在更新）與一個參考模型（保持凍結）來計算 KL 正則化的獎勵，實際上使記憶體需求翻倍。

使用 PEFT 與 8 位元量化減少記憶體占用

8 位元矩陣乘法

  • 8 位元量化 (LLM.int8()) 將權重以每參數 1 位元組儲存，與 float32 相比將模型大小縮減四倍。
  • 此方法將每個線性層分割為處理異常值的 float16 部分與大量的 int8 部分，在保持精度的同時提升速度。

透過 PEFT 的低秩適配 (LoRA)

  • LoRA 凍結預訓練權重，並在注意力模組的 query 與 value 投影中注入低秩矩陣 (A 與 B)。
  • 僅有適配器參數是可訓練的，顯著降低 optimizer 的記憶體需求。
  • 由於額外的矩陣乘法，前向與反向傳播的速度大約變慢兩倍，但記憶體的節省使得在消費者硬體上訓練 20 B 模型成為可能。

針對 20 B 模型在 24 GB GPU 上的端到端流程

步驟 1 – 以 8 位元精度載入模型

model = AutoModelForCausalLM.from_pretrained(
    "EleutherAI/gpt-neox-20b",
    load_in_8bit=True,
    device_map="auto",
)

以 8 位元載入可將記憶體需求從約 80 GB（float32）降低至約 20 GB，能舒適地在 24 GB 顯示卡上運行。

步驟 2 – 使用 PEFT 附加可訓練的 LoRA 適配器

from peft import get_peft_model, LoraConfig
config = LoraConfig(r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], bias="none")
model = get_peft_model(model, config)

僅低秩矩陣會被儲存在 optimizer 狀態中，將 optimizer 的記憶體需求從數十 GB 降至數百 MB。

步驟 3 – 使用單一模型產生參考與主動 logits

PEFT 的 disable_adapters 上下文管理器暫時停用 LoRA 層，允許相同底層模型產生參考 logits:

with model.disable_adapter():
    ref_logits = model(input_ids)
# active logits are computed with adapters enabled
active_logits = model(input_ids)

不需要第二個完整模型的副本，進一步降低記憶體使用量。

訓練腳本概覽

部落格文章連結了三個腳本，展示在 20 B GPT‑NeoX 模型上的完整工作流程：

  1. clm_finetune_peft_imdb.py – 在 IMDB 情感資料集上對 LoRA 適配器進行因果語言模型微調（單 epoch）。
  2. merge_peft_adapter.py – 將 LoRA 權重合併至基礎模型，以供推論或進一步訓練。
  3. gpt-neo-20b_sentiment_peft.py – 使用 IMDB 情感分類器作為獎勵模型的 PPO 微調，以產生正面的電影評論。

所有腳本皆在 NVIDIA RTX 4090（24 GB）上執行。完整的訓練也在 🤗 研究叢集中的單一 A100 上測試過。

結果

  • 損失曲線顯示在 IMDB 上進行一次 epoch 的監督式 LoRA 微調後，收斂穩定。
  • 在 PPO 期間，平均獎勵持續上升，表明模型學會產生更正面的評論。
  • 整個流程在單一 24 GB GPU 上運行，證明 RLHF 不再受限於多 GPU 叢集。

對社群的影響

  • 降低入門門檻 – 研究人員與開發者可以在消費級硬體上嘗試 RLHF。
  • 開源可重現性 – 所有程式碼與適配器皆託管於 Hugging Face Hub，方便分享微調後的產物。
  • 可擴展的基礎 – 一旦加入多 GPU 支援，這種方法即可透過資料平行擴展至更大的模型。

未解決問題與未來工作

  • 多 GPU 擴展 – 此整合在多 GPU 資料平行下的表現如何？
  • 訓練速度 – LoRA 帶來額外開銷；探索更快的 kernel 或混合精度策略或可緩解此問題。
  • 更廣泛的 RL 演算法 – 雖然 PPO 為預設，整合其他 RL 方法（例如 DPO）可能擴大其適用範圍。

參考資料

Sources