使用 DPO 微調 Llama 2

簡介
基于人類反饋的強化學習 (Reinforcement Learning from Human Feedback,RLHF) 事實上已成為 GPT-4 或 Claude 等 LLM 訓練的最后一步,它可以確保語言模型的輸出符合人類在閑聊或安全性等方面的期望。然而,它也給 NLP 引入了一些 RL 相關的復雜性: 既要構建一個好的獎勵函數,并訓練一個模型用以估計每個狀態(tài)的價值 (value); 又要注意最終生成的 LLM 不能與原始模型相差太遠,如果太遠的話會使得模型容易產生亂碼而非有意義的文本。該過程非常復雜,涉及到許多復雜的組件,而這些組件本身在訓練過程中又是動態(tài)變化的,因此把它們料理好并不容易。
Rafailov、Sharma、Mitchell 等人最近發(fā)表了一篇論文 Direct Preference Optimization,論文提出將現有方法使用的基于強化學習的目標轉換為可以通過簡單的二元交叉熵損失直接優(yōu)化的目標,這一做法大大簡化了 LLM 的提純過程。
本文介紹了直接偏好優(yōu)化 (Direct Preference Optimization,DPO) 法,該方法現已集成至 TRL 庫 中。同時,我們還展示了如何在 stack-exchange preference 數據集上微調最新的 Llama v2 7B 模型,?stack-exchange preference
?數據集中包含了各個?stack-exchange
?門戶上的各種問題及其排序后的回答。
DPO 與 PPO
在通過 RL 優(yōu)化人類衍生偏好時,一直以來的傳統(tǒng)做法是使用一個輔助獎勵模型來微調目標模型,以通過 RL 機制最大化目標模型所能獲得的獎勵。直觀上,我們使用獎勵模型向待優(yōu)化模型提供反饋,以促使它多生成高獎勵輸出,少生成低獎勵輸出。同時,我們使用凍結的參考模型來確保輸出偏差不會太大,且繼續(xù)保持輸出的多樣性。這通常需要在目標函數設計時,除了獎勵最大化目標外再添加一個相對于參考模型的 KL 懲罰項,這樣做有助于防止模型學習作弊或鉆營獎勵模型。
DPO 繞過了建模獎勵函數這一步,這源于一個關鍵洞見: 從獎勵函數到最優(yōu) RL 策略的分析映射。這個映射直觀地度量了給定獎勵函數與給定偏好數據的匹配程度。有了它,作者就可與將基于獎勵和參考模型的 RL 損失直接轉換為僅基于參考模型的損失,從而直接在偏好數據上優(yōu)化語言模型!因此,DPO 從尋找最小化 RLHF 損失的最佳方案開始,通過改變參量的方式推導出一個?僅需?參考模型的損失!
有了它,我們可以直接優(yōu)化該似然目標,而不需要獎勵模型或繁瑣的強化學習優(yōu)化過程。
如何使用 TRL 進行訓練
如前所述,一個典型的 RLHF 流水線通常包含以下幾個環(huán)節(jié):
有監(jiān)督微調 (supervised fine-tuning,SFT)
用偏好標簽標注數據
基于偏好數據訓練獎勵模型
RL 優(yōu)化
TRL 庫包含了所有這些環(huán)節(jié)所需的工具程序。而 DPO 訓練直接消滅了獎勵建模和 RL 這兩個環(huán)節(jié) (環(huán)節(jié) 3 和 4),直接根據標注好的偏好數據優(yōu)化 DPO 目標。
使用 DPO,我們仍然需要執(zhí)行環(huán)節(jié) 1,但我們僅需在 TRL 中向?DPOTrainer
?提供環(huán)節(jié) 2 準備好的偏好數據,而不再需要環(huán)節(jié) 3 和 4。標注好的偏好數據需要遵循特定的格式,它是一個含有以下 3 個鍵的字典:
prompt
?: 即推理時輸入給模型的提示chosen
?: 即針對給定提示的較優(yōu)回答rejected
?: ?即針對給定提示的較劣回答或非給定提示的回答
例如,對于?stack-exchange preference
?數據集,我們可以通過以下工具函數將數據集中的樣本映射至上述字典格式并刪除所有原始列:
一旦有了排序數據集,DPO 損失其實本質上就是一種有監(jiān)督損失,其經由參考模型獲得隱式獎勵。因此,從上層來看,DPOTrainer
?需要我們輸入待優(yōu)化的基礎模型以及參考模型:
其中,超參?beta
?是 DPO 損失的溫度,通常在?0.1
?到?0.5
?之間。它控制了我們對參考模型的關注程度,beta
?越小,我們就越忽略參考模型。對訓練器初始化后,我們就可以簡單調用以下方法,使用給定的?training_args
?在給定數據集上進行訓練了:
dpo_trainer.train()
基于 Llama v2 進行實驗
在 TRL 中實現 DPO 訓練器的好處是,人們可以利用 TRL 及其依賴庫 (如 Peft 和 Accelerate) 中已有的 LLM 相關功能。有了這些庫,我們甚至可以使用 bitsandbytes 庫提供的 QLoRA 技術 來訓練 Llama v2 模型。
有監(jiān)督微調
如上文所述,我們先用 TRL 的?SFTTrainer
?在 SFT 數據子集上使用 QLoRA 對 7B Llama v2 模型進行有監(jiān)督微調:
DPO 訓練
SFT 結束后,我們保存好生成的模型。接著,我們繼續(xù)進行 DPO 訓練,我們把 SFT 生成的模型作為 DPO 的基礎模型和參考模型,并在上文生成的?stack-exchange preference
?數據上,以 DPO 為目標函數訓練模型。我們選擇對模型進行 LoRa 微調,因此我們使用 Peft 的?AutoPeftModelForCausalLM
?函數加載模型:?
可以看出,我們以 4 比特的方式加載模型,然后通過?peft_config
?參數選擇 QLora 方法對其進行訓練。訓練器還會用評估數據集評估訓練進度,并報告一些關鍵指標,例如可以選擇通過 WandB 記錄并顯示隱式獎勵。最后,我們可以將訓練好的模型推送到 HuggingFace Hub。
總結
SFT 和 DPO 訓練腳本的完整源代碼可在該目錄 examples/stack_llama_2 處找到,訓好的已合并模型也已上傳至 HF Hub (見 此處)。
你可以在 這兒 找到我們的模型在訓練過程的 WandB 日志,其中包含了?DPOTrainer
?在訓練和評估期間記錄下來的以下獎勵指標:
rewards/chosen (較優(yōu)回答的獎勵)?
: 針對較優(yōu)回答,策略模型與參考模型的對數概率二者之差的均值,按?beta
?縮放。rewards/rejected (較劣回答的獎勵)?
: 針對較劣回答,策略模型與參考模型的對數概率二者之差的均值,按?beta
?縮放。rewards/accuracy (獎勵準確率)?
: 較優(yōu)回答的獎勵大于相應較劣回答的獎勵的頻率的均值rewards/margins (獎勵余裕值)?
: 較優(yōu)回答的獎勵與相應較劣回答的獎勵二者之差的均值。
直觀上講,在訓練過程中,我們希望余裕值增加并且準確率達到 1.0,換句話說,較優(yōu)回答的獎勵高于較劣回答的獎勵 (或余裕值大于零)。隨后,我們還可以在評估數據集上計算這些指標。
我們希望我們代碼的發(fā)布可以降低讀者的入門門檻,讓大家可以在自己的數據集上嘗試這種大語言模型對齊方法,我們迫不及待地想看到你會用它做哪些事情!如果你想試試我們訓練出來的模型,可以玩玩這個 space: trl-lib/stack-llama。
英文原文:?https://hf.co/blog/dpo-trl
原文作者: Kashif Rasul, Younes Belkada, Leandro von Werra
譯者: Matrix Yao (姚偉峰),英特爾深度學習工程師,工作方向為 transformer-family 模型在各模態(tài)數據上的應用及大規(guī)模模型的訓練推理。
審校/排版: zhongdongy (阿東)