StackLLaMA：使用 RLHF 訓練 LLaMA 於 Stack Exchange

Hugging Face 已發布 StackLLaMA，一個專門訓練用於回答 Stack Exchange 問題的模型，採用了完整的人類回饋強化學習（RLHF）流程。此專案展示了如何透過三階段的流程——監督式微調（SFT）、獎勵模型（RM）以及 RLHF 優化——將基礎 LLaMA 模型與人類偏好對齊。

訓練流程

StackLLaMA 以 LLaMA 7B 模型為基礎開發，使用 Hugging Face 的 TRL 函式庫實作以下階段：

1. 監督式微調（SFT）

為了確保模型在 RLHF 之前已熟悉目標領域，首先在 StackExchange 資料集的子集上，以因果語言模型目標進行訓練。為提升訓練效率，團隊使用了「packing」技術，即將多段文字以 EOS 標記串接，並切割成符合模型上下文大小的區塊，以消除填充標記。

2. 獎勵模型（RM）

在 RL 迴圈中，團隊並未依賴即時的人類回饋，而是訓練一個獎勵模型來模擬人類偏好。利用 StackExchange 資料集，團隊根據答案的讚數以及是否被提問者接受來為答案打分。

獎勵模型的訓練目標是預測在給定提示下兩個候選答案的排名。使用 100,000 對的子集作為訓練，並以 50,000 對的保留集進行評估，模型最終達到 67% 的準確率。

3. 人類回饋強化學習（RLHF）

最後階段包含一個由三個步驟組成的 RL 迴圈：產生回應、使用獎勵模型為其評分，並執行策略優化步驟。

為防止模型「利用」獎勵模型（產生雜訊卻獲得高分），會施加 KL 散度懲罰。此懲罰將當前策略的生成結果與凍結的參考模型進行比較，確保輸出仍與原始 SFT 模型相近。

記憶體效能訓練策略

訓練 7B 參數的模型需要大量記憶體；在 bf16 格式下，僅參數與 Adam 優化器就約需 70GB 記憶體。為了讓單顆 GPU 或消費級硬體也能使用，Hugging Face 採用了多項策略：

• PEFT 與 LoRA： 使用 peft 函式庫，團隊套用低秩適應（LoRA）並以 8 位元載入模型。這將記憶體占用降低至每十億參數約 1.2‑1.4GB。
• 資料平行化： 為在多 GPU 上擴展訓練，團隊使用 transformers.Trainer 與 accelerate，可在不修改程式碼的情況下平行化前向與反向傳播。

技術挑戰與不穩定性

在 RL 訓練大型語言模型時，常會出現多種不穩定情況，Hugging Face 團隊在開發 StackLLaMA 時已辨識出以下問題：

• 獎勵利用： PPO 演算法可能會利用獎勵模型的缺陷。例如，模型可能會產生重複的程式碼區塊（```），因為獎勵模型將程式碼區塊與高品質的 Stack Exchange 答案相關聯。
• 負 KL 散度： 雖然 KL 散度理論上應為正值，但 trl 中的估計在某些 token 被強制或抑制時（例如批次填充或抑制 EOS token）可能變為負值，導致 PPO 演算法追逐負懲罰，進而產生不穩定。
• 損失尖峰： 團隊觀察到偶發的損失尖峰，會引發進一步的訓練不穩定，這個問題正持續修正並回報至 trl 函式庫。

模型可取得性

StackLLaMA 已上架於 Hugging Face Hub。因 LLaMA 授權限制，僅釋出適配器權重與模型檢查點；使用者需向 Meta AI 申請取得基礎 LLaMA 權重。